Make word_align_corpus train lazily

machine/translation/corpus_ops.py

@@ -20,17 +20,7 @@ def word_align_corpus(
     progress: Optional[Callable[[ProgressStatus], None]] = None,
 ) -> ParallelTextCorpus:
     if isinstance(aligner, (int, str)):
-        from .thot import create_thot_symmetrized_word_alignment_model
-
-        model = create_thot_symmetrized_word_alignment_model(aligner)
-        model.heuristic = symmetrization_heuristic
-        # Retain the alignments computed during training so that the corpus can be aligned
-        # without a separate, potentially expensive, inference pass.
-        model.emit_training_alignments = True
-        with model.create_trainer(corpus) as trainer:
-            trainer.train(progress)
-            trainer.save()
-        aligner = model
+        return _TrainedWordAlignParallelTextCorpus(corpus, aligner, symmetrization_heuristic, progress)
 
     if isinstance(aligner, TransductiveWordAlignmentModel):
         return _TransductiveWordAlignParallelTextCorpus(corpus, aligner)
@@ -43,18 +33,29 @@ def translate_corpus(
     return _TranslateParallelTextCorpus(corpus, translation_engine, batch_size)
 
 
-class _WordAlignParallelTextCorpus(ParallelTextCorpus):
-    def __init__(self, corpus: ParallelTextCorpus, aligner: WordAligner, batch_size: int) -> None:
+class _WordAlignParallelTextCorpusBase(ParallelTextCorpus):
+    def __init__(self, corpus: ParallelTextCorpus) -> None:
         self._corpus = corpus
-        self._aligner = aligner
-        self._batch_size = batch_size
 
+    @property
     def is_source_tokenized(self) -> bool:
         return self._corpus.is_source_tokenized
 
+    @property
     def is_target_tokenized(self) -> bool:
         return self._corpus.is_target_tokenized
 
+    def count(self, include_empty: bool = True, text_ids: Optional[Iterable[str]] = None) -> int:
+        # Aligning does not add or remove rows, so counting need not align, which may train a model.
+        return self._corpus.count(include_empty, text_ids)
+
+
+class _WordAlignParallelTextCorpus(_WordAlignParallelTextCorpusBase):
+    def __init__(self, corpus: ParallelTextCorpus, aligner: WordAligner, batch_size: int) -> None:
+        super().__init__(corpus)
+        self._aligner = aligner
+        self._batch_size = batch_size
+
     def _get_rows(self, text_ids: Optional[Iterable[str]] = None) -> Generator[ParallelTextRow, None, None]:
         with self._corpus.get_rows(text_ids) as rows:
             for row_batch in batch(rows, self._batch_size):
@@ -73,36 +74,67 @@ def _get_rows(self, text_ids: Optional[Iterable[str]] = None) -> Generator[Paral
                     yield row
 
 
-class _TransductiveWordAlignParallelTextCorpus(ParallelTextCorpus):
+class _TransductiveWordAlignParallelTextCorpus(_WordAlignParallelTextCorpusBase):
     def __init__(self, corpus: ParallelTextCorpus, model: TransductiveWordAlignmentModel) -> None:
-        self._corpus = corpus
+        super().__init__(corpus)
         self._model = model
 
-    def is_source_tokenized(self) -> bool:
-        return self._corpus.is_source_tokenized
-
-    def is_target_tokenized(self) -> bool:
-        return self._corpus.is_target_tokenized
+    def _get_rows(self, text_ids: Optional[Iterable[str]] = None) -> Generator[ParallelTextRow, None, None]:
+        yield from _get_transductive_rows(self._corpus, self._model, text_ids)
+
+
+class _TrainedWordAlignParallelTextCorpus(_WordAlignParallelTextCorpusBase):
+    def __init__(
+        self,
+        corpus: ParallelTextCorpus,
+        aligner: Union[int, str],
+        symmetrization_heuristic: SymmetrizationHeuristic,
+        progress: Optional[Callable[[ProgressStatus], None]],
+    ) -> None:
+        super().__init__(corpus)
+        self._aligner = aligner
+        self._symmetrization_heuristic = symmetrization_heuristic
+        self._progress = progress
 
     def _get_rows(self, text_ids: Optional[Iterable[str]] = None) -> Generator[ParallelTextRow, None, None]:
-        # The training alignments are keyed by the order in which the sentence pairs were added
-        # during training, so the full corpus must be iterated to keep the index in sync; rows that
-        # are not in the requested texts are skipped rather than filtered out of the enumeration.
-        text_id_set = None if text_ids is None else set(text_ids)
-        with self._corpus.get_rows() as rows:
-            for index, row in enumerate(rows):
-                if text_id_set is not None and row.text_id not in text_id_set:
-                    continue
-                alignment = self._model.get_training_alignment(index)
-                known_alignment = WordAlignmentMatrix.from_parallel_text_row(row)
-                if known_alignment is not None:
-                    known_alignment.priority_symmetrize_with(alignment)
-                    alignment = known_alignment
-                word_pairs = alignment.to_aligned_word_pairs()
-                if isinstance(self._model, WordAlignmentModel):
-                    self._model.compute_aligned_word_pair_scores(row.source_segment, row.target_segment, word_pairs)
-                row.aligned_word_pairs = word_pairs
-                yield row
+        from .thot import create_thot_symmetrized_word_alignment_model
+
+        # Training on only the requested texts keeps the training-alignment index in sync with the rows.
+        corpus = self._corpus.filter_texts(text_ids)
+        # Training in the generator ties the model's lifetime to reading the rows, at the cost of
+        # training a new model on each iteration.
+        with create_thot_symmetrized_word_alignment_model(self._aligner) as model:
+            model.heuristic = self._symmetrization_heuristic
+            # Retain the alignments computed during training so that the corpus can be aligned
+            # without a separate, potentially expensive, inference pass.
+            model.emit_training_alignments = True
+            with model.create_trainer(corpus) as trainer:
+                trainer.train(self._progress)
+                trainer.save()
+            yield from _get_transductive_rows(corpus, model, None)
+
+
+def _get_transductive_rows(
+    corpus: ParallelTextCorpus, model: TransductiveWordAlignmentModel, text_ids: Optional[Iterable[str]]
+) -> Generator[ParallelTextRow, None, None]:
+    # The training alignments are keyed by the order in which the sentence pairs were added during
+    # training, so the corpus the model was trained on must be iterated in full to keep the index in
+    # sync; rows outside the requested texts are skipped rather than filtered out.
+    text_id_set = None if text_ids is None else set(text_ids)
+    with corpus.get_rows() as rows:
+        for index, row in enumerate(rows):
+            if text_id_set is not None and row.text_id not in text_id_set:
+                continue
+            alignment = model.get_training_alignment(index)
+            known_alignment = WordAlignmentMatrix.from_parallel_text_row(row)
+            if known_alignment is not None:
+                known_alignment.priority_symmetrize_with(alignment)
+                alignment = known_alignment
+            word_pairs = alignment.to_aligned_word_pairs()
+            if isinstance(model, WordAlignmentModel):
+                model.compute_aligned_word_pair_scores(row.source_segment, row.target_segment, word_pairs)
+            row.aligned_word_pairs = word_pairs
+            yield row
 
 
 class _TranslateParallelTextCorpus(ParallelTextCorpus):
@@ -111,9 +143,11 @@ def __init__(self, corpus: ParallelTextCorpus, translation_engine: TranslationEn
         self._translation_engine = translation_engine
         self._batch_size = batch_size
 
+    @property
     def is_source_tokenized(self) -> bool:
         return self._corpus.is_source_tokenized
 
+    @property
     def is_target_tokenized(self) -> bool:
         return self._corpus.is_target_tokenized
 

samples/word_alignment.ipynb

@@ -59,7 +59,9 @@
    "source": [
     "## Simple word alignment\n",
     "\n",
-    "The easiest way to align a parallel corpus is to use the `word_align_corpus` function. The function will train the model and align the corpus. The alignment will be stored in the `aligned_word_pairs` property as a collection of `AlignedWordPair` instances. By default, the `word_align_corpus` function uses FastAlign. "
+    "The easiest way to align a parallel corpus is to use the `word_align_corpus` function. The function will train the model and align the corpus. The alignment will be stored in the `aligned_word_pairs` property as a collection of `AlignedWordPair` instances. By default, the `word_align_corpus` function uses FastAlign.\n",
+    "\n",
+    "The returned corpus is lazy: the model is trained when the rows are read, and it is released as soon as reading finishes. Each iteration therefore trains a new model, so read the rows once and keep what you need rather than iterating the returned corpus repeatedly."
    ]
   },
   {

tests/translation/test_corpus_ops.py

@@ -1,6 +1,7 @@
 from typing import Iterable, Optional
 
 import pytest
+from decoy import Decoy, matchers
 from testutils.thot_test_helpers import create_test_parallel_corpus
 
 from machine.corpora import (
@@ -11,8 +12,10 @@
     StandardParallelTextCorpus,
     TextRow,
 )
-from machine.translation import SymmetrizationHeuristic, word_align_corpus
-from machine.translation.thot import create_thot_symmetrized_word_alignment_model
+from machine.translation import SymmetrizationHeuristic, Trainer, WordAlignmentMatrix, thot, word_align_corpus
+from machine.translation.thot import ThotSymmetrizedWordAlignmentModel, create_thot_symmetrized_word_alignment_model
+
+_ANY = matchers.AnythingOrNone()
 
 
 def _alignment_strings(corpus: ParallelTextCorpus, text_ids: Optional[Iterable[str]] = None) -> list:
@@ -21,23 +24,29 @@ def _alignment_strings(corpus: ParallelTextCorpus, text_ids: Optional[Iterable[s
     ]
 
 
+def _create_trained_model(corpus: ParallelTextCorpus, aligner: str = "fast_align") -> ThotSymmetrizedWordAlignmentModel:
+    model = create_thot_symmetrized_word_alignment_model(aligner)
+    model.heuristic = SymmetrizationHeuristic.GROW_DIAG_FINAL_AND  # word_align_corpus's default
+    model.emit_training_alignments = True
+    with model.create_trainer(corpus) as trainer:
+        trainer.train()
+        trainer.save()
+    return model
+
+
 @pytest.mark.parametrize("aligner", ["fast_align", "ibm1"])
 def test_word_align_corpus_transductive_matches_inference(aligner: str) -> None:
     # For deterministic models, the alignments retained during training match those produced by a
     # separate inference pass, so the transductive output must equal aligning each row directly.
     transductive = _alignment_strings(word_align_corpus(create_test_parallel_corpus(), aligner=aligner))
 
-    model = create_thot_symmetrized_word_alignment_model(aligner)
-    model.heuristic = SymmetrizationHeuristic.GROW_DIAG_FINAL_AND  # word_align_corpus's default
-    with model.create_trainer(create_test_parallel_corpus()) as trainer:
-        trainer.train()
-        trainer.save()
-    inference = [
-        AlignedWordPair.to_string(
-            model.align(row.source_segment, row.target_segment).to_aligned_word_pairs(), include_scores=False
-        )
-        for row in create_test_parallel_corpus().get_rows()
-    ]
+    with _create_trained_model(create_test_parallel_corpus(), aligner) as model:
+        inference = [
+            AlignedWordPair.to_string(
+                model.align(row.source_segment, row.target_segment).to_aligned_word_pairs(), include_scores=False
+            )
+            for row in create_test_parallel_corpus().get_rows()
+        ]
     assert transductive == inference
 
 
@@ -62,14 +71,110 @@ def _create_two_text_parallel_corpus() -> StandardParallelTextCorpus:
 def test_word_align_corpus_transductive_text_ids_keep_index_in_sync() -> None:
     # Filtering by text must not desync the training-alignment index: the rows for a requested text
     # must get exactly the alignments they got in the unfiltered pass, not those of earlier rows.
-    corpus = word_align_corpus(_create_two_text_parallel_corpus(), aligner="fast_align")
-    full = list(corpus.get_rows())
-    text2_expected = [AlignedWordPair.to_string(r.aligned_word_pairs, include_scores=False) for r in full[2:]]
-    text2_actual = _alignment_strings(corpus, ["text2"])
+    # The model is trained up front so that both passes read the same training alignments.
+    parallel_corpus = _create_two_text_parallel_corpus()
+    with _create_trained_model(parallel_corpus) as model:
+        corpus = word_align_corpus(parallel_corpus, aligner=model)
+        full = list(corpus.get_rows())
+        text2_expected = [AlignedWordPair.to_string(r.aligned_word_pairs, include_scores=False) for r in full[2:]]
+        text2_actual = _alignment_strings(corpus, ["text2"])
     assert text2_actual == text2_expected
 
 
 def test_word_align_corpus_transductive_eflomal() -> None:
     rows = list(word_align_corpus(create_test_parallel_corpus(), aligner="eflomal").get_rows())
     assert len(rows) == 8
     assert any(row.aligned_word_pairs for row in rows)
+
+
+def _create_mock_model(decoy: Decoy) -> ThotSymmetrizedWordAlignmentModel:
+    model = decoy.mock(cls=ThotSymmetrizedWordAlignmentModel)
+    decoy.when(model.__enter__()).then_return(model)
+    decoy.when(model.get_training_alignment(_ANY)).then_return(
+        WordAlignmentMatrix.from_word_pairs(row_count=2, column_count=2, set_values=[(0, 0), (1, 1)])
+    )
+    return model
+
+
+class _TestEnvironment:
+    """Replaces the model that word_align_corpus creates internally with a mock."""
+
+    def __init__(self, decoy: Decoy, monkeypatch: pytest.MonkeyPatch) -> None:
+        self.training_corpus: Optional[ParallelTextCorpus] = None
+
+        self.trainer = decoy.mock(cls=Trainer)
+        decoy.when(self.trainer.__enter__()).then_return(self.trainer)
+
+        self.model = _create_mock_model(decoy)
+        decoy.when(self.model.create_trainer(_ANY)).then_do(self._create_trainer)
+
+        create_model = decoy.mock(func=create_thot_symmetrized_word_alignment_model)
+        decoy.when(create_model(_ANY)).then_return(self.model)
+        monkeypatch.setattr(thot, "create_thot_symmetrized_word_alignment_model", create_model)
+
+    def _create_trainer(self, corpus: ParallelTextCorpus) -> Trainer:
+        self.training_corpus = corpus
+        return self.trainer
+
+
+def test_word_align_corpus_trains_lazily(decoy: Decoy, monkeypatch: pytest.MonkeyPatch) -> None:
+    env = _TestEnvironment(decoy, monkeypatch)
+
+    corpus = word_align_corpus(create_test_parallel_corpus())
+    decoy.verify(env.trainer.train(_ANY), times=0)
+
+    assert len(list(corpus.get_rows())) == 8
+    decoy.verify(env.trainer.train(_ANY), times=1)
+
+
+def test_word_align_corpus_count_does_not_train(decoy: Decoy, monkeypatch: pytest.MonkeyPatch) -> None:
+    env = _TestEnvironment(decoy, monkeypatch)
+
+    assert word_align_corpus(create_test_parallel_corpus()).count() == 8
+    decoy.verify(env.trainer.train(_ANY), times=0)
+
+
+def test_word_align_corpus_closes_trained_model(decoy: Decoy, monkeypatch: pytest.MonkeyPatch) -> None:
+    env = _TestEnvironment(decoy, monkeypatch)
+
+    with word_align_corpus(create_test_parallel_corpus()).get_rows() as rows:
+        assert len(list(rows)) == 8
+    # Exiting the model is what closes it, per WordAlignmentModel.__exit__.
+    decoy.verify(env.model.__exit__(None, None, None), times=1)
+
+
+def test_word_align_corpus_closes_trained_model_on_early_exit(decoy: Decoy, monkeypatch: pytest.MonkeyPatch) -> None:
+    env = _TestEnvironment(decoy, monkeypatch)
+
+    with word_align_corpus(create_test_parallel_corpus()).get_rows() as rows:
+        next(rows)
+    decoy.verify(env.model.__exit__(_ANY, _ANY, _ANY), times=1)
+
+
+def test_word_align_corpus_closes_trained_model_through_chained_operator(
+    decoy: Decoy, monkeypatch: pytest.MonkeyPatch
+) -> None:
+    # Cleanup rides on the row generator, so it still happens when a wrapping operator stops early.
+    env = _TestEnvironment(decoy, monkeypatch)
+
+    assert len(list(word_align_corpus(create_test_parallel_corpus()).take(3))) == 3
+    decoy.verify(env.model.__exit__(_ANY, _ANY, _ANY), times=1)
+
+
+def test_word_align_corpus_trains_on_requested_text_ids_only(decoy: Decoy, monkeypatch: pytest.MonkeyPatch) -> None:
+    env = _TestEnvironment(decoy, monkeypatch)
+    corpus = word_align_corpus(_create_two_text_parallel_corpus())
+
+    rows = list(corpus.get_rows(["text2"]))
+    assert [row.text_id for row in rows] == ["text2", "text2"]
+
+    assert env.training_corpus is not None
+    assert [row.text_id for row in env.training_corpus.get_rows()] == ["text2", "text2"]
+
+
+def test_word_align_corpus_does_not_close_supplied_model(decoy: Decoy) -> None:
+    # A model the caller creates stays the caller's to close, so it can be reused afterward.
+    model = _create_mock_model(decoy)
+
+    assert len(list(word_align_corpus(create_test_parallel_corpus(), aligner=model).get_rows())) == 8
+    decoy.verify(model.__exit__(_ANY, _ANY, _ANY), times=0)