WIP: Add NmtFileService and reading/writing of per-text cached analysis

Author: pmachapman

machine/jobs/__init__.py

@@ -1,6 +1,7 @@
 from .clearml_shared_file_service import ClearMLSharedFileService
 from .local_shared_file_service import LocalSharedFileService
 from .nmt_engine_build_job import NmtEngineBuildJob
+from .nmt_file_service import CachedAnalysisInfo, NmtFileService
 from .nmt_model_factory import NmtModelFactory
 from .shared_file_service_base import DictToJsonWriter, SharedFileServiceBase
 from .smt_engine_build_job import SmtEngineBuildJob
@@ -13,9 +14,11 @@
 from .word_alignment_model_factory import WordAlignmentModelFactory
 
 __all__ = [
+    "CachedAnalysisInfo",
     "ClearMLSharedFileService",
     "LocalSharedFileService",
     "NmtEngineBuildJob",
+    "NmtFileService",
     "NmtModelFactory",
     "DictToJsonWriter",
     "SharedFileServiceBase",

machine/jobs/build_nmt_engine.py

@@ -9,9 +9,9 @@
 from .build_clearml_helper import report_clearml_progress, update_settings
 from .config import SETTINGS
 from .nmt_engine_build_job import NmtEngineBuildJob
+from .nmt_file_service import NmtFileService
 from .nmt_model_factory import NmtModelFactory
 from .shared_file_service_factory import SharedFileServiceType
-from .translation_file_service import TranslationFileService
 
 # Setup logging
 logging.basicConfig(
@@ -47,7 +47,7 @@ def clearml_progress(status: ProgressStatus) -> None:
         logger.info("NMT Engine Build Job started")
         update_settings(SETTINGS, args, task, logger)
 
-        translation_file_service = TranslationFileService(SharedFileServiceType.CLEARML, SETTINGS)
+        translation_file_service = NmtFileService(SharedFileServiceType.CLEARML, SETTINGS)
         nmt_model_factory: NmtModelFactory
         if SETTINGS.model_type == "huggingface":
             from .huggingface.hugging_face_nmt_model_factory import HuggingFaceNmtModelFactory

machine/jobs/nmt_engine_build_job.py

@@ -10,34 +10,36 @@
 from ..utils.phased_progress_reporter import Phase, PhasedProgressReporter
 from ..utils.progress_status import ProgressStatus
 from .eflomal_aligner import EflomalAligner, is_eflomal_available, tokenize
+from .nmt_file_service import NmtFileService
 from .nmt_model_factory import NmtModelFactory
 from .translation_engine_build_job import TranslationEngineBuildJob
-from .translation_file_service import PretranslationInfo, TranslationFileService
+from .translation_file_service import PretranslationInfo
 
 logger = logging.getLogger(__name__)
 
 
 class NmtEngineBuildJob(TranslationEngineBuildJob):
-    def __init__(
-        self, config: Any, nmt_model_factory: NmtModelFactory, translation_file_service: TranslationFileService
-    ) -> None:
+    def __init__(self, config: Any, nmt_model_factory: NmtModelFactory, nmt_file_service: NmtFileService) -> None:
+        self._nmt_file_service = nmt_file_service
         self._nmt_model_factory = nmt_model_factory
         self._nmt_model_factory.init()
-        super().__init__(config, translation_file_service)
+        super().__init__(config, nmt_file_service)
 
     def _get_progress_reporter(
         self, progress: Optional[Callable[[ProgressStatus], None]], corpus_size: int
     ) -> PhasedProgressReporter:
         if corpus_size > 0:
             if self._config.align_pretranslations:
                 phases = [
-                    Phase(message="Training NMT model", percentage=0.8, stage="train"),
+                    Phase(message="Detecting quotation conventions", percentage=0.01),
+                    Phase(message="Training NMT model", percentage=0.79, stage="train"),
                     Phase(message="Pretranslating segments", percentage=0.1, stage="inference"),
                     Phase(message="Aligning segments", percentage=0.1, report_steps=False),
                 ]
             else:
                 phases = [
-                    Phase(message="Training NMT model", percentage=0.9, stage="train"),
+                    Phase(message="Detecting quotation conventions", percentage=0.01),
+                    Phase(message="Training NMT model", percentage=0.89, stage="train"),
                     Phase(message="Pretranslating segments", percentage=0.1, stage="inference"),
                 ]
         else:
@@ -82,6 +84,10 @@ def _train_model(
             if check_canceled is not None:
                 check_canceled()
 
+        logger.info("Detecting quotation conventions")
+        progress_reporter.start_next_phase()
+        # TODO: Detect quotation conventions
+
         logger.info("Training NMT model")
         with (
             progress_reporter.start_next_phase() as phase_progress,

machine/jobs/nmt_file_service.py

@@ -0,0 +1,51 @@
+from contextlib import contextmanager
+from typing import Generator, Iterator, TypedDict
+
+import json_stream
+
+from ..utils.context_managed_generator import ContextManagedGenerator
+from .shared_file_service_base import DictToJsonWriter
+from .translation_file_service import TranslationFileService
+
+
+class CachedAnalysisInfo(TypedDict):
+    corpusId: str  # noqa: N815
+    textId: str  # noqa: N815
+    quoteConvention: str  # noqa: N815
+
+
+SOURCE_CACHED_ANALYSIS_FILENAME = "cached_analysis.src.json"
+TARGET_CACHED_ANALYSIS_FILENAME = "cached_analysis.trg.json"
+
+
+class NmtFileService(TranslationFileService):
+
+    def _get_cached_analysis(self, file_name: str) -> ContextManagedGenerator[CachedAnalysisInfo, None, None]:
+        src_cached_analysis_path = self.shared_file_service.download_file(
+            f"{self.shared_file_service.build_path}/{file_name}"
+        )
+
+        def generator() -> Generator[CachedAnalysisInfo, None, None]:
+            with src_cached_analysis_path.open("r", encoding="utf-8-sig") as file:
+                for pi in json_stream.load(file):
+                    yield CachedAnalysisInfo(
+                        corpusId=pi["corpusId"],
+                        textId=pi["textId"],
+                        quoteConvention=pi["quoteConvention"],
+                    )
+
+        return ContextManagedGenerator(generator())
+
+    def get_source_cached_analysis(self) -> ContextManagedGenerator[CachedAnalysisInfo, None, None]:
+        return self._get_cached_analysis(SOURCE_CACHED_ANALYSIS_FILENAME)
+
+    def get_target_cached_analysis(self) -> ContextManagedGenerator[CachedAnalysisInfo, None, None]:
+        return self._get_cached_analysis(TARGET_CACHED_ANALYSIS_FILENAME)
+
+    @contextmanager
+    def open_source_cached_analysis_writer(self) -> Iterator[DictToJsonWriter]:
+        return self.shared_file_service.open_target_writer(SOURCE_CACHED_ANALYSIS_FILENAME)
+
+    @contextmanager
+    def open_target_cached_analysis_writer(self) -> Iterator[DictToJsonWriter]:
+        return self.shared_file_service.open_target_writer(TARGET_CACHED_ANALYSIS_FILENAME)

tests/jobs/test_nmt_engine_build_job.py

@@ -10,13 +10,7 @@
 
 from machine.annotations import Range
 from machine.corpora import DictionaryTextCorpus
-from machine.jobs import (
-    DictToJsonWriter,
-    NmtEngineBuildJob,
-    NmtModelFactory,
-    PretranslationInfo,
-    TranslationFileService,
-)
+from machine.jobs import DictToJsonWriter, NmtEngineBuildJob, NmtFileService, NmtModelFactory, PretranslationInfo
 from machine.jobs.eflomal_aligner import is_eflomal_available
 from machine.translation import (
     Phrase,
@@ -54,7 +48,8 @@ def test_run(decoy: Decoy) -> None:
         assert pretranslations[0]["sourceTokens"] == []
         assert pretranslations[0]["translationTokens"] == []
         assert len(pretranslations[0]["alignment"]) == 0
-    decoy.verify(env.translation_file_service.save_model(Path("model.tar.gz"), "models/save-model.tar.gz"), times=1)
+
+    decoy.verify(env.nmt_file_service.save_model(Path("model.tar.gz"), "models/save-model.tar.gz"), times=1)
 
 
 def test_cancel(decoy: Decoy) -> None:
@@ -116,12 +111,12 @@ def __init__(self, decoy: Decoy) -> None:
         decoy.when(self.nmt_model_factory.create_engine()).then_return(self.engine)
         decoy.when(self.nmt_model_factory.save_model()).then_return(Path("model.tar.gz"))
 
-        self.translation_file_service = decoy.mock(cls=TranslationFileService)
-        decoy.when(self.translation_file_service.create_source_corpus()).then_return(DictionaryTextCorpus())
-        decoy.when(self.translation_file_service.create_target_corpus()).then_return(DictionaryTextCorpus())
-        decoy.when(self.translation_file_service.exists_source_corpus()).then_return(True)
-        decoy.when(self.translation_file_service.exists_target_corpus()).then_return(True)
-        decoy.when(self.translation_file_service.get_source_pretranslations()).then_do(
+        self.nmt_file_service = decoy.mock(cls=NmtFileService)
+        decoy.when(self.nmt_file_service.create_source_corpus()).then_return(DictionaryTextCorpus())
+        decoy.when(self.nmt_file_service.create_target_corpus()).then_return(DictionaryTextCorpus())
+        decoy.when(self.nmt_file_service.exists_source_corpus()).then_return(True)
+        decoy.when(self.nmt_file_service.exists_target_corpus()).then_return(True)
+        decoy.when(self.nmt_file_service.get_source_pretranslations()).then_do(
             lambda: ContextManagedGenerator(
                 (
                     pi
@@ -150,7 +145,7 @@ def open_target_pretranslation_writer(env: _TestEnvironment) -> Iterator[DictToJ
             file.write("\n]\n")
             env.target_pretranslations = file.getvalue()
 
-        decoy.when(self.translation_file_service.open_target_pretranslation_writer()).then_do(
+        decoy.when(self.nmt_file_service.open_target_pretranslation_writer()).then_do(
             lambda: open_target_pretranslation_writer(self)
         )
 
@@ -165,7 +160,7 @@ def open_target_pretranslation_writer(env: _TestEnvironment) -> Iterator[DictToJ
                 }
             ),
             self.nmt_model_factory,
-            self.translation_file_service,
+            self.nmt_file_service,
         )