refactor(translation-worker): sentence batching by length for hunyuan translation

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(

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",
     )
 
 

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)
 

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,