Move alignment config option to build options, revert to 'translation' in PretranslationInfo

Author: isaac091

machine/jobs/build_nmt_engine.py

@@ -92,13 +92,6 @@ def main() -> None:
     parser.add_argument("--clearml", default=False, action="store_true", help="Initializes a ClearML task")
     parser.add_argument("--build-options", default=None, type=str, help="Build configurations")
     parser.add_argument("--save-model", default=None, type=str, help="Save the model using the specified base name")
-    parser.add_argument(
-        "--align-pretranslations",
-        default=False,
-        action="store_true",
-        help="Aligns source and target pretranslations using Eflomal (linux only) "
-        "and returns the alignments as well as the tokenized source and target with the pretranslations.",
-    )
     args = parser.parse_args()
 
     run({k: v for k, v in vars(args).items() if v is not None})

machine/jobs/nmt_engine_build_job.py

@@ -28,7 +28,7 @@ def _get_progress_reporter(
         self, progress: Optional[Callable[[ProgressStatus], None]], corpus_size: int
     ) -> PhasedProgressReporter:
         if corpus_size > 0:
-            if "align_pretranslations" in self._config and self._config.align_pretranslations:
+            if self._config.align_pretranslations:
                 phases = [
                     Phase(message="Training NMT model", percentage=0.8),
                     Phase(message="Pretranslating segments", percentage=0.1),
@@ -40,7 +40,7 @@ def _get_progress_reporter(
                     Phase(message="Pretranslating segments", percentage=0.1),
                 ]
         else:
-            if "align_pretranslations" in self._config and self._config.align_pretranslations:
+            if self._config.align_pretranslations:
                 phases = [
                     Phase(message="Pretranslating segments", percentage=0.9),
                     Phase(message="Aligning segments", percentage=0.1, report_steps=False),
@@ -128,7 +128,7 @@ def _translate_batch(
     batch: Sequence[PretranslationInfo],
     writer: DictToJsonWriter,
 ) -> None:
-    source_segments = [pi["pretranslation"] for pi in batch]
+    source_segments = [pi["translation"] for pi in batch]
     for i, result in enumerate(engine.translate_batch(source_segments)):
-        batch[i]["pretranslation"] = result.translation
+        batch[i]["translation"] = result.translation
         writer.write(batch[i])

machine/jobs/settings.yaml

@@ -3,6 +3,7 @@ default:
   shared_file_uri: s3:/silnlp/
   shared_file_folder: production
   inference_batch_size: 1024
+  align_pretranslations: false
   huggingface:
     parent_model_name: facebook/nllb-200-distilled-1.3B
     train_params:

machine/jobs/smt_engine_build_job.py

@@ -107,7 +107,7 @@ def _translate_batch(
     batch: Sequence[PretranslationInfo],
     writer: DictToJsonWriter,
 ) -> None:
-    source_segments = [pi["pretranslation"] for pi in batch]
+    source_segments = [pi["translation"] for pi in batch]
     for i, result in enumerate(engine.translate_batch(source_segments)):
-        batch[i]["pretranslation"] = result.translation
+        batch[i]["translation"] = result.translation
         writer.write(batch[i])

machine/jobs/translation_engine_build_job.py

@@ -48,7 +48,7 @@ def run(
         logger.info("Pretranslating segments")
         self._batch_inference(progress_reporter, check_canceled)
 
-        if "align_pretranslations" in self._config and self._config.align_pretranslations and is_eflomal_available():
+        if self._config.align_pretranslations and is_eflomal_available():
             logger.info("Aligning source to pretranslations")
             self._align(progress_reporter, check_canceled)
 
@@ -96,13 +96,13 @@ def _align(
             progress_reporter.start_next_phase()
 
             src_tokenized = [
-                tokenize(s["pretranslation"])
+                tokenize(s["translation"])
                 for s in stack.enter_context(self._translation_file_service.get_source_pretranslations())
             ]
             trg_info = [
                 pt_info for pt_info in stack.enter_context(self._translation_file_service.get_target_pretranslations())
             ]
-            trg_tokenized = [tokenize(pt_info["pretranslation"]) for pt_info in trg_info]
+            trg_tokenized = [tokenize(pt_info["translation"]) for pt_info in trg_info]
 
             with TemporaryDirectory() as td:
                 aligner = EflomalAligner(Path(td))
@@ -125,9 +125,9 @@ def _align(
                         corpusId=trg_pi["corpusId"],
                         textId=trg_pi["textId"],
                         refs=trg_pi["refs"],
-                        pretranslation=trg_pi["pretranslation"],
+                        translation=trg_pi["translation"],
                         source_toks=list(src_toks),
-                        pretranslation_toks=list(trg_toks),
+                        translation_toks=list(trg_toks),
                         alignment=alignment,
                     )
                 )

machine/jobs/translation_file_service.py

@@ -15,9 +15,9 @@ class PretranslationInfo(TypedDict):
     corpusId: str  # noqa: N815
     textId: str  # noqa: N815
     refs: List[str]
-    pretranslation: str
+    translation: str
     source_toks: List[str]
-    pretranslation_toks: List[str]
+    translation_toks: List[str]
     alignment: str
 
 
@@ -62,9 +62,9 @@ def generator() -> Generator[PretranslationInfo, None, None]:
                         corpusId=pi["corpusId"],
                         textId=pi["textId"],
                         refs=list(pi["refs"]),
-                        pretranslation=pi["pretranslation"],
+                        translation=pi["translation"],
                         source_toks=list(pi["source_toks"]),
-                        pretranslation_toks=list(pi["pretranslation_toks"]),
+                        translation_toks=list(pi["translation_toks"]),
                         alignment=pi["alignment"],
                     )
 

tests/jobs/test_nmt_engine_build_job.py

@@ -36,7 +36,7 @@ def test_run(decoy: Decoy) -> None:
 
     pretranslations = json.loads(env.target_pretranslations)
     assert len(pretranslations) == 1
-    assert pretranslations[0]["pretranslation"] == "Please, I have booked a room."
+    assert pretranslations[0]["translation"] == "Please, I have booked a room."
     if is_eflomal_available():
         assert pretranslations[0]["source_toks"] == [
             "Por",
@@ -48,11 +48,11 @@ def test_run(decoy: Decoy) -> None:
             "habitación",
             ".",
         ]
-        assert pretranslations[0]["pretranslation_toks"] == ["Please", ",", "I", "have", "booked", "a", "room", "."]
+        assert pretranslations[0]["translation_toks"] == ["Please", ",", "I", "have", "booked", "a", "room", "."]
         assert len(pretranslations[0]["alignment"]) > 0
     else:
         assert pretranslations[0]["source_toks"] == []
-        assert pretranslations[0]["pretranslation_toks"] == []
+        assert pretranslations[0]["translation_toks"] == []
         assert len(pretranslations[0]["alignment"]) == 0
     decoy.verify(env.translation_file_service.save_model(Path("model.tar.gz"), "models/save-model.tar.gz"), times=1)
 
@@ -130,9 +130,9 @@ def __init__(self, decoy: Decoy) -> None:
                             corpusId="corpus1",
                             textId="text1",
                             refs=["ref1"],
-                            pretranslation="Por favor, tengo reservada una habitación.",
+                            translation="Por favor, tengo reservada una habitación.",
                             source_toks=[],
-                            pretranslation_toks=[],
+                            translation_toks=[],
                             alignment="",
                         )
                     ]
