feature(translation): add independent batcher + tag hunyuan prompt

workers/translation-worker/tests/test_activities.py

@@ -1,4 +1,4 @@
-from collections.abc import AsyncGenerator, Iterable
+from collections.abc import AsyncGenerator, Callable, Iterable
 from functools import partial
 from typing import Any, Self, TypeVar
 from unittest.mock import patch
@@ -24,6 +24,7 @@
     _get_es_docs_by_language,
     _split_sentences,
     _translate_by_length,
+    _translate_in_order,
     _update_docs_translation,
     create_translation_batches_act,
     translate_docs_act,
@@ -32,6 +33,7 @@
 from translation_worker.objects import (
     ArgosTranslatorConfig,
     DefaultSentenceSplitterConfig,
+    SentenceBatcher,
     TranslationModel,
     untranslated_query,
 )
@@ -166,18 +168,28 @@ async def _aiter[T](it: Iterable[T]) -> AsyncGenerator[T, None]:
             [_EN_DOC_TO_SPLIT],
             [["Hello world.", "How are you?", "I'm fine."]],
             [
-                (Document.from_es(_EN_DOC_TO_SPLIT), "Hello world."),
-                (Document.from_es(_EN_DOC_TO_SPLIT), "How are you?"),
-                (Document.from_es(_EN_DOC_TO_SPLIT), "I'm fine."),
+                (
+                    Document.from_es(_EN_DOC_TO_SPLIT),
+                    ["Hello world.", "How are you?", "I'm fine."],
+                ),
             ],
         ),
         # multiple docs sentences
         (
             [EN_DOC_1, EN_DOC_2],
             [[EN_DOC_1_TEXT], [EN_DOC_2_TEXT]],
             [
-                (Document.from_es(EN_DOC_1), EN_DOC_1_TEXT),
-                (Document.from_es(EN_DOC_2), EN_DOC_2_TEXT),
+                (Document.from_es(EN_DOC_1), [EN_DOC_1_TEXT]),
+                (Document.from_es(EN_DOC_2), [EN_DOC_2_TEXT]),
+            ],
+        ),
+        # docs without sentences are kept
+        (
+            [EN_DOC_1, EN_DOC_2],
+            [[], [EN_DOC_2_TEXT]],
+            [
+                (Document.from_es(EN_DOC_1), []),
+                (Document.from_es(EN_DOC_2), [EN_DOC_2_TEXT]),
             ],
         ),
         # no doc, no sentences
@@ -187,7 +199,7 @@ async def _aiter[T](it: Iterable[T]) -> AsyncGenerator[T, None]:
 async def test__iter_sentences__yields_sentences_with_correct_indices(
     docs: list[dict[str, Any]],
     sentences: list[list[str]],
-    expected_doc_sents: list[tuple[Document, str]],
+    expected_doc_sents: list[tuple[Document, list[str]]],
 ) -> None:
     # Given
     sentence_splitter = MockSentenceSplitter(sentences)
@@ -501,16 +513,13 @@ async def test__translate_by_length__batches_by_length_and_restores_order() -> N
     # 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"),
+        (doc_1, ["a long sentence", "short"]),
+        (doc_2, ["the longest sentence", "tiny", "mid sized"]),
     ]
     translator = _UpperTranslator()
     # When
     translated = await _collect_async(
-        _translate_by_length(doc_sents, translator, batch_size=2)
+        _translate_by_length(_aiter(doc_sents), translator, batch_size=2)
     )
     # Then
     assert translator.batches == [
@@ -525,8 +534,87 @@ async def test__translate_by_length__batches_by_length_and_restores_order() -> N
     ]
 
 
