From de98a38da1fafc8c23b806b001a51dd5b80fcaf4 Mon Sep 17 00:00:00 2001 From: Damien Daspit Date: Thu, 6 Aug 2026 11:04:26 -0400 Subject: [PATCH] Make `word_align_corpus` train lazily - Properly cleans up the created model --- machine/translation/corpus_ops.py | 114 ++++++++++++++-------- samples/word_alignment.ipynb | 4 +- tests/translation/test_corpus_ops.py | 139 +++++++++++++++++++++++---- 3 files changed, 199 insertions(+), 58 deletions(-) diff --git a/machine/translation/corpus_ops.py b/machine/translation/corpus_ops.py index f628dc90..35b67300 100644 --- a/machine/translation/corpus_ops.py +++ b/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 diff --git a/samples/word_alignment.ipynb b/samples/word_alignment.ipynb index 67a07ccb..7d6475bc 100644 --- a/samples/word_alignment.ipynb +++ b/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." ] }, { diff --git a/tests/translation/test_corpus_ops.py b/tests/translation/test_corpus_ops.py index 6a3a69d4..a32d3a68 100644 --- a/tests/translation/test_corpus_ops.py +++ b/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,10 +71,13 @@ 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 @@ -73,3 +85,96 @@ 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)