@@ -148,9 +148,9 @@ def __init__(self, decoy: Decoy) -> None:
                             corpusId="corpus1",
                             textId="text1",
                             refs=["ref1"],
-                            pretranslation="Please, I have booked a room.",
+                            translation="Please, I have booked a room.",
                             source_toks=[],
-                            pretranslation_toks=[],
+                            translation_toks=[],
                             alignment="",
                         )
                     ]

tests/jobs/test_smt_engine_build_job.py

@@ -31,7 +31,7 @@ def test_run(decoy: Decoy) -> None:
 
     pretranslations = json.loads(env.target_pretranslations)
     assert len(pretranslations) == 1
-    assert pretranslations[0]["pretranslation"] == "Please, I have booked a room."
+    assert pretranslations[0]["translation"] == "Please, I have booked a room."
     decoy.verify(
         env.translation_file_service.save_model(matchers.Anything(), f"builds/{env.job._config.build_id}/model.zip"),
         times=1,
@@ -136,9 +136,9 @@ def __init__(self, decoy: Decoy) -> None:
                             corpusId="corpus1",
                             textId="text1",
                             refs=["ref1"],
-                            pretranslation="Por favor, tengo reservada una habitación.",
+                            translation="Por favor, tengo reservada una habitación.",
                             source_toks=[],
-                            pretranslation_toks=[],
+                            translation_toks=[],
                             alignment="",
                         )
                     ]
@@ -161,7 +161,14 @@ def open_target_pretranslation_writer(env: _TestEnvironment) -> Iterator[DictToJ
         )
 
         self.job = SmtEngineBuildJob(
-            MockSettings({"build_id": "mybuild", "inference_batch_size": 100, "thot_mt": {"tokenizer": "latin"}}),
+            MockSettings(
+                {
+                    "build_id": "mybuild",
+                    "inference_batch_size": 100,
+                    "thot_mt": {"tokenizer": "latin"},
+                    "align_pretranslations": False,
+                }
+            ),
             self.smt_model_factory,
             self.translation_file_service,
         )