feature[hunyuan translation]: translate chinese dialects with tencent…

datashare-python/datashare_python/dependencies.py

@@ -12,6 +12,7 @@
 from .config import LogLevel, WorkerConfig
 from .exceptions import DependencyInjectionError
 from .logging_ import setup_worker_loggers
+from .objects import Shared
 from .task_client import DatashareTaskClient
 from .types_ import ContextManagerFactory, TemporalClient
 
@@ -23,6 +24,7 @@
 TASK_CLIENT: ContextVar[DatashareTaskClient] = ContextVar("task_client")
 TEMPORAL_CLIENT: ContextVar[TemporalClient] = ContextVar("temporal_client")
 WORKER_CONFIG: ContextVar[WorkerConfig] = ContextVar("worker_config")
+SHARED: ContextVar[Shared] = ContextVar("shared")
 
 
 def set_event_loop(event_loop: AbstractEventLoop) -> None:
@@ -99,6 +101,20 @@ def lifespan_temporal_client() -> TemporalClient:
         raise DependencyInjectionError("temporal client") from e
 
 
+# Setup shared resources
+async def set_shared_resources(shared: Shared) -> Shared:
+    SHARED.set(shared)
+    return shared
+
+
+# Return shared resources
+def shared_resources() -> Shared:
+    try:
+        return SHARED.get()
+    except LookupError as e:
+        raise DependencyInjectionError("shared resources") from e
+
+
 @asynccontextmanager
 async def with_dependencies(
     dependencies: list[ContextManagerFactory], **kwargs

datashare-python/datashare_python/objects.py

@@ -1,7 +1,8 @@
 import logging
 import os
+from asyncio import Lock
 from collections.abc import Awaitable, Callable
-from dataclasses import dataclass
+from dataclasses import dataclass, field
 from datetime import UTC, datetime
 from enum import StrEnum, unique
 from io import BytesIO
@@ -367,3 +368,26 @@ class TaskGroup:
     @classmethod
     def python(cls) -> Self:
         return cls(name="PYTHON")
+
+
+@dataclass(frozen=True)
+class Shared:
+    _resources: dict[str, Any] = field(default_factory=dict)
+    _lock: Lock = field(default_factory=Lock)
+
+    def get_resource(self, key: str, default: Any = None) -> Any:
+        return self._resources.get(key, default)
+
+    def set_resource(self, key: str, value: Any) -> None:
+        self._resources[key] = value
+
+    def pop_resource(self, key: str, default: Any = None) -> Any:
+        return self._resources.pop(key, default)
+
+    async def async_set_resource(self, key: str, value: Any) -> None:
+        async with self._lock:
+            self._resources[key] = value
+
+    async def async_pop_resource(self, key: str, default: Any = None) -> Any:
+        async with self._lock:
+            return self._resources.pop(key, default)

datashare-python/tests/test_dependencies.py

@@ -1,7 +1,72 @@
 from typing import Any
 
 import pytest
-from datashare_python.dependencies import add_missing_args
+from datashare_python.dependencies import (
+    add_missing_args,
+    set_shared_resources,
+    shared_resources,
+)
+from datashare_python.exceptions import DependencyInjectionError
+from datashare_python.objects import Shared
+
+
+async def test_set_shared_resources_and_get() -> None:
+    shared = Shared()
+    shared.set_resource("k", "v")
+    returned = await set_shared_resources(shared)
+    assert returned is shared
+    assert shared_resources() is shared
+
+
+async def test_set_shared_resources_overwrites() -> None:
+    first = Shared()
+    second = Shared()
+    first.set_resource("k", "v1")
+    second.set_resource("k", "v2")
+    await set_shared_resources(first)
+    await set_shared_resources(second)
+    assert shared_resources() is second
+
+
+def test_shared_resources_raises_before_set() -> None:
+    with pytest.raises(DependencyInjectionError):
+        shared_resources()
+
+
+def test_shared_resources_field_is_frozen() -> None:
+    shared = Shared()
+    shared.set_resource("k", "v")
+    with pytest.raises(Exception):  # noqa: B017, PT011
+        shared._resources = {}
+
+
+def test_get_resource_existing_key() -> None:
+    shared = Shared()
+    shared.set_resource("k", "v")
+    assert shared.get_resource("k") == "v"
+
+
+def test_get_resource_missing_key_with_default() -> None:
+    shared = Shared()
+    assert shared.get_resource("missing", "default") == "default"
+
+
+async def test_async_set_resource() -> None:
+    shared = Shared()
+    await shared.async_set_resource("k", "v")
+    assert shared.get_resource("k") == "v"
+
+
+async def test_async_set_resource_overwrites() -> None:
+    shared = Shared()
+    shared.set_resource("k", "old")
+    await shared.async_set_resource("k", "new")
+    assert shared.get_resource("k") == "new"
+
+
+async def test_async_pop_resource_missing_key_with_default() -> None:
+    shared = Shared()
+    assert await shared.async_pop_resource("missing", None) is None
 
 
 @pytest.mark.parametrize(

workers/translation-worker/pyproject.toml

@@ -12,6 +12,7 @@ requires-python = ">=3.11.0, <3.13"
 
 dependencies = [
   "datashare-python~=0.8.20",
+  "language-data>=1.4.0",
   "pydantic-extra-types[pycountry]==2.11.1",
 ]
 
@@ -28,9 +29,14 @@ dependencies = "translation_worker.dependencies:REGISTRY"
 worker_config_cls = "translation_worker.config:WORKER_CONFIG_CLS"
 
 [project.optional-dependencies]
-inference = [
+argos_inference = [
   "argostranslate==1.11.0",
 ]
+hunyuan_inference = [
+  "torch>=2.11.0",
+  "transformers>=5.12.1",
+  "accelerate>=1.14.0",
+]
 
 [tool.uv.sources]
 datashare-python = { path = "../../datashare-python", editable = true }

workers/translation-worker/tests/conftest.py

@@ -24,21 +24,31 @@
     test_worker_config,
     worker_lifetime_deps,
 )
-from datashare_python.objects import DatashareLanguage, Document
+from datashare_python.dependencies import SHARED, shared_resources
+from datashare_python.objects import DatashareLanguage, Document, Shared
 from datashare_python.types_ import ContextManagerFactory, TemporalClient
 from datashare_python.worker import worker_context
 from icij_common.es import ESClient
 from translation_worker.activities import TranslationActivities
 from translation_worker.config import TranslationWorkerConfig
-from translation_worker.dependencies import set_es_client, set_worker_config
+from translation_worker.dependencies import (
+    set_es_client,
+    set_hunyuan_translator,
+    set_worker_config,
+)
 from translation_worker.workflows import TaskQueue, TranslationWorkflow
 
 
 @pytest.fixture(scope="session")
-def test_deps() -> list[ContextManagerFactory]:
+def test_io_deps() -> list[ContextManagerFactory]:
     return [set_worker_config, set_es_client]
 
 
+@pytest.fixture(scope="session")
+def test_inference_deps() -> list[ContextManagerFactory]:
+    return [set_worker_config, set_es_client, set_hunyuan_translator]
+
+
 @pytest.fixture(scope="session")
 def test_worker_config(tmp_path_factory: TempPathFactory) -> TranslationWorkerConfig:  # noqa: ANN001, ARG001, F811
     tmp_path = tmp_path_factory.mktemp("test-")
@@ -67,13 +77,15 @@ def test_worker_config(tmp_path_factory: TempPathFactory) -> TranslationWorkerCo
 ENGLISH = "ENGLISH"
 FRENCH = "FRENCH"
 SPANISH = "SPANISH"
+CHINESE = "CHINESE"
 DOC_ID_1 = "doc_id_1"
 DOC_ID_2 = "doc_id_2"
 ROOT_DOCUMENT_1 = "root_document_1"
 ROOT_DOCUMENT_2 = "root_document_2"
 DS_ENGLISH = DatashareLanguage(ENGLISH)
 DS_FRENCH = DatashareLanguage(FRENCH)
 DS_SPANISH = DatashareLanguage(SPANISH)
+DS_CHINESE = DatashareLanguage(CHINESE)
 
 
 FRENCH_TEXT = (
@@ -87,6 +99,10 @@ def test_worker_config(tmp_path_factory: TempPathFactory) -> TranslationWorkerCo
     "vez. Besame.... Besame mucho.... Que tengo miedo perderte, "
     "perderte otra vez."
 )
+CHINESE_TEXT = (
+    "北京是中国的首都。这是一个美丽的城市，有着悠久的历史和文化。"
+    "每年都有数以百万计的游客来这里参观故宫和长城。"
+)
 
 
 def _create_doc(
@@ -110,6 +126,16 @@ async def index_translation_documents(
     return docs
 
 
+@pytest.fixture
+async def index_chinese_translation_documents(
+    test_es_client: ESClient,  # noqa: F811
+) -> list[Document]:
+    docs = [_create_doc(DOC_ID_1, ROOT_DOCUMENT_1, CHINESE_TEXT, DS_CHINESE)]
+    async for _ in index_docs(test_es_client, docs=docs, index_name=TEST_PROJECT):
+        pass
+    return docs
+
+
 @pytest.fixture(scope="session")
 def translation_worker_config() -> TranslationWorkerConfig:
     return TranslationWorkerConfig()
@@ -142,7 +168,7 @@ async def io_worker(
     test_worker_config: TranslationWorkerConfig,  # noqa: F811
     test_temporal_client_session: TemporalClient,  # noqa: F811
     event_loop: asyncio.AbstractEventLoop,  # noqa: F811
-    test_deps: list[ContextManagerFactory],  # noqa: F811
+    test_io_deps: list[ContextManagerFactory],  # noqa: F811
 ) -> AsyncGenerator[None, None]:
     client = test_temporal_client_session
     worker_id = f"test-translation-io-worker-{uuid.uuid4()}"
@@ -163,7 +189,7 @@ async def io_worker(
         client=client,
         event_loop=event_loop,
         task_queue=task_queue,
-        dependencies=test_deps,
+        dependencies=test_io_deps,
     )
     async with worker_ctx:
         yield
@@ -174,7 +200,7 @@ async def translation_inference_worker(
     test_worker_config: TranslationWorkerConfig,  # noqa: F811
     test_temporal_client_session: TemporalClient,  # noqa: F811
     event_loop: asyncio.AbstractEventLoop,  # noqa: F811
-    test_deps: list[ContextManagerFactory],  # noqa: F811
+    test_inference_deps: list[ContextManagerFactory],  # noqa: F811
 ) -> AsyncGenerator[None, None]:
     client = test_temporal_client_session
     worker_id = f"test-translation-cpu-worker-{uuid.uuid4()}"
@@ -190,7 +216,13 @@ async def translation_inference_worker(
         client=client,
         event_loop=event_loop,
         task_queue=task_queue,
-        dependencies=test_deps,
+        dependencies=test_inference_deps,
     )
     async with worker_ctx:
         yield
+
+
+@pytest.fixture(scope="session")
+def test_shared_resources() -> Shared:
+    SHARED.set(Shared())
+    return shared_resources()

workers/translation-worker/tests/test_activities.py

@@ -26,9 +26,13 @@
     create_translation_batches_act,
     translate_docs_act,
 )
-from translation_worker.config import TranslationWorkerConfig
+from translation_worker.config import (
+    ArgosTranslatorConfig,
+    TranslationModel,
+    TranslationWorkerConfig,
+)
 from translation_worker.constants import DOC_CONTENT_TEXT_LENGTH
-from translation_worker.objects import TranslationModel, untranslated_query
+from translation_worker.objects import untranslated_query
 from translation_worker.processors import SentenceSplitter, Translator
 
 from tests.conftest import (
@@ -80,7 +84,7 @@ class MockTranslator(Translator):
     registered_name = TranslationModel.ARGOS
 
     def __init__(self, translations: list[str]):
-        super().__init__()
+        super().__init__(ArgosTranslatorConfig())
         self._translations = translations
 
     def translate(self, texts: Iterable[str]) -> list[str]:  # noqa: ARG002