Add optional alignment step to nmt jobs, temporary implementation of eflomal

Author: isaac091

.devcontainer/dockerfile

@@ -22,7 +22,7 @@ RUN apt-get update && \
     python$PYTHON_VERSION-distutils \
     git vim curl gdb ca-certificates gnupg2 tar make gcc libssl-dev zlib1g-dev libncurses5-dev \
     libbz2-dev libreadline-dev libreadline6-dev libxml2-dev xz-utils libgdbm-dev libgdbm-compat-dev tk-dev dirmngr \
-    libxmlsec1-dev libsqlite3-dev libffi-dev liblzma-dev lzma lzma-dev uuid-dev && \
+    libxmlsec1-dev libsqlite3-dev libffi-dev liblzma-dev lzma lzma-dev uuid-dev python3.9-dev && \
     rm -rf /var/lib/apt/lists/*
 
 RUN curl -sS https://bootstrap.pypa.io/get-pip.py | python$PYTHON_VERSION
@@ -39,4 +39,6 @@ RUN pip install -U pip setuptools \
 
 COPY ./.devcontainer/clearml.conf /root/clearml.conf
 
+ENV EFLOMAL_PATH=/workspaces/machine.py/.venv/lib/python3.9/site-packages/eflomal/bin
+
 CMD ["bash"]

machine/jobs/build_nmt_engine.py

@@ -92,6 +92,9 @@ 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"
+    )
     args = parser.parse_args()
 
     run({k: v for k, v in vars(args).items() if v is not None})

machine/jobs/eflomal_aligner.py

@@ -0,0 +1,153 @@
+# NOTE: this is a temporary solution to be able to use the eflomal aligner inside of machine.py.
+# The vast majority of this code is taken from the silnlp repository.
+
+import os
+import subprocess
+from contextlib import ExitStack
+from math import sqrt
+from pathlib import Path
+from tempfile import TemporaryDirectory
+from typing import IO, Iterable, List, Sequence, Tuple
+
+from eflomal import read_text, write_text
+
+from ..corpora import AlignedWordPair
+from ..corpora.token_processors import escape_spaces, lowercase, normalize
+from ..tokenization import LatinWordTokenizer
+from ..translation import SymmetrizationHeuristic, WordAlignmentMatrix
+
+# may have to make more dynamic, look at silnlp get_wsl_path, is there something equivalent in machine?
+EFLOMAL_PATH = Path(os.getenv("EFLOMAL_PATH", "."), "eflomal")
+TOKENIZER = LatinWordTokenizer()
+
+
+# From silnlp.alignment.tools
+def execute_eflomal(
+    source_path: Path,
+    target_path: Path,
+    forward_links_path: Path,
+    reverse_links_path: Path,
+    n_iterations: Tuple[int, int, int],
+) -> None:
+    if not EFLOMAL_PATH.is_file():
+        raise RuntimeError("eflomal is not installed.")
+
+    args = [
+        str(EFLOMAL_PATH),
+        "-s",
+        str(source_path),
+        "-t",
+        str(target_path),
+        "-f",
+        str(forward_links_path),
+        "-r",
+        str(reverse_links_path),
+        # "-q",
+        "-m",
+        "3",
+        "-n",
+        "3",
+        "-N",
+        "0.2",
+        "-1",
+        str(n_iterations[0]),
+        "-2",
+        str(n_iterations[1]),
+        "-3",
+        str(n_iterations[2]),
+    ]
+    subprocess.run(args, stderr=subprocess.DEVNULL)
+
+
+# From silnlp.alignment.eflomal
+def to_word_alignment_matrix(alignment_str: str) -> WordAlignmentMatrix:
+    word_pairs = AlignedWordPair.from_string(alignment_str)
+    row_count = 0
+    column_count = 0
+    for pair in word_pairs:
+        if pair.source_index + 1 > row_count:
+            row_count = pair.source_index + 1
+        if pair.target_index + 1 > column_count:
+            column_count = pair.target_index + 1
+    return WordAlignmentMatrix.from_word_pairs(row_count, column_count, word_pairs)
+
+
+# From silnlp.alignment.eflomal
+def to_eflomal_text_file(input: Iterable[str], output_file: IO[bytes], prefix_len: int = 0, suffix_len: int = 0) -> int:
+    sents, index = read_text(input, True, prefix_len, suffix_len)
+    n_sents = len(sents)
+    voc_size = len(index)
+    write_text(output_file, tuple(sents), voc_size)
+    return n_sents
+
+
+# From silnlp.alignment.eflomal
+def prepare_files(
+    src_input: Iterable[str], src_output_file: IO[bytes], trg_input: Iterable[str], trg_output_file: IO[bytes]
+) -> int:
+    n_src_sents = to_eflomal_text_file(src_input, src_output_file)
+    n_trg_sents = to_eflomal_text_file(trg_input, trg_output_file)
+    if n_src_sents != n_trg_sents:
+        raise ValueError("Mismatched file sizes")
+    return n_src_sents
+
+
+def tokenize(sent: str) -> Sequence[str]:
+    return lowercase(normalize("NFC", escape_spaces(list(TOKENIZER.tokenize(sent)))))
+
+
+# From silnlp.alignment.eflomal
+class EflomalAligner:
+    def __init__(self, model_dir: Path) -> None:
+        self._model_dir = model_dir
+
+    def train(self, src_toks: Sequence[Sequence[str]], trg_toks: Sequence[Sequence[str]]) -> None:
+        self._model_dir.mkdir(exist_ok=True)
+        with TemporaryDirectory() as temp_dir:
+            src_eflomal_path = Path(temp_dir, "source")
+            trg_eflomal_path = Path(temp_dir, "target")
+            with ExitStack() as stack:
+                src_output_file = stack.enter_context(src_eflomal_path.open("wb"))
+                trg_output_file = stack.enter_context(trg_eflomal_path.open("wb"))
+                # Write input files for the eflomal binary
+                n_sentences = prepare_files(
+                    [" ".join(s) for s in src_toks], src_output_file, [" ".join(s) for s in trg_toks], trg_output_file
+                )
+
+            iters = max(2, int(round(1.0 * 5000 / sqrt(n_sentences))))
+            iters4 = max(1, iters // 4)
+            n_iterations = (max(2, iters4), iters4, iters)
+
+            # Run wrapper for the eflomal binary
+            execute_eflomal(
+                src_eflomal_path,
+                trg_eflomal_path,
+                self._model_dir / "forward-align.txt",
+                self._model_dir / "reverse-align.txt",
+                n_iterations,
+            )
+
+    def align(self, sym_heuristic: str = "grow-diag-final-and") -> List[str]:
+        forward_align_path = self._model_dir / "forward-align.txt"
+        reverse_align_path = self._model_dir / "reverse-align.txt"
+
+        alignments = []
+        heuristic = SymmetrizationHeuristic[sym_heuristic.upper().replace("-", "_")]
+        with ExitStack() as stack:
+            forward_file = stack.enter_context(forward_align_path.open("r", encoding="utf-8-sig"))
+            reverse_file = stack.enter_context(reverse_align_path.open("r", encoding="utf-8-sig"))
+
+            for forward_line, reverse_line in zip(forward_file, reverse_file):
+                forward_matrix = to_word_alignment_matrix(forward_line.strip())
+                reverse_matrix = to_word_alignment_matrix(reverse_line.strip())
+                src_len = max(forward_matrix.row_count, reverse_matrix.row_count)
+                trg_len = max(forward_matrix.column_count, reverse_matrix.column_count)
+
+                forward_matrix.resize(src_len, trg_len)
+                reverse_matrix.resize(src_len, trg_len)
+
+                forward_matrix.symmetrize_with(reverse_matrix, heuristic)
+
+                alignments.append(str(forward_matrix))
+
+        return alignments

machine/jobs/nmt_engine_build_job.py

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

machine/jobs/smt_engine_build_job.py

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

machine/jobs/translation_engine_build_job.py

@@ -1,12 +1,16 @@
 import logging
 from abc import ABC, abstractmethod
+from contextlib import ExitStack
+from pathlib import Path
+from tempfile import TemporaryDirectory
 from typing import Any, Callable, Optional, Tuple
 
 from ..corpora.parallel_text_corpus import ParallelTextCorpus
 from ..corpora.text_corpus import TextCorpus
 from ..utils.phased_progress_reporter import PhasedProgressReporter
 from ..utils.progress_status import ProgressStatus
-from .translation_file_service import TranslationFileService
+from .eflomal_aligner import EflomalAligner, tokenize
+from .translation_file_service import PretranslationInfo, TranslationFileService
 
 logger = logging.getLogger(__name__)
 
@@ -44,6 +48,10 @@ def run(
         logger.info("Pretranslating segments")
         self._batch_inference(progress_reporter, check_canceled)
 
+        if "align_pretranslations" in self._config and self._config.align_pretranslations:
+            logger.info("Aligning source to pretranslations")
+            self._align(progress_reporter, check_canceled)
+
         self._save_model()
         return train_corpus_size, confidence
 
@@ -74,5 +82,59 @@ def _batch_inference(
         check_canceled: Optional[Callable[[], None]],
     ) -> None: ...
 
+    def _align(
+        self,
+        progress_reporter: PhasedProgressReporter,
+        check_canceled: Optional[Callable[[], None]],
+    ) -> None:
+        if check_canceled is not None:
+            check_canceled()
+
+        logger.info("Aligning source to pretranslations")
+        with ExitStack() as stack:
+            phase_progress = stack.enter_context(progress_reporter.start_next_phase())
+
+            src_tokenized = [
+                tokenize(s["pretranslation"])
+                for s in stack.enter_context(self._translation_file_service.get_source_pretranslations())
+            ]
+            trg_tokenized = [
+                tokenize(s["pretranslation"])
+                for s in stack.enter_context(self._translation_file_service.get_target_pretranslations())
+            ]
+
+            with TemporaryDirectory() as td:
+                aligner = EflomalAligner(Path(td))
+                logger.info("Training aligner")
+                aligner.train(src_tokenized, trg_tokenized)
+
+                if check_canceled is not None:
+                    check_canceled()
+
+                logger.info("Aligning pretranslations")
+                alignments = aligner.align()
+
+            if check_canceled is not None:
+                check_canceled()
+
+            writer = stack.enter_context(self._translation_file_service.open_target_pretranslation_writer())
+            for trg_pi, src_toks, trg_toks, alignment in zip(
+                stack.enter_context(self._translation_file_service.get_target_pretranslations()),
+                src_tokenized,
+                trg_tokenized,
+                alignments,
+            ):
+                writer.write(
+                    PretranslationInfo(
+                        corpusId=trg_pi["corpusId"],
+                        textId=trg_pi["textId"],
+                        refs=trg_pi["refs"],
+                        pretranslation=trg_pi["pretranslation"],
+                        source_toks=list(src_toks),
+                        pretranslation_toks=list(trg_toks),
+                        alignment=alignment,
+                    )
+                )
+
     @abstractmethod
     def _save_model(self) -> None: ...

machine/jobs/translation_file_service.py

@@ -15,7 +15,10 @@ class PretranslationInfo(TypedDict):
     corpusId: str  # noqa: N815
     textId: str  # noqa: N815
     refs: List[str]
-    translation: str
+    pretranslation: str
+    source_toks: List[str]
+    pretranslation_toks: List[str]
+    alignment: str
 
 
 SOURCE_FILENAME = "train.src.txt"
@@ -49,23 +52,30 @@ def exists_source_corpus(self) -> bool:
     def exists_target_corpus(self) -> bool:
         return self.shared_file_service._exists_file(f"{self.shared_file_service.build_path}/{TARGET_FILENAME}")
 
-    def get_source_pretranslations(self) -> ContextManagedGenerator[PretranslationInfo, None, None]:
-        src_pretranslate_path = self.shared_file_service.download_file(
-            f"{self.shared_file_service.build_path}/{SOURCE_PRETRANSLATION_FILENAME}"
-        )
+    def _get_pretranslations(self, filename: str) -> ContextManagedGenerator[PretranslationInfo, None, None]:
+        pretranslate_path = self.shared_file_service.download_file(f"{self.shared_file_service.build_path}/{filename}")
 
         def generator() -> Generator[PretranslationInfo, None, None]:
-            with src_pretranslate_path.open("r", encoding="utf-8-sig") as file:
+            with pretranslate_path.open("r", encoding="utf-8-sig") as file:
                 for pi in json_stream.load(file):
                     yield PretranslationInfo(
                         corpusId=pi["corpusId"],
                         textId=pi["textId"],
                         refs=list(pi["refs"]),
-                        translation=pi["translation"],
+                        pretranslation=pi["pretranslation"],
+                        source_toks=list(pi["source_toks"]),
+                        pretranslation_toks=list(pi["pretranslation_toks"]),
+                        alignment=pi["alignment"],
                     )
 
         return ContextManagedGenerator(generator())
 
+    def get_source_pretranslations(self) -> ContextManagedGenerator[PretranslationInfo, None, None]:
+        return self._get_pretranslations(SOURCE_PRETRANSLATION_FILENAME)
+
+    def get_target_pretranslations(self) -> ContextManagedGenerator[PretranslationInfo, None, None]:
+        return self._get_pretranslations(TARGET_PRETRANSLATION_FILENAME)
+
     def save_model(self, model_path: Path, destination: str) -> None:
         self.shared_file_service.upload_path(model_path, destination)