diff --git a/workers/translation-worker/tests/test_activities.py b/workers/translation-worker/tests/test_activities.py index 1194c75b..cbcb6d75 100644 --- a/workers/translation-worker/tests/test_activities.py +++ b/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,9 +168,10 @@ 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 @@ -176,8 +179,17 @@ async def _aiter[T](it: Iterable[T]) -> AsyncGenerator[T, None]: [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( diff --git a/workers/translation-worker/tests/test_hunyuan_translator.py b/workers/translation-worker/tests/test_hunyuan_translator.py index 405fc13e..f01af68a 100644 --- a/workers/translation-worker/tests/test_hunyuan_translator.py +++ b/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"{TRANSLATED_TEXT}" + 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 = ( + "将以下之间的文本翻译为{}，注意只需要输出翻译后的结果，" + "不要额外解释。输出格式为：str\n\n{}" +) _XX_PROMPT = ( - "Translate the following segment into {}, without additional explanation.\n\n{}" # noqa: E501 + "Translate the text between into {}, without additional " + "explanation. The output format is: str\n\n{}" ) @@ -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"), + ("translated text", "translated text"), + ("translated text", "translated text"), + ("translated text", "translated text"), + (" translated text\n", "translated text"), + # anything after the closing tag is dropped + ('translated text""', 'translated text"'), + # tags inside the translation might come from the source, they're kept + ( + "a tag and a tag", + "a tag and a tag", + ), + ("", ""), + ], +) +def test__strip_target_tags(translation: str, expected: str) -> None: + assert _strip_target_tags(translation) == expected diff --git a/workers/translation-worker/tests/test_objects.py b/workers/translation-worker/tests/test_objects.py index ea055daa..5857bcb0 100644 --- a/workers/translation-worker/tests/test_objects.py +++ b/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: diff --git a/workers/translation-worker/translation_worker/activities.py b/workers/translation-worker/translation_worker/activities.py index 9ff015d3..caa93868 100644 --- a/workers/translation-worker/translation_worker/activities.py +++ b/workers/translation-worker/translation_worker/activities.py @@ -1,6 +1,7 @@ import asyncio import logging -from collections.abc import AsyncGenerator, AsyncIterator, Iterable +from collections import deque +from collections.abc import AsyncGenerator, AsyncIterable, AsyncIterator, Iterable from enum import StrEnum from functools import partial from typing import Any, cast @@ -37,7 +38,12 @@ from .config import TranslationWorkerConfig from .dependencies import lifespan_sentence_splitter_cache, lifespan_translator_cache -from .objects import SentenceSplitterConfig, TranslationConfig, TranslatorConfig +from .objects import ( + SentenceBatcher, + SentenceSplitterConfig, + TranslationConfig, + TranslatorConfig, +) from .processors import SentenceSplitter, Translator logger = logging.getLogger(__name__) @@ -128,6 +134,7 @@ async def translate_docs( worker_config=worker_config, translator=translator, sentence_splitter=sentence_splitter, + sentence_batcher=config.sentence_batcher, ) logger.info("done translating !") return n_translated @@ -216,6 +223,7 @@ async def translate_docs_act( worker_config: TranslationWorkerConfig, es_client: ESClient, progress: AsyncProgressRateHandler | None = None, # noqa: F821 + sentence_batcher: SentenceBatcher = SentenceBatcher.BY_LENGTH, ) -> int: # TODO: this should not happen es_queue = asyncio.Queue() @@ -228,6 +236,7 @@ async def translate_docs_act( worker_config, es_client, progress, + sentence_batcher, ) publisher = asyncio.create_task(publisher) publisher_callback = lambda: es_queue.put_nowait(None) # noqa: E731 @@ -251,6 +260,7 @@ async def _translate_and_queue( # noqa: PLR0917 worker_config: TranslationWorkerConfig, es_client: ESClient, progress: AsyncProgressRateHandler | None = None, # noqa: F821 + sentence_batcher: SentenceBatcher = SentenceBatcher.BY_LENGTH, ) -> int: n_docs = sum(len(b) for b in batches) if not n_docs: @@ -263,6 +273,13 @@ async def _translate_and_queue( # noqa: PLR0917 translation_factory = partial( Translation, source_language=source, target_language=target, translator=model ) + match sentence_batcher: + case SentenceBatcher.BY_LENGTH: + translate_sentences = _translate_by_length + case SentenceBatcher.IN_ORDER: + translate_sentences = _translate_in_order + case _: + raise ValueError(f"unsupported sentence batcher {sentence_batcher}") seen = 0 buffer = [] n_batches = len(batches) @@ -274,8 +291,8 @@ async def _translate_and_queue( # noqa: PLR0917 body={QUERY: has_id(doc_ids)}, source_includes=_TRANSLATION_DOC_SOURCES, ) - doc_sents = [ds async for ds in _split_sentences(docs, sentence_splitter)] - translated_docs = _translate_by_length( + doc_sents = _split_sentences(docs, sentence_splitter) + translated_docs = translate_sentences( doc_sents, translator, batch_size=worker_config.batch_size ) async for doc, translated_sents in translated_docs: @@ -293,34 +310,76 @@ async def _translate_and_queue( # noqa: PLR0917 async def _translate_by_length( - doc_sents: list[tuple[Document, str]], translator: Translator, *, batch_size: int + doc_sents: AsyncIterable[tuple[Document, list[str]]], + translator: Translator, + *, + batch_size: int, ) -> AsyncGenerator[tuple[Document, list[str]], None]: + # all sentences are kept in memory to be sorted by length 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) + sents: list[tuple[str, int, str]] = [] + async for doc, doc_sentences in doc_sents: + if not doc_sentences: + yield doc, [] + continue + docs[doc.id] = doc + translations[doc.id] = [None] * len(doc_sentences) + remaining[doc.id] = len(doc_sentences) + sents.extend((doc.id, i, sent) for i, sent in enumerate(doc_sentences)) + sents.sort(key=lambda s: len(s[2])) + for start in range(0, len(sents), batch_size): + batch = sents[start : start + batch_size] + translated_sents = await asyncio.to_thread( + translator.translate, [sent for _, _, sent in batch] + ) completed = [] - for i, translated_sent in zip(batch, translated_sents, strict=True): - doc_id, position = positions[i] - translations[doc_id][position] = translated_sent + for (doc_id, i, _), translated_sent in zip( + batch, translated_sents, strict=True + ): + translations[doc_id][i] = 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) + yield docs.pop(doc_id), translations.pop(doc_id) + + +async def _translate_in_order( + doc_sents: AsyncIterable[tuple[Document, list[str]]], + translator: Translator, + *, + batch_size: int, +) -> AsyncGenerator[tuple[Document, list[str]], None]: + # sentences are streamed, only the current doc's sentences and batch are held in + # memory + pending: deque[tuple[Document, int]] = deque() + to_translate: list[str] = [] + translated: list[str] = [] + async for doc, doc_sentences in doc_sents: + pending.append((doc, len(doc_sentences))) + to_translate.extend(doc_sentences) + while len(to_translate) >= batch_size: + batch, to_translate = to_translate[:batch_size], to_translate[batch_size:] + translated += await asyncio.to_thread(translator.translate, batch) + for translated_doc in _pop_translated_docs(pending, translated): + yield translated_doc + if to_translate: + translated += await asyncio.to_thread(translator.translate, to_translate) + for translated_doc in _pop_translated_docs(pending, translated): + yield translated_doc + + +def _pop_translated_docs( + pending: deque[tuple[Document, int]], translated: list[str] +) -> list[tuple[Document, list[str]]]: + docs = [] + while pending and pending[0][1] <= len(translated): + doc, n_sents = pending.popleft() + docs.append((doc, translated[:n_sents])) + del translated[:n_sents] + return docs async def _write_translations_to_es( @@ -340,14 +399,13 @@ async def _write_translations_to_es( async def _split_sentences( doc_iter: AsyncGenerator[dict, None], sentence_splitter: SentenceSplitter -) -> AsyncGenerator[tuple[Document, str], None]: +) -> AsyncGenerator[tuple[Document, list[str]], None]: async for doc in doc_iter: es_doc = Document.from_es(doc) sentences = await asyncio.to_thread( sentence_splitter.split_sentences, es_doc.content ) - for sentence in sentences: - yield es_doc, sentence + yield es_doc, sentences async def _get_es_docs_by_language( diff --git a/workers/translation-worker/translation_worker/objects.py b/workers/translation-worker/translation_worker/objects.py index fe2dc9f3..15b991d9 100644 --- a/workers/translation-worker/translation_worker/objects.py +++ b/workers/translation-worker/translation_worker/objects.py @@ -40,6 +40,13 @@ class SentenceSplitterModel(StrEnum): SAT = "SAT" +class SentenceBatcher(StrEnum): + # sort and batch sentences by length to minimize padding + BY_LENGTH = "BY_LENGTH" + # stream sentences in document order + IN_ORDER = "IN_ORDER" + + class TranslationModel(StrEnum): ARGOS = "ARGOS" HUNYUAN = "HUNYUAN" @@ -206,6 +213,7 @@ class TranslationConfig(DatashareModel): Language, Annotated[_SentenceSplitterConfig, Discriminator(splitter_discriminator)], ] = Field(default_factory=_default_language_sentence_splitters) + sentence_batcher: SentenceBatcher = SentenceBatcher.BY_LENGTH translator: TranslatorConfig = Field( discriminator=Discriminator(translator_discriminator), default_factory=ArgosTranslatorConfig, diff --git a/workers/translation-worker/translation_worker/translators/hunyuan.py b/workers/translation-worker/translation_worker/translators/hunyuan.py index 7db2a03b..80ee499a 100644 --- a/workers/translation-worker/translation_worker/translators/hunyuan.py +++ b/workers/translation-worker/translation_worker/translators/hunyuan.py @@ -28,22 +28,40 @@ def _message_template(text: str, source: Language, target: Language) -> dict[str is_chinese_pair = bool( {source_tag.language, target_tag.language} & _CHINESE_LANGUAGES ) - # tencent recommends using a chinese prompt when translating from chinese + # tencent recommends using a chinese prompt when translating from chinese. + # tags are taken from tencent's guidelines if is_chinese_pair: target_lang = target_tag.display_name("zh") context = ( - f"将以下文本翻译为{target_lang}，注意只需要输出翻译后的结果，不要额外解释：" - f"\n\n{text}" + f"将以下之间的文本翻译为{target_lang}，" + "注意只需要输出翻译后的结果，不要额外解释。输出格式为：str" + f"\n\n{text}" ) else: target_lang = target_tag.display_name("en") context = ( - f"Translate the following segment into {target_lang}, without additional " - f"explanation.\n\n{text}" + f"Translate the text between into {target_lang}, " + "without additional explanation. The output format is: " + f"str\n\n{text}" ) return {"role": "user", "content": context} +_TARGET_START = "" +_TARGET_END = "" + + +def _strip_target_tags(translation: str) -> str: + # strip erroneous target tags (often included by hunyuan) from translations. + # only strip affixed tags; tags inside content may be native. + translation = translation.strip() + translation = translation.removeprefix(_TARGET_START) + end = translation.rfind(_TARGET_END) + if end != -1: + translation = translation[:end] + return translation.strip() + + @Translator.register(TranslationModel.HUNYUAN) class HunyuanMtTranslator(Translator): def __init__( @@ -109,11 +127,13 @@ def translate(self, texts: Iterable[str]) -> list[str]: temperature=self._config.temperature, repetition_penalty=self._config.repetition_penalty, ) - # generate returns the prompt followed by the completion, only decode the + # generate returns the prompt followed by the completion. only decode the # completion prompt_length = input_ids.shape[-1] return [ - self._tokenizer.decode(out[prompt_length:], skip_special_tokens=True) + _strip_target_tags( + self._tokenizer.decode(out[prompt_length:], skip_special_tokens=True) + ) for out in outputs ]