diff --git a/machine/jobs/nmt_engine_build_job.py b/machine/jobs/nmt_engine_build_job.py index cf25d4b7..3b6dc464 100644 --- a/machine/jobs/nmt_engine_build_job.py +++ b/machine/jobs/nmt_engine_build_job.py @@ -148,7 +148,7 @@ def _translate( 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 _save_model(self) -> None: ) +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: diff --git a/tests/jobs/test_nmt_engine_build_job.py b/tests/jobs/test_nmt_engine_build_job.py index 57905ce9..d6d36051 100644 --- a/tests/jobs/test_nmt_engine_build_job.py +++ b/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="",