Working test, minor updates to naming

sillsdev · Jan 17, 2025 · e862df2 · e862df2
1 parent 672c630
commit e862df2
Show file tree

Hide file tree

Showing 3 changed files with 31 additions and 5 deletions.
diff --git a/machine/jobs/build_word_alignment_model.py b/machine/jobs/build_word_alignment_model.py
@@ -89,7 +89,7 @@ def run(args: dict):
 
 
 def main() -> None:
-    parser = argparse.ArgumentParser(description="Trains an SMT model.")
+    parser = argparse.ArgumentParser(description="Trains a word alignment model.")
     parser.add_argument("--model-type", required=True, type=str, help="Model type")
     parser.add_argument("--build-id", required=True, type=str, help="Build id")
     parser.add_argument("--clearml", default=False, action="store_true", help="Initializes a ClearML task")

diff --git a/machine/jobs/word_alignment_file_service.py b/machine/jobs/word_alignment_file_service.py
@@ -25,8 +25,8 @@ def __init__(
         config: Any,
         source_filename: str = "train.src.txt",
         target_filename: str = "train.trg.txt",
-        word_alignment_input_filename: str = "word_alignment_inputs.json",
-        word_alignment_output_filename: str = "word_alignment_outputs.json",
+        word_alignment_input_filename: str = "word_alignments.inputs.json",
+        word_alignment_output_filename: str = "word_alignments.outputs.json",
     ) -> None:
 
         self._source_filename = source_filename

diff --git a/tests/jobs/test_word_alignment_build_job.py b/tests/jobs/test_word_alignment_build_job.py
@@ -10,7 +10,7 @@
 
 from machine.corpora import DictionaryTextCorpus, MemoryText, TextRow
 from machine.jobs import DictToJsonWriter, WordAlignmentBuildJob, WordAlignmentModelFactory
-from machine.jobs.word_alignment_file_service import WordAlignmentFileService
+from machine.jobs.word_alignment_file_service import WordAlignmentFileService, WordAlignmentInput
 from machine.translation import Trainer, TrainStats, WordAlignmentMatrix
 from machine.translation.word_alignment_model import WordAlignmentModel
 from machine.utils import CanceledError
@@ -90,6 +90,32 @@ def __init__(self, decoy: Decoy) -> None:
         decoy.when(self.word_alignment_file_service.exists_source_corpus()).then_return(True)
         decoy.when(self.word_alignment_file_service.exists_target_corpus()).then_return(True)
 
+        decoy.when(self.word_alignment_file_service.get_word_alignment_inputs()).then_return(
+            [
+                WordAlignmentInput(
+                    corpusId="corpus1",
+                    textId="text1",
+                    refs=["1"],
+                    source="¿Le importaría darnos las llaves de la habitación, por favor?",
+                    target="Would you mind giving us the room keys, please?",
+                ),
+                WordAlignmentInput(
+                    corpusId="corpus1",
+                    textId="text1",
+                    refs=["2"],
+                    source="¿Le importaría cambiarme a otra habitación más tranquila?",
+                    target="Would you mind moving me to another quieter room?",
+                ),
+                WordAlignmentInput(
+                    corpusId="corpus1",
+                    textId="text1",
+                    refs=["3"],
+                    source="Me parece que existe un problema.",
+                    target="I think there is a problem.",
+                ),
+            ]
+        )
+
         self.alignment_json = ""
 
         @contextmanager
@@ -100,7 +126,7 @@ def open_target_alignment_writer(env: _TestEnvironment) -> Iterator[DictToJsonWr
             file.write("\n]\n")
             env.alignment_json = file.getvalue()
 
-        decoy.when(self.word_alignment_file_service.open_target_alignment_writer()).then_do(
+        decoy.when(self.word_alignment_file_service.open_alignment_output_writer()).then_do(
             lambda: open_target_alignment_writer(self)
         )