diff --git a/machine/translation/huggingface/hugging_face_nmt_engine.py b/machine/translation/huggingface/hugging_face_nmt_engine.py index 80219d6a..a407df5e 100644 --- a/machine/translation/huggingface/hugging_face_nmt_engine.py +++ b/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__( diff --git a/tests/testutils/hugging_face_nmt_test_helpers.py b/tests/testutils/hugging_face_nmt_test_helpers.py new file mode 100644 index 00000000..1732823a --- /dev/null +++ b/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) diff --git a/tests/translation/huggingface/test_hugging_face_nmt_engine.py b/tests/translation/huggingface/test_hugging_face_nmt_engine.py index d1d34be4..96793917 100644 --- a/tests/translation/huggingface/test_hugging_face_nmt_engine.py +++ b/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( diff --git a/tests/translation/huggingface/test_hugging_face_nmt_model_trainer.py b/tests/translation/huggingface/test_hugging_face_nmt_model_trainer.py index cbbc95e7..602357c7 100644 --- a/tests/translation/huggingface/test_hugging_face_nmt_model_trainer.py +++ b/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)