Fix surrogate pairs crashing tokenization and saving to UTF-8

machine/jobs/nmt_engine_build_job.py

@@ -148,7 +148,7 @@
         for pt_batch in batch(src_pretranslations, batch_size):
             if check_canceled is not None:
                 check_canceled()
-            source_segments = [pt_info["translation"] for pt_info in pt_batch]
+            source_segments = [_join_surrogate_pairs_if_needed(pt_info["translation"]) for pt_info in pt_batch]
             t_batch = engine.translate_batch(source_segments)
             for pt_info, result in zip(pt_batch, t_batch, strict=True):
                 pt_info["translation"] = result.translation
@@ -253,6 +253,14 @@
             )
 
 
+def _join_surrogate_pairs_if_needed(text: str) -> str:
+    # Check if any characters are in the UTF-16 surrogate range, and join pairs and drop unpaired halves if needed.
+    # These are not supported unjoined by UTF-8 encoding or the MBartTokenizer's Rust implementation.
+    if any(0xD800 <= ord(char) <= 0xDFFF for char in text):
+        return text.encode("utf-16", "surrogatepass").decode("utf-16", "ignore")
+    return text
+
+
 def _read_pretranslations(path: Path) -> Generator[PretranslationInfo, None, None]:
     with path.open("r", encoding="utf-8") as file:
         for line in file:

tests/jobs/test_nmt_engine_build_job.py

@@ -56,6 +56,53 @@ def test_run(decoy: Decoy) -> None:
     decoy.verify(env.translation_file_service.save_model(Path("model.tar.gz"), "models/save-model.tar.gz"), times=1)
 
 
+def test_run_with_surrogate_pair(decoy: Decoy) -> None:
+    env = _TestEnvironment(
+        decoy,
+        transductive=True,
+        source_text="亚撒黑仍不肯转开。故此，押尼珥就用枪\ud863\udc54刺入他的肚腹，甚至枪从背后透出，亚撒黑就在那里仆倒而死。众人赶到亚撒黑仆倒而死的地方，就都站住。",
+    )
+    env.job.run()
+
+    pretranslations = json.loads(env.target_pretranslations)
+    assert len(pretranslations) == 1
+    assert pretranslations[0]["translation"] == "Please, I have booked a room."
+    assert pretranslations[0]["sourceTokens"] == [
+        "亚撒黑仍不肯转开",
+        "。",
+        "故此",
+        "，",
+        "押尼珥就用枪𨱔刺入他的肚腹",
+        "，",
+        "甚至枪从背后透出",
+        "，",
+        "亚撒黑就在那里仆倒而死",
+        "。",
+        "众人赶到亚撒黑仆倒而死的地方",
+        "，",
+        "就都站住",
+        "。",
+    ]
+    assert pretranslations[0]["translationTokens"] == ["Please", ",", "I", "have", "booked", "a", "room", "."]
+    assert len(pretranslations[0]["alignment"]) > 0
+    assert pretranslations[0]["sequenceConfidence"] == 0.5
+    decoy.verify(env.model.align_batch(matchers.Anything()), times=0)
+    decoy.verify(env.translation_file_service.save_model(Path("model.tar.gz"), "models/save-model.tar.gz"), times=1)
+
+
+def test_run_with_lone_surrogate(decoy: Decoy) -> None:
+    env = _TestEnvironment(
+        decoy,
+        transductive=True,
+        source_text="lone\ud863surrogate",
+    )
+    env.job.run()
+
+    pretranslations = json.loads(env.target_pretranslations)
+    assert len(pretranslations) == 1
+    assert pretranslations[0]["sourceTokens"] == ["lonesurrogate"]
+
+
 def test_run_inductive_fallback(decoy: Decoy) -> None:
     env = _TestEnvironment(decoy, transductive=False)
     env.job.run()
@@ -131,7 +178,12 @@ def _create_translation_result() -> TranslationResult:
 
 class _TestEnvironment:
     def __init__(
-        self, decoy: Decoy, transductive: bool = False, num_pretranslations: int = 1, batch_size: int = 100
+        self,
+        decoy: Decoy,
+        transductive: bool = False,
+        num_pretranslations: int = 1,
+        batch_size: int = 100,
+        source_text: str = "Por favor, tengo reservada una habitación.",
     ) -> None:
         self.source_tokenizer_trainer = decoy.mock(cls=Trainer)
         self.target_tokenizer_trainer = decoy.mock(cls=Trainer)
@@ -180,7 +232,7 @@ def _translate_batch(segments: List[str]) -> List[TranslationResult]:
                             textId=f"text{i + 1}",
                             sourceRefs=[f"ref{i + 1}"],
                             targetRefs=[f"ref{i + 1}"],
-                            translation="Por favor, tengo reservada una habitación.",
+                            translation=source_text,
                             sourceTokens=[],
                             translationTokens=[],
                             alignment="",