Pass the MosesPunctNormalizer to the new Pipeline on OOM

machine/translation/huggingface/hugging_face_nmt_engine.py

@@ -114,13 +114,7 @@ def __init__(
 
         self._oom_batch_size_backoff_mult = oom_batch_size_backoff_mult
 
-        self._pipeline = SilTranslationPipeline(
-            model=self._model,
-            tokenizer=self.tokenizer,
-            mpn=self._mpn,
-            batch_size=self._batch_size,
-            **self._pipeline_kwargs,
-        )
+        self._pipeline = self._create_pipeline()
 
     @property
     def tokenizer(self) -> PreTrainedTokenizer | PreTrainedTokenizerFast:
@@ -154,12 +148,7 @@ def translate_n_batch(
                     raise
                 self._batch_size = max(int(round(self._batch_size * self._oom_batch_size_backoff_mult)), 1)
                 logger.warning(f"Out of memory error caught. Reducing batch size to {self._batch_size} and retrying.")
-                self._pipeline = SilTranslationPipeline(
-                    model=self._model,
-                    tokenizer=self._tokenizer,
-                    batch_size=self._batch_size,
-                    **self._pipeline_kwargs,
-                )
+                self._pipeline = self._create_pipeline()
 
     def _try_translate_n_batch(
         self, n: int, segments: Sequence[Union[str, Sequence[str]]]
@@ -203,6 +192,15 @@ def close(self) -> None:
         with torch.no_grad():
             torch.cuda.empty_cache()
 
+    def _create_pipeline(self) -> SilTranslationPipeline:
+        return SilTranslationPipeline(
+            model=self._model,
+            tokenizer=self._tokenizer,
+            mpn=self._mpn,
+            batch_size=self._batch_size,
+            **self._pipeline_kwargs,
+        )
+
 
 class SilTranslationPipeline(TranslationPipeline):
     def __init__(

tests/testutils/hugging_face_nmt_test_helpers.py

@@ -0,0 +1,19 @@
+from typing import cast
+
+from transformers import M2M100Config, M2M100ForConditionalGeneration, NllbTokenizer, set_seed
+
+
+def create_tiny_nllb_model(model_dir: str) -> None:
+    # hf-internal-testing/tiny-random-nllb declares MBartTokenizer, which in transformers 5 discards the normalizer in
+    # tokenizer.json. Pair its tiny config with the real NLLB tokenizer so that input normalization is exercised.
+    tokenizer = cast(NllbTokenizer, NllbTokenizer.from_pretrained("facebook/nllb-200-distilled-600M"))
+    config = cast(M2M100Config, M2M100Config.from_pretrained("hf-internal-testing/tiny-random-nllb"))
+    config.vocab_size = len(tokenizer)
+    config.bos_token_id = cast(int, tokenizer.bos_token_id)
+    config.pad_token_id = cast(int, tokenizer.pad_token_id)
+    config.eos_token_id = cast(int, tokenizer.eos_token_id)
+    config.decoder_start_token_id = config.eos_token_id
+    set_seed(0)
+    model = M2M100ForConditionalGeneration(config)
+    model.save_pretrained(model_dir)
+    tokenizer.save_pretrained(model_dir)

tests/translation/huggingface/test_hugging_face_nmt_engine.py

@@ -1,14 +1,18 @@
 import sys
 
+from testutils.hugging_face_nmt_test_helpers import create_tiny_nllb_model
+
 if sys.platform == "darwin":
     from pytest import skip
 
     skip("skipping Hugging Face tests on MacOS", allow_module_level=True)
 
 from math import exp, log
+from tempfile import TemporaryDirectory
 from typing import Any, Dict, Optional
 
-from pytest import approx, mark, raises
+import torch  # pyright: ignore[reportMissingImports]
+from pytest import MonkeyPatch, approx, mark, raises
 from transformers import AutoModelForSeq2SeqLM, AutoTokenizer
 
 from machine.translation.huggingface import HuggingFaceNmtEngine, SilTranslationPipeline
@@ -53,6 +57,39 @@ def test_translate_n_batch_beam(output_attentions: bool) -> None:
         assert results[1][1].sequence_confidence == approx(_get_sequence_confidence(results[0][0]), 0.01)
 
 
+def test_translate_n_batch_oom_retains_mpn(monkeypatch: MonkeyPatch) -> None:
+    with TemporaryDirectory() as model_dir:
+        create_tiny_nllb_model(model_dir)
+        with HuggingFaceNmtEngine(
+            model_dir,
+            src_lang="eng_Latn",
+            tgt_lang="spa_Latn",
+            batch_size=2,
+            oom_batch_size_backoff_mult=0.5,
+            output_attentions=False,
+        ) as engine:
+            original_mpn = engine._mpn
+            assert original_mpn is not None
+            assert engine._pipeline._mpn is original_mpn
+
+            calls = 0
+
+            def fake_try_translate_n_batch(n: int, segments: Any) -> Any:
+                nonlocal calls
+                calls += 1
+                if calls == 1:
+                    raise torch.cuda.OutOfMemoryError
+                assert engine._pipeline._mpn is original_mpn
+                return [[] for _ in segments]
+
+            monkeypatch.setattr(engine, "_try_translate_n_batch", fake_try_translate_n_batch)
+
+            results = engine.translate_n_batch(1, ["This", "is", "a", "test", "."])
+            assert engine._pipeline._mpn is original_mpn
+            assert calls == 6  # 1 OOM call + 5 segments at batch size 1
+            assert results == [[], [], [], [], []]
+
+
 @mark.parametrize("output_attentions", [True, False])
 def test_translate_greedy(output_attentions: bool) -> None:
     with HuggingFaceNmtEngine(

tests/translation/huggingface/test_hugging_face_nmt_model_trainer.py

@@ -1,5 +1,7 @@
 import sys
 
+from testutils.hugging_face_nmt_test_helpers import create_tiny_nllb_model
+
 if sys.platform == "darwin":
     from pytest import skip
 
@@ -9,15 +11,12 @@
 from typing import cast
 
 from transformers import (
-    M2M100Config,
-    M2M100ForConditionalGeneration,
     M2M100Tokenizer,
     MBart50Tokenizer,
     MBartTokenizer,
     NllbTokenizer,
     PreTrainedTokenizerFast,
     Seq2SeqTrainingArguments,
-    set_seed,
 )
 
 from machine.corpora import DictionaryTextCorpus, MemoryText, TextRow
@@ -103,7 +102,7 @@ def test_train_non_empty_corpus() -> None:
 
 def test_update_tokenizer_missing_char() -> None:
     with TemporaryDirectory() as model_dir, TemporaryDirectory() as temp_dir:
-        _create_tiny_nllb_model(model_dir)
+        create_tiny_nllb_model(model_dir)
         source_corpus = DictionaryTextCorpus(
             [
                 MemoryText(
@@ -541,19 +540,3 @@ def test_m2m_100_tokenizer_add_lang_code() -> None:
 
 def _row(row_ref: int, text: str) -> TextRow:
     return TextRow("text1", row_ref, segment=[text])
-
-
-def _create_tiny_nllb_model(model_dir: str) -> None:
-    # hf-internal-testing/tiny-random-nllb declares MBartTokenizer, which in transformers 5 discards the normalizer in
-    # tokenizer.json. Pair its tiny config with the real NLLB tokenizer so that input normalization is exercised.
-    tokenizer = cast(NllbTokenizer, NllbTokenizer.from_pretrained("facebook/nllb-200-distilled-600M"))
-    config = cast(M2M100Config, M2M100Config.from_pretrained("hf-internal-testing/tiny-random-nllb"))
-    config.vocab_size = len(tokenizer)
-    config.bos_token_id = cast(int, tokenizer.bos_token_id)
-    config.pad_token_id = cast(int, tokenizer.pad_token_id)
-    config.eos_token_id = cast(int, tokenizer.eos_token_id)
-    config.decoder_start_token_id = config.eos_token_id
-    set_seed(0)
-    model = M2M100ForConditionalGeneration(config)
-    model.save_pretrained(model_dir)
-    tokenizer.save_pretrained(model_dir)