+async def test__translate_in_order__batches_in_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", "short", "the longest sentence"]),
+        (doc_2, ["tiny", "mid sized"]),
+    ]
+    translator = _UpperTranslator()
+    # When
+    translated = await _collect_async(
+        _translate_in_order(_aiter(doc_sents), translator, batch_size=2)
+    )
+    # Then
+    assert translator.batches == [
+        ["a long sentence", "short"],
+        ["the longest sentence", "tiny"],
+        ["mid sized"],
+    ]
+    assert translated == [
+        (doc_1, ["A LONG SENTENCE", "SHORT", "THE LONGEST SENTENCE"]),
+        (doc_2, ["TINY", "MID SIZED"]),
+    ]
+
+
+async def test__translate_in_order__streams_sentences() -> None:
+    # Given
+    doc = Document.from_es(FR_DOC_1)
+    n_read = 0
+
+    async def doc_sents() -> AsyncGenerator[tuple[Document, list[str]], None]:
+        nonlocal n_read
+        for sents in [["one", "two"], ["three", "four"], ["five", "six", "seven"]]:
+            n_read += 1
+            yield doc, sents
+
+    class _ReadTranslator(_UpperTranslator):
+        def __init__(self):
+            super().__init__()
+            self.n_read = []
+
+        def translate(self, texts: Iterable[str]) -> list[str]:
+            self.n_read.append(n_read)
+            return super().translate(texts)
+
+    translator = _ReadTranslator()
+    # When
+    await _collect_async(_translate_in_order(doc_sents(), translator, batch_size=3))
+    # Then
+    # batches are translated before the next docs are read
+    assert translator.n_read == [2, 3, 3]
+    assert translator.batches == [
+        ["one", "two", "three"],
+        ["four", "five", "six"],
+        ["seven"],
+    ]
+
+
+@pytest.mark.parametrize(
+    "translate_sentences", [_translate_by_length, _translate_in_order]
+)
+async def test__translate_sentences__yields_docs_without_sentences(
+    translate_sentences: Callable,
+) -> None:
+    # Given
+    doc_1, doc_2 = Document.from_es(FR_DOC_1), Document.from_es(FR_DOC_2)
+    doc_sents = [(doc_1, []), (doc_2, ["oui"])]
+    translator = _UpperTranslator()
+    # When
+    translated = await _collect_async(
+        translate_sentences(_aiter(doc_sents), translator, batch_size=2)
+    )
+    # Then
+    assert translator.batches == [["oui"]]
+    assert translated == [(doc_1, []), (doc_2, ["OUI"])]
+
+
+@pytest.mark.parametrize("sentence_batcher", list(SentenceBatcher))
 async def test_translate_docs_act__es_update_keeps_sentence_order(
-    monkeypatch: MonkeyPatch, test_worker_config: TranslationWorkerConfig
+    monkeypatch: MonkeyPatch,
+    test_worker_config: TranslationWorkerConfig,
+    sentence_batcher: SentenceBatcher,
 ) -> None:
     # Given
     worker_config = test_worker_config
@@ -549,6 +637,7 @@ async def test_translate_docs_act__es_update_keeps_sentence_order(
             worker_config=worker_config,
             translator=translator,
             sentence_splitter=sentence_splitter,
+            sentence_batcher=sentence_batcher,
         )
     # Then
     assert n_translated == 2
@@ -559,6 +648,48 @@ async def test_translate_docs_act__es_update_keeps_sentence_order(
     }
 
 
+@pytest.mark.parametrize("sentence_batcher", list(SentenceBatcher))
+async def test_translate_docs_act__es_update_docs_without_sentences(
+    monkeypatch: MonkeyPatch,
+    test_worker_config: TranslationWorkerConfig,
+    sentence_batcher: SentenceBatcher,
+) -> None:
+    # Given
+    worker_config = test_worker_config
+    batches = [[DOC_ID_1, DOC_ID_2]]
+    sentences = [[], ["Oui."]]
+    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])
+    progresses = []
+
+    async def progress(p: float, *, force: bool = False) -> None:  # noqa: ARG001
+        progresses.append(p)
+
+    # 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,
+            sentence_batcher=sentence_batcher,
+            progress=progress,
+        )
+    # Then
+    assert n_translated == 2
+    contents = {doc.id: translation.content for doc, translation in captured}
+    assert contents == {DOC_ID_1: "", DOC_ID_2: "OUI."}
+    assert progresses[-1] == 1.0
+
+
 async def test__update_docs() -> None:
     # Given
     doc_1 = Document(

workers/translation-worker/tests/test_hunyuan_translator.py

@@ -9,6 +9,7 @@
 from translation_worker.translators.hunyuan import (
     HunyuanMtTranslator,
     _message_template,
+    _strip_target_tags,
 )
 
 from tests.conftest import DS_CHINESE, DS_ENGLISH, DS_FRENCH
@@ -53,17 +54,13 @@ def test__translate__wraps_each_text_in_prompt() -> None:
             [
                 {
                     "role": "user",
-                    "content": "将以下文本翻译为英语，"
-                    "注意只需要输出翻译后的结果，不要额外解释：\n\n"
-                    f"{TEXT_TO_TRANSLATE}",
+                    "content": _ZH_PROMPT.format("英语", TEXT_TO_TRANSLATE),
                 }
             ],
             [
                 {
                     "role": "user",
-                    "content": "将以下文本翻译为英语，"
-                    "注意只需要输出翻译后的结果，不要额外解释：\n\n"
-                    f"{MORE_TEXT_TO_TRANSLATE}",
+                    "content": _ZH_PROMPT.format("英语", MORE_TEXT_TO_TRANSLATE),
                 }
             ],
         ],
