feature(translation-worker): add rules-based sentence-splitting

workers/translation-worker/pyproject.toml

@@ -32,11 +32,18 @@ worker_config_cls = "translation_worker.config:WORKER_CONFIG_CLS"
 inference = ['datashare-translation-worker[c2translate]']
 c2translate = [
   "argostranslate==1.11.0",
+  "datashare-translation-worker[sat]",
 ]
 torch = [
   "torch==2.11.0",
   "transformers==5.12.1",
   "accelerate==1.14.0",
+  "datashare-translation-worker[sat]",
+]
+# Sentence splitting for languages without sentence punctuation (Thai), runs on the
+# torch installed by the c2translate (through argostranslate) or torch extras
+sat = [
+  "wtpsplit==2.2.2",
 ]
 
 [tool.uv.sources]

workers/translation-worker/tests/test_activities.py

@@ -30,6 +30,7 @@
 from translation_worker.config import TranslationWorkerConfig
 from translation_worker.objects import (
     ArgosTranslatorConfig,
+    DefaultSentenceSplitterConfig,
     TranslationModel,
     untranslated_query,
 )
@@ -71,11 +72,12 @@ async def poll_search_pages(
 
 class MockSentenceSplitter(SentenceSplitter):
     def __init__(self, splits: list[list[str]]):
+        super().__init__(DefaultSentenceSplitterConfig())
         self._splits = iter(splits)
 
     def load(self, language: Language) -> Self: ...
 
-    def split_sentences(self, text: str) -> list[str]:  # noqa: ARG002
+    def _split_sentences(self, text: str) -> list[str]:  # noqa: ARG002
         return next(self._splits)
 
     @classmethod

workers/translation-worker/tests/test_objects.py

@@ -2,17 +2,24 @@
 
 import pytest
 from datashare_python.conftest import TEST_PROJECT
+from datashare_python.objects import DatashareLanguage, Language
 from translation_worker.objects import (
     ArgosSentenceSplitterConfig,
     ArgosSentencizer,
     ArgosTranslatorConfig,
+    BaseTranslatorConfig,
+    DefaultSentenceSplitterConfig,
     DocId,
     DocumentSearchQuery,
+    HunyuanMtTranslatorConfig,
+    RulesSentenceSplitterConfig,
+    SaTSentenceSplitterConfig,
+    SentenceSplitterConfig,
     TranslationArgs,
     TranslationConfig,
 )
 
-from .conftest import DS_ENGLISH
+from .conftest import DS_CHINESE, DS_ENGLISH
 
 
 def test_config_deser() -> None:
@@ -31,6 +38,114 @@ def test_config_deser() -> None:
     assert deser == expected
 
 
+def test_config_deser_defaults() -> None:
+    # When
+    deser = TranslationConfig.model_validate({})
+    # Then
+    assert deser.sentence_splitter == RulesSentenceSplitterConfig()
+    assert deser.language_sentence_splitters == {"th": SaTSentenceSplitterConfig()}
+
+
+def test_config_deser_language_sentence_splitters() -> None:
+    # Given
+    config = {
+        "languageSentenceSplitters": {
+            "km": {"model": "SAT", "modelRef": "sat-12l-sm"},
+            "GERMAN": {"model": "ARGOS"},
+        }
+    }
+    # When
+    deser = TranslationConfig.model_validate(config)
+    # Then
+    expected = {
+        "km": SaTSentenceSplitterConfig(model_ref="sat-12l-sm"),
+        "GERMAN": ArgosSentenceSplitterConfig(),
+    }
+    assert deser.language_sentence_splitters == expected
+    assert TranslationConfig.model_validate(deser.model_dump()) == deser
+
+
+@pytest.mark.parametrize(
+    ("overrides", "language", "expected"),
+    [
+        (None, DatashareLanguage("THAI"), SaTSentenceSplitterConfig()),
+        (None, "th", SaTSentenceSplitterConfig()),
+        (None, DS_ENGLISH, RulesSentenceSplitterConfig()),
+        (None, DS_CHINESE, RulesSentenceSplitterConfig()),
+        (
+            {"zh": DefaultSentenceSplitterConfig()},
+            DS_CHINESE,
+            DefaultSentenceSplitterConfig(),
+        ),
+        (
+            {"zh": DefaultSentenceSplitterConfig()},
+            "zh-Hant",
+            DefaultSentenceSplitterConfig(),
+        ),
+        (
+            {"ENGLISH": SaTSentenceSplitterConfig()},
+            "en-US",
+            SaTSentenceSplitterConfig(),
+        ),
+        ({}, DatashareLanguage("THAI"), RulesSentenceSplitterConfig()),
+    ],
+)
+def test_config_sentence_splitter_for(
+    overrides: dict | None, language: Language, expected: SentenceSplitterConfig
+) -> None:
+    # Given
+    config = TranslationConfig()
+    if overrides is not None:
+        config = TranslationConfig(language_sentence_splitters=overrides)
+    # When
+    splitter_config = config.sentence_splitter_for(language)
+    # Then
+    assert splitter_config.model_copy(update={"max_length": None}) == expected
+
+
+@pytest.mark.parametrize(
+    ("splitter", "translator", "expected_max_length"),
+    [
+        (RulesSentenceSplitterConfig(), ArgosTranslatorConfig(), 1024),
+        (RulesSentenceSplitterConfig(), HunyuanMtTranslatorConfig(), 1024),
+        (
+            RulesSentenceSplitterConfig(),
+            HunyuanMtTranslatorConfig(max_new_tokens=512),
+            256,
+        ),
+        # explicit lengths take precedence
+        (
+            RulesSentenceSplitterConfig(max_length=100),
+            HunyuanMtTranslatorConfig(max_new_tokens=512),
+            100,
+        ),
+    ],
+)
+def test_config_sentence_splitter_for_max_length(
+    splitter: SentenceSplitterConfig,
+    translator: BaseTranslatorConfig,
+    expected_max_length: int,
+) -> None:
+    # Given
+    config = TranslationConfig(sentence_splitter=splitter, translator=translator)
+    # When
+    splitter_config = config.sentence_splitter_for(DS_CHINESE)
+    # Then
+    assert splitter_config.max_length == expected_max_length
+
+
+def test_config_sentence_splitter_for_language_max_length() -> None:
+    # Given
+    config = TranslationConfig(
+        language_sentence_splitters={"zh": SaTSentenceSplitterConfig()},
+        translator=HunyuanMtTranslatorConfig(max_new_tokens=300),
+    )
+    # When
+    splitter_config = config.sentence_splitter_for(DS_CHINESE)
+    # Then
+    assert splitter_config == SaTSentenceSplitterConfig(max_length=150)
+
+
 _NO_ES_TRANSLATED_CONTENT = {
     "bool": {
         "must_not": [