From 23f3376625b61e4a9ab1d47388d74f78db81915d Mon Sep 17 00:00:00 2001 From: "copilot-swe-agent[bot]" <198982749+Copilot@users.noreply.github.com> Date: Wed, 20 May 2026 17:20:51 +0000 Subject: [PATCH 01/16] Initial plan From 1623e4bb8485776e2ad4f00a524ace3bb836cbb7 Mon Sep 17 00:00:00 2001 From: "copilot-swe-agent[bot]" <198982749+Copilot@users.noreply.github.com> Date: Wed, 20 May 2026 17:26:23 +0000 Subject: [PATCH 02/16] fix: cap tensor inference batch size to reduce HF OOM risk Agent-Logs-Url: https://github.com/sillsdev/silnlp/sessions/1d43fe05-c924-4848-872d-3ad357553aef Co-authored-by: benjaminking <1214233+benjaminking@users.noreply.github.com> --- silnlp/nmt/hugging_face_config.py | 4 ++++ tests/smoke_tests/test_experiment.py | 28 +++++++++++++++++++++++++++- 2 files changed, 31 insertions(+), 1 deletion(-) diff --git a/silnlp/nmt/hugging_face_config.py b/silnlp/nmt/hugging_face_config.py index 0399de9c..448759e0 100644 --- a/silnlp/nmt/hugging_face_config.py +++ b/silnlp/nmt/hugging_face_config.py @@ -396,6 +396,7 @@ def __init__(self, exp_dir: Path, config: dict) -> None: }, "infer": { "infer_batch_size": 16, + "infer_batch_size_with_tensors": 1, "num_beams": 2, "num_drafts": 3, "multiple_translations_method": "hybrid", @@ -1489,6 +1490,9 @@ def _translate_sentences( return_tensors: bool = False, ) -> Iterable[ModelOutputGroup]: batch_size: int = self._config.infer["infer_batch_size"] + if return_tensors: + batch_size_with_tensors = max(1, int(self._config.infer.get("infer_batch_size_with_tensors", 1))) + batch_size = min(batch_size, batch_size_with_tensors) dictionary = self._get_dictionary() if vrefs is None or len(dictionary) == 0: diff --git a/tests/smoke_tests/test_experiment.py b/tests/smoke_tests/test_experiment.py index 1b8b76c7..49c6bcb7 100644 --- a/tests/smoke_tests/test_experiment.py +++ b/tests/smoke_tests/test_experiment.py @@ -1,13 +1,14 @@ import shutil from pathlib import Path from typing import cast +from unittest.mock import Mock import torch from silnlp.common.environment import SilNlpEnv from silnlp.nmt.config_utils import load_config_from_exp_dir from silnlp.nmt.experiment import SILExperiment -from silnlp.nmt.hugging_face_config import HuggingFaceConfig +from silnlp.nmt.hugging_face_config import HuggingFaceConfig, HuggingFaceNMTModel from tests.smoke_tests.mock_pretrained_model import ( MockModelOutput, MockPreTrainedModelProviderFactory, @@ -32,6 +33,31 @@ def test_experiment_full_pipeline(): clean_experiment_directory(environment.get_mt_exp_dir(EXPERIMENT_NAME)) +def test_translate_sentences_uses_tensor_batch_size_cap(): + model = cast(HuggingFaceNMTModel, HuggingFaceNMTModel.__new__(HuggingFaceNMTModel)) + model._config = cast(HuggingFaceConfig, Mock(infer={"infer_batch_size": 16, "infer_batch_size_with_tensors": 2})) + model._get_dictionary = Mock(return_value={}) + captured_batch_size = {} + + def fake_translate_sentence_helper(*args, **kwargs): + captured_batch_size["value"] = args[2] + return iter(()) + + model._translate_sentence_helper = fake_translate_sentence_helper + + list( + model._translate_sentences( + tokenizer=Mock(), + pipeline=Mock(), + sentences=[["token"]], + vrefs=None, + return_tensors=True, + ) + ) + + assert captured_batch_size["value"] == 2 + + def set_up_environment() -> SilNlpEnv: # To avoid keeping large numbers of files in the repository, the tests assume that there is an # active connection to the MinIO bucket and use file from the "Scripture" and "Paratext" directories. From d122e90fb7292a087139e92869e5dc320771775c Mon Sep 17 00:00:00 2001 From: "copilot-swe-agent[bot]" <198982749+Copilot@users.noreply.github.com> Date: Wed, 20 May 2026 17:27:39 +0000 Subject: [PATCH 03/16] test: harden tensor batch-size cap smoke test setup Agent-Logs-Url: https://github.com/sillsdev/silnlp/sessions/1d43fe05-c924-4848-872d-3ad357553aef Co-authored-by: benjaminking <1214233+benjaminking@users.noreply.github.com> --- tests/smoke_tests/test_experiment.py | 11 +++++++---- 1 file changed, 7 insertions(+), 4 deletions(-) diff --git a/tests/smoke_tests/test_experiment.py b/tests/smoke_tests/test_experiment.py index 49c6bcb7..5c34937c 100644 --- a/tests/smoke_tests/test_experiment.py +++ b/tests/smoke_tests/test_experiment.py @@ -34,19 +34,22 @@ def test_experiment_full_pipeline(): def test_translate_sentences_uses_tensor_batch_size_cap(): - model = cast(HuggingFaceNMTModel, HuggingFaceNMTModel.__new__(HuggingFaceNMTModel)) + model = cast(HuggingFaceNMTModel, Mock(spec=HuggingFaceNMTModel)) model._config = cast(HuggingFaceConfig, Mock(infer={"infer_batch_size": 16, "infer_batch_size_with_tensors": 2})) model._get_dictionary = Mock(return_value={}) captured_batch_size = {} - def fake_translate_sentence_helper(*args, **kwargs): - captured_batch_size["value"] = args[2] + def fake_translate_sentence_helper( + pipeline, sentences, batch_size, return_tensors, force_words_ids=None, produce_multiple_translations=False + ): + captured_batch_size["value"] = batch_size return iter(()) model._translate_sentence_helper = fake_translate_sentence_helper list( - model._translate_sentences( + HuggingFaceNMTModel._translate_sentences( + model, tokenizer=Mock(), pipeline=Mock(), sentences=[["token"]], From e5b944dbad183199cb1e44a0c7fbb746bab84dd0 Mon Sep 17 00:00:00 2001 From: "copilot-swe-agent[bot]" <198982749+Copilot@users.noreply.github.com> Date: Wed, 20 May 2026 17:28:32 +0000 Subject: [PATCH 04/16] refactor: simplify tensor batch size cap calculation Agent-Logs-Url: https://github.com/sillsdev/silnlp/sessions/1d43fe05-c924-4848-872d-3ad357553aef Co-authored-by: benjaminking <1214233+benjaminking@users.noreply.github.com> --- silnlp/nmt/hugging_face_config.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/silnlp/nmt/hugging_face_config.py b/silnlp/nmt/hugging_face_config.py index 448759e0..e3cc486c 100644 --- a/silnlp/nmt/hugging_face_config.py +++ b/silnlp/nmt/hugging_face_config.py @@ -1491,7 +1491,7 @@ def _translate_sentences( ) -> Iterable[ModelOutputGroup]: batch_size: int = self._config.infer["infer_batch_size"] if return_tensors: - batch_size_with_tensors = max(1, int(self._config.infer.get("infer_batch_size_with_tensors", 1))) + batch_size_with_tensors = int(self._config.infer.get("infer_batch_size_with_tensors", 1)) batch_size = min(batch_size, batch_size_with_tensors) dictionary = self._get_dictionary() From 27fa7d448813650805f9984527dbd48111096acd Mon Sep 17 00:00:00 2001 From: "copilot-swe-agent[bot]" <198982749+Copilot@users.noreply.github.com> Date: Wed, 20 May 2026 17:29:22 +0000 Subject: [PATCH 05/16] refactor: remove unnecessary cast in tensor batch cap logic Agent-Logs-Url: https://github.com/sillsdev/silnlp/sessions/1d43fe05-c924-4848-872d-3ad357553aef Co-authored-by: benjaminking <1214233+benjaminking@users.noreply.github.com> --- silnlp/nmt/hugging_face_config.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/silnlp/nmt/hugging_face_config.py b/silnlp/nmt/hugging_face_config.py index e3cc486c..95ae5575 100644 --- a/silnlp/nmt/hugging_face_config.py +++ b/silnlp/nmt/hugging_face_config.py @@ -1491,7 +1491,7 @@ def _translate_sentences( ) -> Iterable[ModelOutputGroup]: batch_size: int = self._config.infer["infer_batch_size"] if return_tensors: - batch_size_with_tensors = int(self._config.infer.get("infer_batch_size_with_tensors", 1)) + batch_size_with_tensors = self._config.infer.get("infer_batch_size_with_tensors", 1) batch_size = min(batch_size, batch_size_with_tensors) dictionary = self._get_dictionary() From a92abd1d891bb64804bbd05f2d49e67f02dabcc8 Mon Sep 17 00:00:00 2001 From: "copilot-swe-agent[bot]" <198982749+Copilot@users.noreply.github.com> Date: Wed, 20 May 2026 17:55:43 +0000 Subject: [PATCH 06/16] feat: add adaptive OOM batch retry for translation inference Agent-Logs-Url: https://github.com/sillsdev/silnlp/sessions/318d55d1-ef9f-41ea-962b-48430ef1e230 Co-authored-by: benjaminking <1214233+benjaminking@users.noreply.github.com> --- silnlp/nmt/hugging_face_config.py | 159 ++++++++++++++++++--------- tests/smoke_tests/test_experiment.py | 22 ++++ 2 files changed, 128 insertions(+), 53 deletions(-) diff --git a/silnlp/nmt/hugging_face_config.py b/silnlp/nmt/hugging_face_config.py index 95ae5575..be5957b5 100644 --- a/silnlp/nmt/hugging_face_config.py +++ b/silnlp/nmt/hugging_face_config.py @@ -396,7 +396,7 @@ def __init__(self, exp_dir: Path, config: dict) -> None: }, "infer": { "infer_batch_size": 16, - "infer_batch_size_with_tensors": 1, + "infer_batch_size_with_tensors": 16, "num_beams": 2, "num_drafts": 3, "multiple_translations_method": "hybrid", @@ -1491,34 +1491,25 @@ def _translate_sentences( ) -> Iterable[ModelOutputGroup]: batch_size: int = self._config.infer["infer_batch_size"] if return_tensors: - batch_size_with_tensors = self._config.infer.get("infer_batch_size_with_tensors", 1) + batch_size_with_tensors = self._config.infer.get("infer_batch_size_with_tensors", batch_size) batch_size = min(batch_size, batch_size_with_tensors) dictionary = self._get_dictionary() - if vrefs is None or len(dictionary) == 0: + for batch, force_words in batch_sentences(sentences, vrefs, batch_size, dictionary): + if force_words is None: + force_words_ids = None + else: + force_words_ids = [[tokenizer.convert_tokens_to_ids(v) for v in vs] for vs in force_words] + force_words_ids = prune_sublists(force_words_ids) + yield from self._translate_sentence_helper( pipeline, - sentences, + batch, batch_size, return_tensors, + force_words_ids, produce_multiple_translations=produce_multiple_translations, ) - else: - for batch, force_words in batch_sentences(sentences, vrefs, batch_size, dictionary): - if force_words is None: - force_words_ids = None - else: - force_words_ids = [[tokenizer.convert_tokens_to_ids(v) for v in vs] for vs in force_words] - force_words_ids = prune_sublists(force_words_ids) - - yield from self._translate_sentence_helper( - pipeline, - batch, - batch_size, - return_tensors, - force_words_ids, - produce_multiple_translations=produce_multiple_translations, - ) def _translate_sentence_helper( self, @@ -1529,30 +1520,38 @@ def _translate_sentence_helper( force_words_ids: List[List[List[int]]] = None, produce_multiple_translations: bool = False, ) -> Iterable[ModelOutputGroup]: - + sentences = list(sentences) num_drafts = self.get_num_drafts() if produce_multiple_translations and num_drafts > 1: multiple_translations_method: str = self._config.infer.get("multiple_translations_method") - sentences = list(sentences) - if multiple_translations_method == "hybrid": - beam_search_results: List[dict] = self._translate_with_beam_search( - pipeline, + beam_search_results: List[dict] = self._translate_with_adaptive_batch_size( + lambda batch, current_batch_size, batch_force_words_ids: self._translate_with_beam_search( + pipeline, + batch, + current_batch_size, + return_tensors, + num_return_sequences=1, + force_words_ids=batch_force_words_ids, + ), sentences, batch_size, - return_tensors, - num_return_sequences=1, - force_words_ids=force_words_ids, + force_words_ids, ) - sampling_results: List[dict] = self._translate_with_sampling( - pipeline, + sampling_results: List[dict] = self._translate_with_adaptive_batch_size( + lambda batch, current_batch_size, batch_force_words_ids: self._translate_with_sampling( + pipeline, + batch, + current_batch_size, + return_tensors, + num_return_sequences=num_drafts - 1, + force_words_ids=batch_force_words_ids, + ), sentences, batch_size, - return_tensors, - num_return_sequences=num_drafts - 1, - force_words_ids=force_words_ids, + force_words_ids, ) # concatenate the beam search results with the sampling results @@ -1564,39 +1563,54 @@ def _translate_sentence_helper( elif multiple_translations_method == "sampling": yield from [ ModelOutputGroup(result) - for result in self._translate_with_sampling( - pipeline, + for result in self._translate_with_adaptive_batch_size( + lambda batch, current_batch_size, batch_force_words_ids: self._translate_with_sampling( + pipeline, + batch, + current_batch_size, + return_tensors, + num_return_sequences=num_drafts, + force_words_ids=batch_force_words_ids, + ), sentences, batch_size, - return_tensors, - num_return_sequences=num_drafts, - force_words_ids=force_words_ids, + force_words_ids, ) ] elif multiple_translations_method == "beam_search": yield from [ ModelOutputGroup(result) - for result in self._translate_with_beam_search( - pipeline, + for result in self._translate_with_adaptive_batch_size( + lambda batch, current_batch_size, batch_force_words_ids: self._translate_with_beam_search( + pipeline, + batch, + current_batch_size, + return_tensors, + num_return_sequences=num_drafts, + force_words_ids=batch_force_words_ids, + ), sentences, batch_size, - return_tensors, - num_return_sequences=num_drafts, - force_words_ids=force_words_ids, + force_words_ids, ) ] elif multiple_translations_method == "diverse_beam_search": yield from [ ModelOutputGroup(result) - for result in self._translate_with_diverse_beam_search( - pipeline, + for result in self._translate_with_adaptive_batch_size( + lambda batch, current_batch_size, batch_force_words_ids: self._translate_with_diverse_beam_search( + pipeline, + batch, + current_batch_size, + return_tensors, + num_return_sequences=num_drafts, + force_words_ids=batch_force_words_ids, + ), sentences, batch_size, - return_tensors, - num_return_sequences=num_drafts, - force_words_ids=force_words_ids, + force_words_ids, ) ] else: @@ -1605,16 +1619,55 @@ def _translate_sentence_helper( else: yield from [ ModelOutputGroup([translated_sentence[0]]) - for translated_sentence in self._translate_with_beam_search( - pipeline, + for translated_sentence in self._translate_with_adaptive_batch_size( + lambda batch, current_batch_size, batch_force_words_ids: self._translate_with_beam_search( + pipeline, + batch, + current_batch_size, + return_tensors, + num_return_sequences=1, + force_words_ids=batch_force_words_ids, + ), sentences, batch_size, - return_tensors, - num_return_sequences=1, - force_words_ids=force_words_ids, + force_words_ids, ) ] + def _translate_with_adaptive_batch_size( + self, + translate: Callable[[List[TSent], int, Optional[List[List[List[int]]]]], List[List[dict]]], + sentences: List[TSent], + batch_size: int, + force_words_ids: Optional[List[List[List[int]]]] = None, + ) -> List[List[dict]]: + current_batch_size = batch_size + translated_sentences: List[List[dict]] = [] + index = 0 + + while index < len(sentences): + current_batch_size = min(current_batch_size, len(sentences) - index) + batch_sentences = sentences[index : index + current_batch_size] + batch_force_words_ids = ( + force_words_ids[index : index + current_batch_size] if force_words_ids is not None else None + ) + try: + translated_sentences.extend(translate(batch_sentences, current_batch_size, batch_force_words_ids)) + index += current_batch_size + except Exception as e: + if not _should_reduce_batch_size(e) or current_batch_size <= 1: + raise + current_batch_size //= 2 + LOGGER.warning( + "OOM during translation inference; reducing batch size to %d and retrying current batch.", + current_batch_size, + ) + gc.collect() + if torch.cuda.is_available(): + torch.cuda.empty_cache() + + return translated_sentences + # When translating tokenized sentences, for some reason the Huggingface pipeline # returns List[List[dict]] instead of List[dict]. Each nested list is a # singleton. This function flattens the structure. diff --git a/tests/smoke_tests/test_experiment.py b/tests/smoke_tests/test_experiment.py index 5c34937c..10d8c431 100644 --- a/tests/smoke_tests/test_experiment.py +++ b/tests/smoke_tests/test_experiment.py @@ -61,6 +61,28 @@ def fake_translate_sentence_helper( assert captured_batch_size["value"] == 2 +def test_translate_with_adaptive_batch_size_retries_with_smaller_batch(): + model = cast(HuggingFaceNMTModel, Mock(spec=HuggingFaceNMTModel)) + call_batch_sizes: list[int] = [] + + def fake_translate(batch_sentences, batch_size, batch_force_words_ids): + call_batch_sizes.append(batch_size) + if batch_size > 2: + raise RuntimeError("CUDA out of memory. Tried to allocate 7.00 GiB.") + return [[{"translation_text": str(sentence[0])}] for sentence in batch_sentences] + + sentences = [["a"], ["b"], ["c"], ["d"]] + translated = HuggingFaceNMTModel._translate_with_adaptive_batch_size( + model, + fake_translate, + sentences, + batch_size=4, + ) + + assert [translation[0]["translation_text"] for translation in translated] == ["a", "b", "c", "d"] + assert call_batch_sizes == [4, 2, 2] + + def set_up_environment() -> SilNlpEnv: # To avoid keeping large numbers of files in the repository, the tests assume that there is an # active connection to the MinIO bucket and use file from the "Scripture" and "Paratext" directories. From a22b5e8825113d31f72559e97654009743a3be99 Mon Sep 17 00:00:00 2001 From: "copilot-swe-agent[bot]" <198982749+Copilot@users.noreply.github.com> Date: Wed, 20 May 2026 17:57:03 +0000 Subject: [PATCH 07/16] refactor: tighten adaptive OOM handler and test constant Agent-Logs-Url: https://github.com/sillsdev/silnlp/sessions/318d55d1-ef9f-41ea-962b-48430ef1e230 Co-authored-by: benjaminking <1214233+benjaminking@users.noreply.github.com> --- silnlp/nmt/hugging_face_config.py | 3 ++- tests/smoke_tests/test_experiment.py | 3 ++- 2 files changed, 4 insertions(+), 2 deletions(-) diff --git a/silnlp/nmt/hugging_face_config.py b/silnlp/nmt/hugging_face_config.py index be5957b5..0e272c4d 100644 --- a/silnlp/nmt/hugging_face_config.py +++ b/silnlp/nmt/hugging_face_config.py @@ -1641,6 +1641,7 @@ def _translate_with_adaptive_batch_size( batch_size: int, force_words_ids: Optional[List[List[List[int]]]] = None, ) -> List[List[dict]]: + """Translate input sentences and reduce batch size on OOM until the current batch succeeds.""" current_batch_size = batch_size translated_sentences: List[List[dict]] = [] index = 0 @@ -1654,7 +1655,7 @@ def _translate_with_adaptive_batch_size( try: translated_sentences.extend(translate(batch_sentences, current_batch_size, batch_force_words_ids)) index += current_batch_size - except Exception as e: + except (RuntimeError, torch.OutOfMemoryError) as e: if not _should_reduce_batch_size(e) or current_batch_size <= 1: raise current_batch_size //= 2 diff --git a/tests/smoke_tests/test_experiment.py b/tests/smoke_tests/test_experiment.py index 10d8c431..18fb906c 100644 --- a/tests/smoke_tests/test_experiment.py +++ b/tests/smoke_tests/test_experiment.py @@ -17,6 +17,7 @@ TEST_MT_DIR = Path(__file__).parent EXPERIMENT_NAME = "test_experiment" +OOM_ERROR_MESSAGE = "CUDA out of memory. Tried to allocate 7.00 GiB." def test_experiment_full_pipeline(): @@ -68,7 +69,7 @@ def test_translate_with_adaptive_batch_size_retries_with_smaller_batch(): def fake_translate(batch_sentences, batch_size, batch_force_words_ids): call_batch_sizes.append(batch_size) if batch_size > 2: - raise RuntimeError("CUDA out of memory. Tried to allocate 7.00 GiB.") + raise RuntimeError(OOM_ERROR_MESSAGE) return [[{"translation_text": str(sentence[0])}] for sentence in batch_sentences] sentences = [["a"], ["b"], ["c"], ["d"]] From 60a3d3e7a26ef1a91ece60df39708c018e610c60 Mon Sep 17 00:00:00 2001 From: "copilot-swe-agent[bot]" <198982749+Copilot@users.noreply.github.com> Date: Wed, 20 May 2026 17:58:06 +0000 Subject: [PATCH 08/16] docs: clarify adaptive translation batch retry helper contract Agent-Logs-Url: https://github.com/sillsdev/silnlp/sessions/318d55d1-ef9f-41ea-962b-48430ef1e230 Co-authored-by: benjaminking <1214233+benjaminking@users.noreply.github.com> --- silnlp/nmt/hugging_face_config.py | 12 +++++++++++- 1 file changed, 11 insertions(+), 1 deletion(-) diff --git a/silnlp/nmt/hugging_face_config.py b/silnlp/nmt/hugging_face_config.py index 0e272c4d..6783fde5 100644 --- a/silnlp/nmt/hugging_face_config.py +++ b/silnlp/nmt/hugging_face_config.py @@ -1641,7 +1641,17 @@ def _translate_with_adaptive_batch_size( batch_size: int, force_words_ids: Optional[List[List[List[int]]]] = None, ) -> List[List[dict]]: - """Translate input sentences and reduce batch size on OOM until the current batch succeeds.""" + """Translate sentences in batches, reducing batch size and retrying when OOM occurs. + + Args: + translate: Translation callback that accepts a sentence batch, batch size, and optional force words IDs. + sentences: Input sentences to translate. + batch_size: Initial batch size to use. + force_words_ids: Optional force words IDs aligned with ``sentences``. + + Returns: + A list of translated outputs aligned with the input sentences. + """ current_batch_size = batch_size translated_sentences: List[List[dict]] = [] index = 0 From 0658859ab7c2aa4d78343ab16685dd00cbce1024 Mon Sep 17 00:00:00 2001 From: "copilot-swe-agent[bot]" <198982749+Copilot@users.noreply.github.com> Date: Wed, 20 May 2026 17:59:11 +0000 Subject: [PATCH 09/16] refactor: handle adaptive retry OOM through RuntimeError path Agent-Logs-Url: https://github.com/sillsdev/silnlp/sessions/318d55d1-ef9f-41ea-962b-48430ef1e230 Co-authored-by: benjaminking <1214233+benjaminking@users.noreply.github.com> --- silnlp/nmt/hugging_face_config.py | 6 ++++-- 1 file changed, 4 insertions(+), 2 deletions(-) diff --git a/silnlp/nmt/hugging_face_config.py b/silnlp/nmt/hugging_face_config.py index 6783fde5..2de81eda 100644 --- a/silnlp/nmt/hugging_face_config.py +++ b/silnlp/nmt/hugging_face_config.py @@ -1644,7 +1644,9 @@ def _translate_with_adaptive_batch_size( """Translate sentences in batches, reducing batch size and retrying when OOM occurs. Args: - translate: Translation callback that accepts a sentence batch, batch size, and optional force words IDs. + translate: Translation callback with signature + ``(batch_sentences: List[TSent], batch_size: int, batch_force_words_ids: Optional[List[List[List[int]]]])`` + ``-> List[List[dict]]``. sentences: Input sentences to translate. batch_size: Initial batch size to use. force_words_ids: Optional force words IDs aligned with ``sentences``. @@ -1665,7 +1667,7 @@ def _translate_with_adaptive_batch_size( try: translated_sentences.extend(translate(batch_sentences, current_batch_size, batch_force_words_ids)) index += current_batch_size - except (RuntimeError, torch.OutOfMemoryError) as e: + except RuntimeError as e: if not _should_reduce_batch_size(e) or current_batch_size <= 1: raise current_batch_size //= 2 From 27d0b448c4c151131434db6c16e7db0770c89705 Mon Sep 17 00:00:00 2001 From: "copilot-swe-agent[bot]" <198982749+Copilot@users.noreply.github.com> Date: Wed, 20 May 2026 18:15:07 +0000 Subject: [PATCH 10/16] refactor: move adaptive OOM retry into each translate method Agent-Logs-Url: https://github.com/sillsdev/silnlp/sessions/ee62a80c-fe7c-4ffd-b732-fda53446b7c4 Co-authored-by: benjaminking <1214233+benjaminking@users.noreply.github.com> --- silnlp/nmt/hugging_face_config.py | 282 +++++++++++++-------------- tests/smoke_tests/test_experiment.py | 10 +- 2 files changed, 143 insertions(+), 149 deletions(-) diff --git a/silnlp/nmt/hugging_face_config.py b/silnlp/nmt/hugging_face_config.py index 2de81eda..052d8abd 100644 --- a/silnlp/nmt/hugging_face_config.py +++ b/silnlp/nmt/hugging_face_config.py @@ -1526,32 +1526,22 @@ def _translate_sentence_helper( multiple_translations_method: str = self._config.infer.get("multiple_translations_method") if multiple_translations_method == "hybrid": - beam_search_results: List[dict] = self._translate_with_adaptive_batch_size( - lambda batch, current_batch_size, batch_force_words_ids: self._translate_with_beam_search( - pipeline, - batch, - current_batch_size, - return_tensors, - num_return_sequences=1, - force_words_ids=batch_force_words_ids, - ), + beam_search_results: List[dict] = self._translate_with_beam_search( + pipeline, sentences, batch_size, - force_words_ids, + return_tensors, + num_return_sequences=1, + force_words_ids=force_words_ids, ) - sampling_results: List[dict] = self._translate_with_adaptive_batch_size( - lambda batch, current_batch_size, batch_force_words_ids: self._translate_with_sampling( - pipeline, - batch, - current_batch_size, - return_tensors, - num_return_sequences=num_drafts - 1, - force_words_ids=batch_force_words_ids, - ), + sampling_results: List[dict] = self._translate_with_sampling( + pipeline, sentences, batch_size, - force_words_ids, + return_tensors, + num_return_sequences=num_drafts - 1, + force_words_ids=force_words_ids, ) # concatenate the beam search results with the sampling results @@ -1563,54 +1553,39 @@ def _translate_sentence_helper( elif multiple_translations_method == "sampling": yield from [ ModelOutputGroup(result) - for result in self._translate_with_adaptive_batch_size( - lambda batch, current_batch_size, batch_force_words_ids: self._translate_with_sampling( - pipeline, - batch, - current_batch_size, - return_tensors, - num_return_sequences=num_drafts, - force_words_ids=batch_force_words_ids, - ), + for result in self._translate_with_sampling( + pipeline, sentences, batch_size, - force_words_ids, + return_tensors, + num_return_sequences=num_drafts, + force_words_ids=force_words_ids, ) ] elif multiple_translations_method == "beam_search": yield from [ ModelOutputGroup(result) - for result in self._translate_with_adaptive_batch_size( - lambda batch, current_batch_size, batch_force_words_ids: self._translate_with_beam_search( - pipeline, - batch, - current_batch_size, - return_tensors, - num_return_sequences=num_drafts, - force_words_ids=batch_force_words_ids, - ), + for result in self._translate_with_beam_search( + pipeline, sentences, batch_size, - force_words_ids, + return_tensors, + num_return_sequences=num_drafts, + force_words_ids=force_words_ids, ) ] elif multiple_translations_method == "diverse_beam_search": yield from [ ModelOutputGroup(result) - for result in self._translate_with_adaptive_batch_size( - lambda batch, current_batch_size, batch_force_words_ids: self._translate_with_diverse_beam_search( - pipeline, - batch, - current_batch_size, - return_tensors, - num_return_sequences=num_drafts, - force_words_ids=batch_force_words_ids, - ), + for result in self._translate_with_diverse_beam_search( + pipeline, sentences, batch_size, - force_words_ids, + return_tensors, + num_return_sequences=num_drafts, + force_words_ids=force_words_ids, ) ] else: @@ -1619,45 +1594,39 @@ def _translate_sentence_helper( else: yield from [ ModelOutputGroup([translated_sentence[0]]) - for translated_sentence in self._translate_with_adaptive_batch_size( - lambda batch, current_batch_size, batch_force_words_ids: self._translate_with_beam_search( - pipeline, - batch, - current_batch_size, - return_tensors, - num_return_sequences=1, - force_words_ids=batch_force_words_ids, - ), + for translated_sentence in self._translate_with_beam_search( + pipeline, sentences, batch_size, - force_words_ids, + return_tensors, + num_return_sequences=1, + force_words_ids=force_words_ids, ) ] - def _translate_with_adaptive_batch_size( + # When translating tokenized sentences, for some reason the Huggingface pipeline + # returns List[List[dict]] instead of List[dict]. Each nested list is a + # singleton. This function flattens the structure. + def _flatten_tokenized_translations(self, pipeline_output) -> List[dict]: + return [[i if isinstance(i, dict) else i[0] for i in translation] for translation in pipeline_output] + + def _translate_with_beam_search( self, - translate: Callable[[List[TSent], int, Optional[List[List[List[int]]]]], List[List[dict]]], - sentences: List[TSent], + pipeline: TranslationPipeline, + sentences: Iterable[TSent], batch_size: int, - force_words_ids: Optional[List[List[List[int]]]] = None, + return_tensors: bool, + num_return_sequences: int = 1, + force_words_ids: List[List[List[int]]] = None, ) -> List[List[dict]]: - """Translate sentences in batches, reducing batch size and retrying when OOM occurs. - - Args: - translate: Translation callback with signature - ``(batch_sentences: List[TSent], batch_size: int, batch_force_words_ids: Optional[List[List[List[int]]]])`` - ``-> List[List[dict]]``. - sentences: Input sentences to translate. - batch_size: Initial batch size to use. - force_words_ids: Optional force words IDs aligned with ``sentences``. - - Returns: - A list of translated outputs aligned with the input sentences. - """ + num_beams: Optional[int] = self._config.infer.get("num_beams") + if num_beams is None: + num_beams = self._config.params.get("generation_num_beams") + + sentences = list(sentences) current_batch_size = batch_size - translated_sentences: List[List[dict]] = [] + translations: List[List[dict]] = [] index = 0 - while index < len(sentences): current_batch_size = min(current_batch_size, len(sentences) - index) batch_sentences = sentences[index : index + current_batch_size] @@ -1665,55 +1634,32 @@ def _translate_with_adaptive_batch_size( force_words_ids[index : index + current_batch_size] if force_words_ids is not None else None ) try: - translated_sentences.extend(translate(batch_sentences, current_batch_size, batch_force_words_ids)) + batch_translations = pipeline( + batch_sentences, + num_beams=num_beams, + num_return_sequences=num_return_sequences, + force_words_ids=batch_force_words_ids, + batch_size=current_batch_size, + return_text=not return_tensors, + return_tensors=return_tensors, + ) + if num_return_sequences == 1: + batch_translations = [[t] for t in batch_translations] + translations.extend(self._flatten_tokenized_translations(batch_translations)) index += current_batch_size except RuntimeError as e: if not _should_reduce_batch_size(e) or current_batch_size <= 1: raise current_batch_size //= 2 LOGGER.warning( - "OOM during translation inference; reducing batch size to %d and retrying current batch.", + "OOM during beam-search translation inference; reducing batch size to %d and retrying current batch.", current_batch_size, ) gc.collect() if torch.cuda.is_available(): torch.cuda.empty_cache() - return translated_sentences - - # When translating tokenized sentences, for some reason the Huggingface pipeline - # returns List[List[dict]] instead of List[dict]. Each nested list is a - # singleton. This function flattens the structure. - def _flatten_tokenized_translations(self, pipeline_output) -> List[dict]: - return [[i if isinstance(i, dict) else i[0] for i in translation] for translation in pipeline_output] - - def _translate_with_beam_search( - self, - pipeline: TranslationPipeline, - sentences: Iterable[TSent], - batch_size: int, - return_tensors: bool, - num_return_sequences: int = 1, - force_words_ids: List[List[List[int]]] = None, - ) -> List[List[dict]]: - num_beams: Optional[int] = self._config.infer.get("num_beams") - if num_beams is None: - num_beams = self._config.params.get("generation_num_beams") - - translations = pipeline( - sentences, - num_beams=num_beams, - num_return_sequences=num_return_sequences, - force_words_ids=force_words_ids, - batch_size=batch_size, - return_text=not return_tensors, - return_tensors=return_tensors, - ) - - if num_return_sequences == 1: - translations = [[t] for t in translations] - - return self._flatten_tokenized_translations(translations) + return translations def _translate_with_sampling( self, @@ -1727,21 +1673,44 @@ def _translate_with_sampling( temperature: Optional[int] = self._config.infer.get("temperature") - translations = pipeline( - sentences, - do_sample=True, - temperature=temperature, - num_return_sequences=num_return_sequences, - force_words_ids=force_words_ids, - batch_size=batch_size, - return_text=not return_tensors, - return_tensors=return_tensors, - ) - - if num_return_sequences == 1: - translations = [[t] for t in translations] + sentences = list(sentences) + current_batch_size = batch_size + translations: List[List[dict]] = [] + index = 0 + while index < len(sentences): + current_batch_size = min(current_batch_size, len(sentences) - index) + batch_sentences = sentences[index : index + current_batch_size] + batch_force_words_ids = ( + force_words_ids[index : index + current_batch_size] if force_words_ids is not None else None + ) + try: + batch_translations = pipeline( + batch_sentences, + do_sample=True, + temperature=temperature, + num_return_sequences=num_return_sequences, + force_words_ids=batch_force_words_ids, + batch_size=current_batch_size, + return_text=not return_tensors, + return_tensors=return_tensors, + ) + if num_return_sequences == 1: + batch_translations = [[t] for t in batch_translations] + translations.extend(self._flatten_tokenized_translations(batch_translations)) + index += current_batch_size + except RuntimeError as e: + if not _should_reduce_batch_size(e) or current_batch_size <= 1: + raise + current_batch_size //= 2 + LOGGER.warning( + "OOM during sampling translation inference; reducing batch size to %d and retrying current batch.", + current_batch_size, + ) + gc.collect() + if torch.cuda.is_available(): + torch.cuda.empty_cache() - return self._flatten_tokenized_translations(translations) + return translations def _translate_with_diverse_beam_search( self, @@ -1757,22 +1726,45 @@ def _translate_with_diverse_beam_search( num_beams = self._config.params.get("generation_num_beams") diversity_penalty: Optional[float] = self._config.infer.get("diversity_penalty") - translations = pipeline( - sentences, - num_beams=num_beams, - num_beam_groups=num_beams, - num_return_sequences=num_return_sequences, - diversity_penalty=diversity_penalty, - force_words_ids=force_words_ids, - batch_size=batch_size, - return_text=not return_tensors, - return_tensors=return_tensors, - ) - - if num_return_sequences == 1: - translations = [[t] for t in translations] + sentences = list(sentences) + current_batch_size = batch_size + translations: List[List[dict]] = [] + index = 0 + while index < len(sentences): + current_batch_size = min(current_batch_size, len(sentences) - index) + batch_sentences = sentences[index : index + current_batch_size] + batch_force_words_ids = ( + force_words_ids[index : index + current_batch_size] if force_words_ids is not None else None + ) + try: + batch_translations = pipeline( + batch_sentences, + num_beams=num_beams, + num_beam_groups=num_beams, + num_return_sequences=num_return_sequences, + diversity_penalty=diversity_penalty, + force_words_ids=batch_force_words_ids, + batch_size=current_batch_size, + return_text=not return_tensors, + return_tensors=return_tensors, + ) + if num_return_sequences == 1: + batch_translations = [[t] for t in batch_translations] + translations.extend(self._flatten_tokenized_translations(batch_translations)) + index += current_batch_size + except RuntimeError as e: + if not _should_reduce_batch_size(e) or current_batch_size <= 1: + raise + current_batch_size //= 2 + LOGGER.warning( + "OOM during diverse-beam translation inference; reducing batch size to %d and retrying current batch.", + current_batch_size, + ) + gc.collect() + if torch.cuda.is_available(): + torch.cuda.empty_cache() - return self._flatten_tokenized_translations(translations) + return translations def _create_inference_model( self, diff --git a/tests/smoke_tests/test_experiment.py b/tests/smoke_tests/test_experiment.py index 18fb906c..539ed0f6 100644 --- a/tests/smoke_tests/test_experiment.py +++ b/tests/smoke_tests/test_experiment.py @@ -62,22 +62,24 @@ def fake_translate_sentence_helper( assert captured_batch_size["value"] == 2 -def test_translate_with_adaptive_batch_size_retries_with_smaller_batch(): +def test_translate_with_beam_search_retries_with_smaller_batch(): model = cast(HuggingFaceNMTModel, Mock(spec=HuggingFaceNMTModel)) + model._config = cast(HuggingFaceConfig, Mock(infer={"num_beams": 2}, params={})) call_batch_sizes: list[int] = [] - def fake_translate(batch_sentences, batch_size, batch_force_words_ids): + def fake_pipeline(batch_sentences, batch_size, **kwargs): call_batch_sizes.append(batch_size) if batch_size > 2: raise RuntimeError(OOM_ERROR_MESSAGE) return [[{"translation_text": str(sentence[0])}] for sentence in batch_sentences] sentences = [["a"], ["b"], ["c"], ["d"]] - translated = HuggingFaceNMTModel._translate_with_adaptive_batch_size( + translated = HuggingFaceNMTModel._translate_with_beam_search( model, - fake_translate, + fake_pipeline, sentences, batch_size=4, + return_tensors=False, ) assert [translation[0]["translation_text"] for translation in translated] == ["a", "b", "c", "d"] From bda48334106d9e6c3c5134f5e333a8b3a634588f Mon Sep 17 00:00:00 2001 From: "copilot-swe-agent[bot]" <198982749+Copilot@users.noreply.github.com> Date: Wed, 20 May 2026 18:16:31 +0000 Subject: [PATCH 11/16] test: align beam-search retry mock output with pipeline format Agent-Logs-Url: https://github.com/sillsdev/silnlp/sessions/ee62a80c-fe7c-4ffd-b732-fda53446b7c4 Co-authored-by: benjaminking <1214233+benjaminking@users.noreply.github.com> --- silnlp/nmt/hugging_face_config.py | 3 --- tests/smoke_tests/test_experiment.py | 2 +- 2 files changed, 1 insertion(+), 4 deletions(-) diff --git a/silnlp/nmt/hugging_face_config.py b/silnlp/nmt/hugging_face_config.py index 052d8abd..bc924562 100644 --- a/silnlp/nmt/hugging_face_config.py +++ b/silnlp/nmt/hugging_face_config.py @@ -1623,7 +1623,6 @@ def _translate_with_beam_search( if num_beams is None: num_beams = self._config.params.get("generation_num_beams") - sentences = list(sentences) current_batch_size = batch_size translations: List[List[dict]] = [] index = 0 @@ -1673,7 +1672,6 @@ def _translate_with_sampling( temperature: Optional[int] = self._config.infer.get("temperature") - sentences = list(sentences) current_batch_size = batch_size translations: List[List[dict]] = [] index = 0 @@ -1726,7 +1724,6 @@ def _translate_with_diverse_beam_search( num_beams = self._config.params.get("generation_num_beams") diversity_penalty: Optional[float] = self._config.infer.get("diversity_penalty") - sentences = list(sentences) current_batch_size = batch_size translations: List[List[dict]] = [] index = 0 diff --git a/tests/smoke_tests/test_experiment.py b/tests/smoke_tests/test_experiment.py index 539ed0f6..229756b4 100644 --- a/tests/smoke_tests/test_experiment.py +++ b/tests/smoke_tests/test_experiment.py @@ -71,7 +71,7 @@ def fake_pipeline(batch_sentences, batch_size, **kwargs): call_batch_sizes.append(batch_size) if batch_size > 2: raise RuntimeError(OOM_ERROR_MESSAGE) - return [[{"translation_text": str(sentence[0])}] for sentence in batch_sentences] + return [{"translation_text": str(sentence[0])} for sentence in batch_sentences] sentences = [["a"], ["b"], ["c"], ["d"]] translated = HuggingFaceNMTModel._translate_with_beam_search( From 2864bd652205b9c9c61cb462472cccc4bc54bae2 Mon Sep 17 00:00:00 2001 From: "copilot-swe-agent[bot]" <198982749+Copilot@users.noreply.github.com> Date: Wed, 20 May 2026 18:28:51 +0000 Subject: [PATCH 12/16] refactor: extract shared translate batching helpers Agent-Logs-Url: https://github.com/sillsdev/silnlp/sessions/e0d13ac1-561a-44de-b3a9-3a9f1f3dab80 Co-authored-by: benjaminking <1214233+benjaminking@users.noreply.github.com> --- silnlp/nmt/hugging_face_config.py | 106 ++++++++++++++++-------------- 1 file changed, 55 insertions(+), 51 deletions(-) diff --git a/silnlp/nmt/hugging_face_config.py b/silnlp/nmt/hugging_face_config.py index bc924562..63af6953 100644 --- a/silnlp/nmt/hugging_face_config.py +++ b/silnlp/nmt/hugging_face_config.py @@ -1610,6 +1610,37 @@ def _translate_sentence_helper( def _flatten_tokenized_translations(self, pipeline_output) -> List[dict]: return [[i if isinstance(i, dict) else i[0] for i in translation] for translation in pipeline_output] + def _get_translation_batch( + self, + sentences: List[TSent], + force_words_ids: Optional[List[List[List[int]]]], + index: int, + current_batch_size: int, + ) -> Tuple[int, List[TSent], Optional[List[List[List[int]]]]]: + current_batch_size = min(current_batch_size, len(sentences) - index) + batch_sentences = sentences[index : index + current_batch_size] + batch_force_words_ids = force_words_ids[index : index + current_batch_size] if force_words_ids is not None else None + return current_batch_size, batch_sentences, batch_force_words_ids + + def _normalize_batch_translations(self, batch_translations, num_return_sequences: int) -> List[List[dict]]: + if num_return_sequences == 1: + batch_translations = [[t] for t in batch_translations] + return self._flatten_tokenized_translations(batch_translations) + + def _handle_translation_oom(self, error: RuntimeError, current_batch_size: int, method_name: str) -> int: + if not _should_reduce_batch_size(error) or current_batch_size <= 1: + raise error + next_batch_size = current_batch_size // 2 + LOGGER.warning( + "OOM during %s translation inference; reducing batch size to %d and retrying current batch.", + method_name, + next_batch_size, + ) + gc.collect() + if torch.cuda.is_available(): + torch.cuda.empty_cache() + return next_batch_size + def _translate_with_beam_search( self, pipeline: TranslationPipeline, @@ -1623,14 +1654,16 @@ def _translate_with_beam_search( if num_beams is None: num_beams = self._config.params.get("generation_num_beams") + sentences = list(sentences) current_batch_size = batch_size translations: List[List[dict]] = [] index = 0 while index < len(sentences): - current_batch_size = min(current_batch_size, len(sentences) - index) - batch_sentences = sentences[index : index + current_batch_size] - batch_force_words_ids = ( - force_words_ids[index : index + current_batch_size] if force_words_ids is not None else None + current_batch_size, batch_sentences, batch_force_words_ids = self._get_translation_batch( + sentences, + force_words_ids, + index, + current_batch_size, ) try: batch_translations = pipeline( @@ -1642,21 +1675,10 @@ def _translate_with_beam_search( return_text=not return_tensors, return_tensors=return_tensors, ) - if num_return_sequences == 1: - batch_translations = [[t] for t in batch_translations] - translations.extend(self._flatten_tokenized_translations(batch_translations)) + translations.extend(self._normalize_batch_translations(batch_translations, num_return_sequences)) index += current_batch_size except RuntimeError as e: - if not _should_reduce_batch_size(e) or current_batch_size <= 1: - raise - current_batch_size //= 2 - LOGGER.warning( - "OOM during beam-search translation inference; reducing batch size to %d and retrying current batch.", - current_batch_size, - ) - gc.collect() - if torch.cuda.is_available(): - torch.cuda.empty_cache() + current_batch_size = self._handle_translation_oom(e, current_batch_size, "beam-search") return translations @@ -1672,14 +1694,16 @@ def _translate_with_sampling( temperature: Optional[int] = self._config.infer.get("temperature") + sentences = list(sentences) current_batch_size = batch_size translations: List[List[dict]] = [] index = 0 while index < len(sentences): - current_batch_size = min(current_batch_size, len(sentences) - index) - batch_sentences = sentences[index : index + current_batch_size] - batch_force_words_ids = ( - force_words_ids[index : index + current_batch_size] if force_words_ids is not None else None + current_batch_size, batch_sentences, batch_force_words_ids = self._get_translation_batch( + sentences, + force_words_ids, + index, + current_batch_size, ) try: batch_translations = pipeline( @@ -1692,21 +1716,10 @@ def _translate_with_sampling( return_text=not return_tensors, return_tensors=return_tensors, ) - if num_return_sequences == 1: - batch_translations = [[t] for t in batch_translations] - translations.extend(self._flatten_tokenized_translations(batch_translations)) + translations.extend(self._normalize_batch_translations(batch_translations, num_return_sequences)) index += current_batch_size except RuntimeError as e: - if not _should_reduce_batch_size(e) or current_batch_size <= 1: - raise - current_batch_size //= 2 - LOGGER.warning( - "OOM during sampling translation inference; reducing batch size to %d and retrying current batch.", - current_batch_size, - ) - gc.collect() - if torch.cuda.is_available(): - torch.cuda.empty_cache() + current_batch_size = self._handle_translation_oom(e, current_batch_size, "sampling") return translations @@ -1724,14 +1737,16 @@ def _translate_with_diverse_beam_search( num_beams = self._config.params.get("generation_num_beams") diversity_penalty: Optional[float] = self._config.infer.get("diversity_penalty") + sentences = list(sentences) current_batch_size = batch_size translations: List[List[dict]] = [] index = 0 while index < len(sentences): - current_batch_size = min(current_batch_size, len(sentences) - index) - batch_sentences = sentences[index : index + current_batch_size] - batch_force_words_ids = ( - force_words_ids[index : index + current_batch_size] if force_words_ids is not None else None + current_batch_size, batch_sentences, batch_force_words_ids = self._get_translation_batch( + sentences, + force_words_ids, + index, + current_batch_size, ) try: batch_translations = pipeline( @@ -1745,21 +1760,10 @@ def _translate_with_diverse_beam_search( return_text=not return_tensors, return_tensors=return_tensors, ) - if num_return_sequences == 1: - batch_translations = [[t] for t in batch_translations] - translations.extend(self._flatten_tokenized_translations(batch_translations)) + translations.extend(self._normalize_batch_translations(batch_translations, num_return_sequences)) index += current_batch_size except RuntimeError as e: - if not _should_reduce_batch_size(e) or current_batch_size <= 1: - raise - current_batch_size //= 2 - LOGGER.warning( - "OOM during diverse-beam translation inference; reducing batch size to %d and retrying current batch.", - current_batch_size, - ) - gc.collect() - if torch.cuda.is_available(): - torch.cuda.empty_cache() + current_batch_size = self._handle_translation_oom(e, current_batch_size, "diverse-beam") return translations From 7ca4455285094e8da4b0a63e2e4d34ebaab74965 Mon Sep 17 00:00:00 2001 From: "copilot-swe-agent[bot]" <198982749+Copilot@users.noreply.github.com> Date: Wed, 20 May 2026 18:30:27 +0000 Subject: [PATCH 13/16] perf: remove redundant sentence list conversion in helper Agent-Logs-Url: https://github.com/sillsdev/silnlp/sessions/e0d13ac1-561a-44de-b3a9-3a9f1f3dab80 Co-authored-by: benjaminking <1214233+benjaminking@users.noreply.github.com> --- silnlp/nmt/hugging_face_config.py | 1 - 1 file changed, 1 deletion(-) diff --git a/silnlp/nmt/hugging_face_config.py b/silnlp/nmt/hugging_face_config.py index 63af6953..40fb3fc3 100644 --- a/silnlp/nmt/hugging_face_config.py +++ b/silnlp/nmt/hugging_face_config.py @@ -1520,7 +1520,6 @@ def _translate_sentence_helper( force_words_ids: List[List[List[int]]] = None, produce_multiple_translations: bool = False, ) -> Iterable[ModelOutputGroup]: - sentences = list(sentences) num_drafts = self.get_num_drafts() if produce_multiple_translations and num_drafts > 1: multiple_translations_method: str = self._config.infer.get("multiple_translations_method") From de611902ffb548177d51e4deb944f7f7484b5c14 Mon Sep 17 00:00:00 2001 From: "copilot-swe-agent[bot]" <198982749+Copilot@users.noreply.github.com> Date: Wed, 20 May 2026 19:38:22 +0000 Subject: [PATCH 14/16] refactor: move adaptive batch sizing loop up into _translate_sentences Agent-Logs-Url: https://github.com/sillsdev/silnlp/sessions/b110343c-0608-42c4-82ec-6d4f0ef3857d Co-authored-by: benjaminking <1214233+benjaminking@users.noreply.github.com> --- silnlp/nmt/hugging_face_config.py | 199 ++++++++++-------------------- 1 file changed, 66 insertions(+), 133 deletions(-) diff --git a/silnlp/nmt/hugging_face_config.py b/silnlp/nmt/hugging_face_config.py index 40fb3fc3..5cfeea2a 100644 --- a/silnlp/nmt/hugging_face_config.py +++ b/silnlp/nmt/hugging_face_config.py @@ -1495,6 +1495,7 @@ def _translate_sentences( batch_size = min(batch_size, batch_size_with_tensors) dictionary = self._get_dictionary() + current_batch_size = batch_size for batch, force_words in batch_sentences(sentences, vrefs, batch_size, dictionary): if force_words is None: force_words_ids = None @@ -1502,20 +1503,39 @@ def _translate_sentences( force_words_ids = [[tokenizer.convert_tokens_to_ids(v) for v in vs] for vs in force_words] force_words_ids = prune_sublists(force_words_ids) - yield from self._translate_sentence_helper( - pipeline, - batch, - batch_size, - return_tensors, - force_words_ids, - produce_multiple_translations=produce_multiple_translations, - ) + batch_list = list(batch) + index = 0 + while index < len(batch_list): + effective_size = min(current_batch_size, len(batch_list) - index) + sub_batch = batch_list[index : index + effective_size] + sub_force_words_ids = ( + force_words_ids[index : index + effective_size] if force_words_ids is not None else None + ) + try: + yield from self._translate_sentence_helper( + pipeline, + sub_batch, + return_tensors, + sub_force_words_ids, + produce_multiple_translations=produce_multiple_translations, + ) + index += effective_size + except RuntimeError as e: + if not _should_reduce_batch_size(e) or current_batch_size <= 1: + raise + current_batch_size //= 2 + LOGGER.warning( + "OOM during translation inference; reducing batch size to %d and retrying current batch.", + current_batch_size, + ) + gc.collect() + if torch.cuda.is_available(): + torch.cuda.empty_cache() def _translate_sentence_helper( self, pipeline: TranslationPipeline, - sentences: Iterable[TSent], - batch_size: int, + sentences: List[TSent], return_tensors: bool, force_words_ids: List[List[List[int]]] = None, produce_multiple_translations: bool = False, @@ -1528,7 +1548,6 @@ def _translate_sentence_helper( beam_search_results: List[dict] = self._translate_with_beam_search( pipeline, sentences, - batch_size, return_tensors, num_return_sequences=1, force_words_ids=force_words_ids, @@ -1537,7 +1556,6 @@ def _translate_sentence_helper( sampling_results: List[dict] = self._translate_with_sampling( pipeline, sentences, - batch_size, return_tensors, num_return_sequences=num_drafts - 1, force_words_ids=force_words_ids, @@ -1555,7 +1573,6 @@ def _translate_sentence_helper( for result in self._translate_with_sampling( pipeline, sentences, - batch_size, return_tensors, num_return_sequences=num_drafts, force_words_ids=force_words_ids, @@ -1568,7 +1585,6 @@ def _translate_sentence_helper( for result in self._translate_with_beam_search( pipeline, sentences, - batch_size, return_tensors, num_return_sequences=num_drafts, force_words_ids=force_words_ids, @@ -1581,7 +1597,6 @@ def _translate_sentence_helper( for result in self._translate_with_diverse_beam_search( pipeline, sentences, - batch_size, return_tensors, num_return_sequences=num_drafts, force_words_ids=force_words_ids, @@ -1596,7 +1611,6 @@ def _translate_sentence_helper( for translated_sentence in self._translate_with_beam_search( pipeline, sentences, - batch_size, return_tensors, num_return_sequences=1, force_words_ids=force_words_ids, @@ -1609,42 +1623,15 @@ def _translate_sentence_helper( def _flatten_tokenized_translations(self, pipeline_output) -> List[dict]: return [[i if isinstance(i, dict) else i[0] for i in translation] for translation in pipeline_output] - def _get_translation_batch( - self, - sentences: List[TSent], - force_words_ids: Optional[List[List[List[int]]]], - index: int, - current_batch_size: int, - ) -> Tuple[int, List[TSent], Optional[List[List[List[int]]]]]: - current_batch_size = min(current_batch_size, len(sentences) - index) - batch_sentences = sentences[index : index + current_batch_size] - batch_force_words_ids = force_words_ids[index : index + current_batch_size] if force_words_ids is not None else None - return current_batch_size, batch_sentences, batch_force_words_ids - def _normalize_batch_translations(self, batch_translations, num_return_sequences: int) -> List[List[dict]]: if num_return_sequences == 1: batch_translations = [[t] for t in batch_translations] return self._flatten_tokenized_translations(batch_translations) - def _handle_translation_oom(self, error: RuntimeError, current_batch_size: int, method_name: str) -> int: - if not _should_reduce_batch_size(error) or current_batch_size <= 1: - raise error - next_batch_size = current_batch_size // 2 - LOGGER.warning( - "OOM during %s translation inference; reducing batch size to %d and retrying current batch.", - method_name, - next_batch_size, - ) - gc.collect() - if torch.cuda.is_available(): - torch.cuda.empty_cache() - return next_batch_size - def _translate_with_beam_search( self, pipeline: TranslationPipeline, - sentences: Iterable[TSent], - batch_size: int, + sentences: List[TSent], return_tensors: bool, num_return_sequences: int = 1, force_words_ids: List[List[List[int]]] = None, @@ -1653,80 +1640,43 @@ def _translate_with_beam_search( if num_beams is None: num_beams = self._config.params.get("generation_num_beams") - sentences = list(sentences) - current_batch_size = batch_size - translations: List[List[dict]] = [] - index = 0 - while index < len(sentences): - current_batch_size, batch_sentences, batch_force_words_ids = self._get_translation_batch( - sentences, - force_words_ids, - index, - current_batch_size, - ) - try: - batch_translations = pipeline( - batch_sentences, - num_beams=num_beams, - num_return_sequences=num_return_sequences, - force_words_ids=batch_force_words_ids, - batch_size=current_batch_size, - return_text=not return_tensors, - return_tensors=return_tensors, - ) - translations.extend(self._normalize_batch_translations(batch_translations, num_return_sequences)) - index += current_batch_size - except RuntimeError as e: - current_batch_size = self._handle_translation_oom(e, current_batch_size, "beam-search") - - return translations + batch_translations = pipeline( + sentences, + num_beams=num_beams, + num_return_sequences=num_return_sequences, + force_words_ids=force_words_ids, + batch_size=len(sentences), + return_text=not return_tensors, + return_tensors=return_tensors, + ) + return self._normalize_batch_translations(batch_translations, num_return_sequences) def _translate_with_sampling( self, pipeline: TranslationPipeline, - sentences: Iterable[TSent], - batch_size: int, + sentences: List[TSent], return_tensors: bool, num_return_sequences: int = 1, force_words_ids: List[List[List[int]]] = None, ) -> List[List[dict]]: - temperature: Optional[int] = self._config.infer.get("temperature") - sentences = list(sentences) - current_batch_size = batch_size - translations: List[List[dict]] = [] - index = 0 - while index < len(sentences): - current_batch_size, batch_sentences, batch_force_words_ids = self._get_translation_batch( - sentences, - force_words_ids, - index, - current_batch_size, - ) - try: - batch_translations = pipeline( - batch_sentences, - do_sample=True, - temperature=temperature, - num_return_sequences=num_return_sequences, - force_words_ids=batch_force_words_ids, - batch_size=current_batch_size, - return_text=not return_tensors, - return_tensors=return_tensors, - ) - translations.extend(self._normalize_batch_translations(batch_translations, num_return_sequences)) - index += current_batch_size - except RuntimeError as e: - current_batch_size = self._handle_translation_oom(e, current_batch_size, "sampling") - - return translations + batch_translations = pipeline( + sentences, + do_sample=True, + temperature=temperature, + num_return_sequences=num_return_sequences, + force_words_ids=force_words_ids, + batch_size=len(sentences), + return_text=not return_tensors, + return_tensors=return_tensors, + ) + return self._normalize_batch_translations(batch_translations, num_return_sequences) def _translate_with_diverse_beam_search( self, pipeline: TranslationPipeline, - sentences: Iterable[TSent], - batch_size: int, + sentences: List[TSent], return_tensors: bool, num_return_sequences: int = 1, force_words_ids: List[List[List[int]]] = None, @@ -1736,35 +1686,18 @@ def _translate_with_diverse_beam_search( num_beams = self._config.params.get("generation_num_beams") diversity_penalty: Optional[float] = self._config.infer.get("diversity_penalty") - sentences = list(sentences) - current_batch_size = batch_size - translations: List[List[dict]] = [] - index = 0 - while index < len(sentences): - current_batch_size, batch_sentences, batch_force_words_ids = self._get_translation_batch( - sentences, - force_words_ids, - index, - current_batch_size, - ) - try: - batch_translations = pipeline( - batch_sentences, - num_beams=num_beams, - num_beam_groups=num_beams, - num_return_sequences=num_return_sequences, - diversity_penalty=diversity_penalty, - force_words_ids=batch_force_words_ids, - batch_size=current_batch_size, - return_text=not return_tensors, - return_tensors=return_tensors, - ) - translations.extend(self._normalize_batch_translations(batch_translations, num_return_sequences)) - index += current_batch_size - except RuntimeError as e: - current_batch_size = self._handle_translation_oom(e, current_batch_size, "diverse-beam") - - return translations + batch_translations = pipeline( + sentences, + num_beams=num_beams, + num_beam_groups=num_beams, + num_return_sequences=num_return_sequences, + diversity_penalty=diversity_penalty, + force_words_ids=force_words_ids, + batch_size=len(sentences), + return_text=not return_tensors, + return_tensors=return_tensors, + ) + return self._normalize_batch_translations(batch_translations, num_return_sequences) def _create_inference_model( self, From 8ca9d2451f3c5161d7c216a3edecc6c030abefbc Mon Sep 17 00:00:00 2001 From: "copilot-swe-agent[bot]" <198982749+Copilot@users.noreply.github.com> Date: Wed, 20 May 2026 19:40:11 +0000 Subject: [PATCH 15/16] test: update tests for new _translate_sentences adaptive batch sizing Agent-Logs-Url: https://github.com/sillsdev/silnlp/sessions/b110343c-0608-42c4-82ec-6d4f0ef3857d Co-authored-by: benjaminking <1214233+benjaminking@users.noreply.github.com> --- tests/smoke_tests/test_experiment.py | 51 ++++++++++++++++------------ 1 file changed, 30 insertions(+), 21 deletions(-) diff --git a/tests/smoke_tests/test_experiment.py b/tests/smoke_tests/test_experiment.py index 229756b4..ab9987e1 100644 --- a/tests/smoke_tests/test_experiment.py +++ b/tests/smoke_tests/test_experiment.py @@ -38,12 +38,12 @@ def test_translate_sentences_uses_tensor_batch_size_cap(): model = cast(HuggingFaceNMTModel, Mock(spec=HuggingFaceNMTModel)) model._config = cast(HuggingFaceConfig, Mock(infer={"infer_batch_size": 16, "infer_batch_size_with_tensors": 2})) model._get_dictionary = Mock(return_value={}) - captured_batch_size = {} + captured_sub_batch_size = {} def fake_translate_sentence_helper( - pipeline, sentences, batch_size, return_tensors, force_words_ids=None, produce_multiple_translations=False + pipeline, sentences, return_tensors, force_words_ids=None, produce_multiple_translations=False ): - captured_batch_size["value"] = batch_size + captured_sub_batch_size["value"] = len(sentences) return iter(()) model._translate_sentence_helper = fake_translate_sentence_helper @@ -59,31 +59,40 @@ def fake_translate_sentence_helper( ) ) - assert captured_batch_size["value"] == 2 + assert captured_sub_batch_size["value"] == 1 -def test_translate_with_beam_search_retries_with_smaller_batch(): +def test_translate_sentences_retries_with_smaller_batch(): model = cast(HuggingFaceNMTModel, Mock(spec=HuggingFaceNMTModel)) - model._config = cast(HuggingFaceConfig, Mock(infer={"num_beams": 2}, params={})) - call_batch_sizes: list[int] = [] + model._config = cast( + HuggingFaceConfig, Mock(infer={"infer_batch_size": 4, "num_beams": 2}, params={}) + ) + model._get_dictionary = Mock(return_value={}) + call_sub_batch_sizes: list[int] = [] - def fake_pipeline(batch_sentences, batch_size, **kwargs): - call_batch_sizes.append(batch_size) - if batch_size > 2: + def fake_translate_sentence_helper( + pipeline, sentences, return_tensors, force_words_ids=None, produce_multiple_translations=False + ): + call_sub_batch_sizes.append(len(sentences)) + if len(sentences) > 2: raise RuntimeError(OOM_ERROR_MESSAGE) - return [{"translation_text": str(sentence[0])} for sentence in batch_sentences] - - sentences = [["a"], ["b"], ["c"], ["d"]] - translated = HuggingFaceNMTModel._translate_with_beam_search( - model, - fake_pipeline, - sentences, - batch_size=4, - return_tensors=False, + return iter(()) + + model._translate_sentence_helper = fake_translate_sentence_helper + + list( + HuggingFaceNMTModel._translate_sentences( + model, + tokenizer=Mock(), + pipeline=Mock(), + sentences=[["a"], ["b"], ["c"], ["d"]], + vrefs=None, + return_tensors=False, + ) ) - assert [translation[0]["translation_text"] for translation in translated] == ["a", "b", "c", "d"] - assert call_batch_sizes == [4, 2, 2] + # First call is batch-of-4 (OOM), then two calls of batch-of-2 succeed + assert call_sub_batch_sizes == [4, 2, 2] def set_up_environment() -> SilNlpEnv: From e7cc3eedb1450a0793e0132a94db4e8b385c76ac Mon Sep 17 00:00:00 2001 From: "copilot-swe-agent[bot]" <198982749+Copilot@users.noreply.github.com> Date: Wed, 20 May 2026 20:15:27 +0000 Subject: [PATCH 16/16] refactor: rename translate batch helper and drop tensor cap/pipeline batch_size Agent-Logs-Url: https://github.com/sillsdev/silnlp/sessions/06f3a41e-075a-43a6-87b1-2887369bed21 Co-authored-by: benjaminking <1214233+benjaminking@users.noreply.github.com> --- silnlp/nmt/hugging_face_config.py | 37 +++++++++++----------------- tests/smoke_tests/test_experiment.py | 26 +++++++++---------- 2 files changed, 28 insertions(+), 35 deletions(-) diff --git a/silnlp/nmt/hugging_face_config.py b/silnlp/nmt/hugging_face_config.py index 5cfeea2a..324fafc5 100644 --- a/silnlp/nmt/hugging_face_config.py +++ b/silnlp/nmt/hugging_face_config.py @@ -396,7 +396,6 @@ def __init__(self, exp_dir: Path, config: dict) -> None: }, "infer": { "infer_batch_size": 16, - "infer_batch_size_with_tensors": 16, "num_beams": 2, "num_drafts": 3, "multiple_translations_method": "hybrid", @@ -1490,9 +1489,6 @@ def _translate_sentences( return_tensors: bool = False, ) -> Iterable[ModelOutputGroup]: batch_size: int = self._config.infer["infer_batch_size"] - if return_tensors: - batch_size_with_tensors = self._config.infer.get("infer_batch_size_with_tensors", batch_size) - batch_size = min(batch_size, batch_size_with_tensors) dictionary = self._get_dictionary() current_batch_size = batch_size @@ -1512,7 +1508,7 @@ def _translate_sentences( force_words_ids[index : index + effective_size] if force_words_ids is not None else None ) try: - yield from self._translate_sentence_helper( + yield from self._translate_batch( pipeline, sub_batch, return_tensors, @@ -1532,10 +1528,10 @@ def _translate_sentences( if torch.cuda.is_available(): torch.cuda.empty_cache() - def _translate_sentence_helper( + def _translate_batch( self, pipeline: TranslationPipeline, - sentences: List[TSent], + batch: List[TSent], return_tensors: bool, force_words_ids: List[List[List[int]]] = None, produce_multiple_translations: bool = False, @@ -1547,7 +1543,7 @@ def _translate_sentence_helper( if multiple_translations_method == "hybrid": beam_search_results: List[dict] = self._translate_with_beam_search( pipeline, - sentences, + batch, return_tensors, num_return_sequences=1, force_words_ids=force_words_ids, @@ -1555,7 +1551,7 @@ def _translate_sentence_helper( sampling_results: List[dict] = self._translate_with_sampling( pipeline, - sentences, + batch, return_tensors, num_return_sequences=num_drafts - 1, force_words_ids=force_words_ids, @@ -1572,7 +1568,7 @@ def _translate_sentence_helper( ModelOutputGroup(result) for result in self._translate_with_sampling( pipeline, - sentences, + batch, return_tensors, num_return_sequences=num_drafts, force_words_ids=force_words_ids, @@ -1584,7 +1580,7 @@ def _translate_sentence_helper( ModelOutputGroup(result) for result in self._translate_with_beam_search( pipeline, - sentences, + batch, return_tensors, num_return_sequences=num_drafts, force_words_ids=force_words_ids, @@ -1596,7 +1592,7 @@ def _translate_sentence_helper( ModelOutputGroup(result) for result in self._translate_with_diverse_beam_search( pipeline, - sentences, + batch, return_tensors, num_return_sequences=num_drafts, force_words_ids=force_words_ids, @@ -1610,7 +1606,7 @@ def _translate_sentence_helper( ModelOutputGroup([translated_sentence[0]]) for translated_sentence in self._translate_with_beam_search( pipeline, - sentences, + batch, return_tensors, num_return_sequences=1, force_words_ids=force_words_ids, @@ -1631,7 +1627,7 @@ def _normalize_batch_translations(self, batch_translations, num_return_sequences def _translate_with_beam_search( self, pipeline: TranslationPipeline, - sentences: List[TSent], + batch: List[TSent], return_tensors: bool, num_return_sequences: int = 1, force_words_ids: List[List[List[int]]] = None, @@ -1641,11 +1637,10 @@ def _translate_with_beam_search( num_beams = self._config.params.get("generation_num_beams") batch_translations = pipeline( - sentences, + batch, num_beams=num_beams, num_return_sequences=num_return_sequences, force_words_ids=force_words_ids, - batch_size=len(sentences), return_text=not return_tensors, return_tensors=return_tensors, ) @@ -1654,7 +1649,7 @@ def _translate_with_beam_search( def _translate_with_sampling( self, pipeline: TranslationPipeline, - sentences: List[TSent], + batch: List[TSent], return_tensors: bool, num_return_sequences: int = 1, force_words_ids: List[List[List[int]]] = None, @@ -1662,12 +1657,11 @@ def _translate_with_sampling( temperature: Optional[int] = self._config.infer.get("temperature") batch_translations = pipeline( - sentences, + batch, do_sample=True, temperature=temperature, num_return_sequences=num_return_sequences, force_words_ids=force_words_ids, - batch_size=len(sentences), return_text=not return_tensors, return_tensors=return_tensors, ) @@ -1676,7 +1670,7 @@ def _translate_with_sampling( def _translate_with_diverse_beam_search( self, pipeline: TranslationPipeline, - sentences: List[TSent], + batch: List[TSent], return_tensors: bool, num_return_sequences: int = 1, force_words_ids: List[List[List[int]]] = None, @@ -1687,13 +1681,12 @@ def _translate_with_diverse_beam_search( diversity_penalty: Optional[float] = self._config.infer.get("diversity_penalty") batch_translations = pipeline( - sentences, + batch, num_beams=num_beams, num_beam_groups=num_beams, num_return_sequences=num_return_sequences, diversity_penalty=diversity_penalty, force_words_ids=force_words_ids, - batch_size=len(sentences), return_text=not return_tensors, return_tensors=return_tensors, ) diff --git a/tests/smoke_tests/test_experiment.py b/tests/smoke_tests/test_experiment.py index ab9987e1..9dd3638a 100644 --- a/tests/smoke_tests/test_experiment.py +++ b/tests/smoke_tests/test_experiment.py @@ -34,32 +34,32 @@ def test_experiment_full_pipeline(): clean_experiment_directory(environment.get_mt_exp_dir(EXPERIMENT_NAME)) -def test_translate_sentences_uses_tensor_batch_size_cap(): +def test_translate_sentences_does_not_cap_batch_size_with_tensors(): model = cast(HuggingFaceNMTModel, Mock(spec=HuggingFaceNMTModel)) - model._config = cast(HuggingFaceConfig, Mock(infer={"infer_batch_size": 16, "infer_batch_size_with_tensors": 2})) + model._config = cast(HuggingFaceConfig, Mock(infer={"infer_batch_size": 4})) model._get_dictionary = Mock(return_value={}) captured_sub_batch_size = {} - def fake_translate_sentence_helper( - pipeline, sentences, return_tensors, force_words_ids=None, produce_multiple_translations=False + def fake_translate_batch( + pipeline, batch, return_tensors, force_words_ids=None, produce_multiple_translations=False ): - captured_sub_batch_size["value"] = len(sentences) + captured_sub_batch_size["value"] = len(batch) return iter(()) - model._translate_sentence_helper = fake_translate_sentence_helper + model._translate_batch = fake_translate_batch list( HuggingFaceNMTModel._translate_sentences( model, tokenizer=Mock(), pipeline=Mock(), - sentences=[["token"]], + sentences=[["a"], ["b"], ["c"], ["d"]], vrefs=None, return_tensors=True, ) ) - assert captured_sub_batch_size["value"] == 1 + assert captured_sub_batch_size["value"] == 4 def test_translate_sentences_retries_with_smaller_batch(): @@ -70,15 +70,15 @@ def test_translate_sentences_retries_with_smaller_batch(): model._get_dictionary = Mock(return_value={}) call_sub_batch_sizes: list[int] = [] - def fake_translate_sentence_helper( - pipeline, sentences, return_tensors, force_words_ids=None, produce_multiple_translations=False + def fake_translate_batch( + pipeline, batch, return_tensors, force_words_ids=None, produce_multiple_translations=False ): - call_sub_batch_sizes.append(len(sentences)) - if len(sentences) > 2: + call_sub_batch_sizes.append(len(batch)) + if len(batch) > 2: raise RuntimeError(OOM_ERROR_MESSAGE) return iter(()) - model._translate_sentence_helper = fake_translate_sentence_helper + model._translate_batch = fake_translate_batch list( HuggingFaceNMTModel._translate_sentences(