From d2969c6656fe7749d2d1d3354d8dfd16407f5e7e Mon Sep 17 00:00:00 2001 From: Lion Summerbell Date: Thu, 15 Oct 2026 14:43:25 +0200 Subject: [PATCH] refactor(translation-worker): sentence batching by length for hunyuan translation --- .../tests/test_activities.py | 74 ++++++++++++++++ .../tests/test_hunyuan_translator.py | 4 +- .../translation_worker/activities.py | 84 +++++++++++-------- .../translation_worker/translators/hunyuan.py | 4 +- 4 files changed, 126 insertions(+), 40 deletions(-) diff --git a/workers/translation-worker/tests/test_activities.py b/workers/translation-worker/tests/test_activities.py index 4977beb9..1194c75b 100644 --- a/workers/translation-worker/tests/test_activities.py +++ b/workers/translation-worker/tests/test_activities.py @@ -23,6 +23,7 @@ _DOC_CONTENT_TEXT_LENGTH, _get_es_docs_by_language, _split_sentences, + _translate_by_length, _update_docs_translation, create_translation_batches_act, translate_docs_act, @@ -485,6 +486,79 @@ async def test_translate_docs_act__es_update( assert captured == expected +class _UpperTranslator(MockTranslator): + def __init__(self): + super().__init__([]) + self.batches = [] + + def translate(self, texts: Iterable[str]) -> list[str]: + texts = list(texts) + self.batches.append(texts) + return [t.upper() for t in texts] + + +async def test__translate_by_length__batches_by_length_and_restores_order() -> None: + # Given + doc_1, doc_2 = Document.from_es(FR_DOC_1), Document.from_es(FR_DOC_2) + doc_sents = [ + (doc_1, "a long sentence"), + (doc_1, "short"), + (doc_2, "the longest sentence"), + (doc_2, "tiny"), + (doc_2, "mid sized"), + ] + translator = _UpperTranslator() + # When + translated = await _collect_async( + _translate_by_length(doc_sents, translator, batch_size=2) + ) + # Then + assert translator.batches == [ + ["tiny", "short"], + ["mid sized", "a long sentence"], + ["the longest sentence"], + ] + # Documents are yielded as soon as all their sentences are translated + assert translated == [ + (doc_1, ["A LONG SENTENCE", "SHORT"]), + (doc_2, ["THE LONGEST SENTENCE", "TINY", "MID SIZED"]), + ] + + +async def test_translate_docs_act__es_update_keeps_sentence_order( + monkeypatch: MonkeyPatch, test_worker_config: TranslationWorkerConfig +) -> None: + # Given + worker_config = test_worker_config + batches = [[DOC_ID_1, DOC_ID_2]] + sentences = [["Une longue phrase.", "Oui."], ["Non.", "Une phrase."]] + translator = _UpperTranslator() + sentence_splitter = MockSentenceSplitter(sentences) + captured = [] + update_doc = partial(_capturing_es_update, captured=captured) + monkeypatch.setattr(activities, "_update_docs_translation", update_doc) + es_client = MockESClient([FR_DOC_1, FR_DOC_2]) + # When + with translator.load_cm( + source=DS_FRENCH, target=DS_ENGLISH, worker_config=worker_config + ): + n_translated = await translate_docs_act( + batches, + project=TEST_PROJECT, + es_client=es_client, + worker_config=worker_config, + translator=translator, + sentence_splitter=sentence_splitter, + ) + # Then + assert n_translated == 2 + contents = {doc.id: translation.content for doc, translation in captured} + assert contents == { + DOC_ID_1: "UNE LONGUE PHRASE. OUI.", + DOC_ID_2: "NON. UNE PHRASE.", + } + + async def test__update_docs() -> None: # Given doc_1 = Document( diff --git a/workers/translation-worker/tests/test_hunyuan_translator.py b/workers/translation-worker/tests/test_hunyuan_translator.py index c99322e7..405fc13e 100644 --- a/workers/translation-worker/tests/test_hunyuan_translator.py +++ b/workers/translation-worker/tests/test_hunyuan_translator.py @@ -68,7 +68,7 @@ def test__translate__wraps_each_text_in_prompt() -> None: ], ], tokenize=True, - add_generation_prompt=False, + add_generation_prompt=True, return_tensors="pt", padding=True, return_dict=True, @@ -134,7 +134,7 @@ def test__load__initialises_tokenizer_and_model_from_config_model_ref_with_devic dummy_tf.AutoModelForCausalLM.from_pretrained.assert_called_once_with( DEFAULT_HUNYUAN_MODEL_REF, device_map="cuda:0", - torch_dtype="bfloat16", + dtype="bfloat16", ) diff --git a/workers/translation-worker/translation_worker/activities.py b/workers/translation-worker/translation_worker/activities.py index e8cb9238..9ff015d3 100644 --- a/workers/translation-worker/translation_worker/activities.py +++ b/workers/translation-worker/translation_worker/activities.py @@ -33,7 +33,7 @@ has_id, has_type, ) -from icij_common.iter_utils import async_batches, before_and_after, once +from icij_common.iter_utils import before_and_after, once from .config import TranslationWorkerConfig from .dependencies import lifespan_sentence_splitter_cache, lifespan_translator_cache @@ -96,7 +96,6 @@ async def translate_docs( ) -> int: es_client = lifespan_es_client() worker_config = cast(TranslationWorkerConfig, lifespan_worker_config()) - # Load the translator first to install the logger.debug("loading %s -> %s translator...", source, target) translator_factory = partial( _load_translator_from_config, @@ -107,7 +106,6 @@ async def translate_docs( ) translator_key = config_cache_key(config.translator) translator_cache = lifespan_translator_cache() - # SBD, then load the splitter splitter_config = config.sentence_splitter_for(source) splitter_factory = partial( _load_splitter_from_config, config=splitter_config, language=source @@ -170,7 +168,7 @@ async def create_translation_batches_act( max_docs: int | None = None, es_client: ESClient | None = None, ) -> AsyncGenerator[tuple[DatashareLanguage, list[Batch]], None]: - # Retrieve unprocessed docs. + # retrieve unprocessed docs query = _with_doc_type(query) es_docs = _get_es_docs_by_language( es_client, project, query, source_includes=_BATCHING_DOC_SOURCES @@ -267,8 +265,6 @@ async def _translate_and_queue( # noqa: PLR0917 ) seen = 0 buffer = [] - current_doc = None - current_doc_translation = [] n_batches = len(batches) for batch_i, doc_ids in enumerate(batches): logger.debug("translating batch %s / %s", batch_i, n_batches) @@ -278,39 +274,55 @@ async def _translate_and_queue( # noqa: PLR0917 body={QUERY: has_id(doc_ids)}, source_includes=_TRANSLATION_DOC_SOURCES, ) - doc_sents = _split_sentences(docs, sentence_splitter) - # TODO: ideally we should aim at having almost constant size batches, - # by using some sort of binarization / binning. That will also improve add - # context to translate short sentences. A split strategy - # like adding min and max batch_item length should help - doc_sent_batches = async_batches(doc_sents, batch_size=worker_config.batch_size) - async for batch in doc_sent_batches: - batch_docs, sents = zip(*batch, strict=False) - # Run translation 1 batch at the time, parallelization is controlled - # via the batch_size - translated_sents = await asyncio.to_thread(translator.translate, sents) - for doc, translated_sent in zip(batch_docs, translated_sents, strict=True): - if current_doc is not None and doc.id != current_doc.id: - translation = translation_factory(content=current_doc_translation) - buffer.append((current_doc, translation)) - if len(buffer) >= worker_config.es_buffer_size: - queue.put_nowait(buffer) - buffer = [] - seen += 1 - if progress is not None: - await progress(seen) - current_doc_translation = [] - current_doc = doc - current_doc_translation.append(translated_sent) + doc_sents = [ds async for ds in _split_sentences(docs, sentence_splitter)] + translated_docs = _translate_by_length( + doc_sents, translator, batch_size=worker_config.batch_size + ) + async for doc, translated_sents in translated_docs: + buffer.append((doc, translation_factory(content=translated_sents))) + if len(buffer) >= worker_config.es_buffer_size: + queue.put_nowait(buffer) + buffer = [] + seen += 1 + if progress is not None: + await progress(seen) logger.debug("batch %s / %s translated !", batch_i, n_batches) - # Empty the buffer - if current_doc_translation: - translation = translation_factory(content=current_doc_translation) - buffer.append((current_doc, translation)) + if buffer: queue.put_nowait(buffer) return n_docs +async def _translate_by_length( + doc_sents: list[tuple[Document, str]], translator: Translator, *, batch_size: int +) -> AsyncGenerator[tuple[Document, list[str]], None]: + docs: dict[str, Document] = {} + translations: dict[str, list[str | None]] = {} + remaining: dict[str, int] = {} + positions = [] + for doc, _ in doc_sents: + if doc.id not in docs: + docs[doc.id] = doc + translations[doc.id] = [] + remaining[doc.id] = 0 + positions.append((doc.id, len(translations[doc.id]))) + translations[doc.id].append(None) + remaining[doc.id] += 1 + by_length = sorted(range(len(doc_sents)), key=lambda i: len(doc_sents[i][1])) + for start in range(0, len(by_length), batch_size): + batch = by_length[start : start + batch_size] + sents = [doc_sents[i][1] for i in batch] + translated_sents = await asyncio.to_thread(translator.translate, sents) + completed = [] + for i, translated_sent in zip(batch, translated_sents, strict=True): + doc_id, position = positions[i] + translations[doc_id][position] = translated_sent + remaining[doc_id] -= 1 + if not remaining[doc_id]: + completed.append(doc_id) + for doc_id in completed: + yield docs[doc_id], translations.pop(doc_id) + + async def _write_translations_to_es( es_client: ESClient, queue: asyncio.Queue, project: str ) -> None: @@ -357,11 +369,11 @@ async def _get_es_docs_by_language( except StopAsyncIteration: return current_language = next_doc[SOURCE][DOC_LANGUAGE] - # Consume the iterator until we find a doc with a different language + # consume the iterator until we find a doc with a different language language_docs, docs = before_and_after( docs, predicate=partial(_has_language, language=current_language) ) - # Group all docs of same language + # group all docs of same language grouped_docs = chain(once(next_doc), language_docs) yield aiter(grouped_docs) diff --git a/workers/translation-worker/translation_worker/translators/hunyuan.py b/workers/translation-worker/translation_worker/translators/hunyuan.py index b17c84c7..7db2a03b 100644 --- a/workers/translation-worker/translation_worker/translators/hunyuan.py +++ b/workers/translation-worker/translation_worker/translators/hunyuan.py @@ -77,7 +77,7 @@ def load( translator = AutoModelForCausalLM.from_pretrained( self._config.model_ref, device_map=self._config.device_map, - torch_dtype=self._config.torch_dtype, + dtype=self._config.torch_dtype, ) self._translator = translator self._tokenizer = AutoTokenizer.from_pretrained( @@ -92,7 +92,7 @@ def translate(self, texts: Iterable[str]) -> list[str]: tokenized = self._tokenizer.apply_chat_template( conversations, tokenize=True, - add_generation_prompt=False, + add_generation_prompt=True, return_tensors="pt", padding=True, return_dict=True,