@@ -114,6 +111,18 @@ def test__translate__decodes_only_generated_tokens() -> None:
     assert result == [TRANSLATED_TEXT]
 
 
+def test__translate__strips_target_tags() -> None:
+    translator = _translator()
+    translator._tokenizer = MagicMock()
+    translator._translator = MagicMock()
+    input_ids = translator._tokenizer.apply_chat_template.return_value["input_ids"]
+    input_ids.to.return_value.shape = (1, 0)
+    translator._translator.generate.return_value = [[0]]
+    translator._tokenizer.decode.return_value = f"<target>{TRANSLATED_TEXT}</target>"
+    result = translator.translate([TEXT_TO_TRANSLATE])
+    assert result == [TRANSLATED_TEXT]
+
+
 def test__load__initialises_tokenizer_and_model_from_config_model_ref_with_device_map_and_dtype() -> (  # noqa: E501
     None
 ):
@@ -138,9 +147,13 @@ def test__load__initialises_tokenizer_and_model_from_config_model_ref_with_devic
     )
 
 
-_ZH_PROMPT = "将以下文本翻译为{}，注意只需要输出翻译后的结果，不要额外解释：\n\n{}"
+_ZH_PROMPT = (
+    "将以下<source></source>之间的文本翻译为{}，注意只需要输出翻译后的结果，"
+    "不要额外解释。输出格式为：<target>str</target>\n\n<source>{}</source>"
+)
 _XX_PROMPT = (
-    "Translate the following segment into {}, without additional explanation.\n\n{}"  # noqa: E501
+    "Translate the text between <source></source> into {}, without additional "
+    "explanation. The output format is: <target>str</target>\n\n<source>{}</source>"
 )
 
 
@@ -164,3 +177,25 @@ def test__message_template__picks_prompt_from_language_pair(
 ) -> None:
     message = _message_template(TEXT_TO_TRANSLATE, source, target)
     assert message == {"role": "user", "content": expected}
+
+
+@pytest.mark.parametrize(
+    ("translation", "expected"),
+    [
+        ("translated text", "translated text"),
+        ("<target>translated text</target>", "translated text"),
+        ("translated text</target>", "translated text"),
+        ("<target>translated text", "translated text"),
+        ("  <target>translated text</target>\n", "translated text"),
+        # anything after the closing tag is dropped
+        ('translated text"</target>"', 'translated text"'),
+        # tags inside the translation might come from the source, they're kept
+        (
+            "a <target> tag and a </target> tag</target>",
+            "a <target> tag and a </target> tag",
+        ),
+        ("", ""),
+    ],
+)
+def test__strip_target_tags(translation: str, expected: str) -> None:
+    assert _strip_target_tags(translation) == expected

workers/translation-worker/tests/test_objects.py

@@ -14,6 +14,7 @@
     HunyuanMtTranslatorConfig,
     RulesSentenceSplitterConfig,
     SaTSentenceSplitterConfig,
+    SentenceBatcher,
     SentenceSplitterConfig,
     TranslationArgs,
     TranslationConfig,
@@ -44,6 +45,17 @@ def test_config_deser_defaults() -> None:
     # Then
     assert deser.sentence_splitter == RulesSentenceSplitterConfig()
     assert deser.language_sentence_splitters == {"th": SaTSentenceSplitterConfig()}
+    assert deser.sentence_batcher == SentenceBatcher.BY_LENGTH
+
+
+def test_config_deser_sentence_batcher() -> None:
+    # Given
+    config = {"sentenceBatcher": "IN_ORDER"}
+    # When
+    deser = TranslationConfig.model_validate(config)
+    # Then
+    assert deser.sentence_batcher == SentenceBatcher.IN_ORDER
+    assert TranslationConfig.model_validate(deser.model_dump()) == deser
 
 
 def test_config_deser_language_sentence_splitters() -> None: