Address reviewer comment; add key terms as partial words

Author: Enkidu93

machine/corpora/key_term_row.py machine/corpora/key_term.pymachine/corpora/key_term_row.py renamed to machine/corpora/key_term.py

@@ -3,7 +3,7 @@
 
 from ..utils.typeshed import StrPath
 from .dictionary_text_corpus import DictionaryTextCorpus
-from .key_term_row import KeyTerm
+from .key_term import KeyTerm
 from .memory_text import MemoryText
 from .text_row import TextRow
 from .zip_paratext_project_settings_parser import ZipParatextProjectSettingsParser

machine/corpora/paratext_backup_terms_corpus.py

@@ -9,7 +9,7 @@
 
 from ..scripture.constants import ORIGINAL_VERSIFICATION
 from ..scripture.verse_ref import VerseRef
-from .key_term_row import KeyTerm
+from .key_term import KeyTerm
 from .paratext_project_file_handler import ParatextProjectFileHandler
 from .paratext_project_settings import ParatextProjectSettings
 from .paratext_project_settings_parser_base import ParatextProjectSettingsParserBase

machine/corpora/paratext_project_terms_parser_base.py

@@ -78,11 +78,12 @@ def create_source_tokenizer_trainer(self, corpus: TextCorpus) -> Trainer:
     def create_target_tokenizer_trainer(self, corpus: TextCorpus) -> Trainer:
         return NullTrainer()
 
-    def create_model_trainer(self, corpus: ParallelTextCorpus) -> Trainer:
+    def create_model_trainer(self, corpus: ParallelTextCorpus, terms_corpus: ParallelTextCorpus) -> Trainer:
         return HuggingFaceNmtModelTrainer(
             self._model,
             self._training_args,
             corpus,
+            terms_corpus,
             src_lang=self._config.src_lang,
             tgt_lang=self._config.trg_lang,
             add_unk_src_tokens=self._config.huggingface.tokenizer.add_unk_src_tokens,

machine/jobs/huggingface/hugging_face_nmt_model_factory.py

@@ -59,6 +59,7 @@ def _train_model(
         source_corpus: TextCorpus,
         target_corpus: TextCorpus,
         parallel_corpus: ParallelTextCorpus,
+        parallel_terms_corpus: ParallelTextCorpus,
         progress_reporter: PhasedProgressReporter,
         check_canceled: Optional[Callable[[], None]],
     ) -> Tuple[int, float]:
@@ -85,7 +86,7 @@ def _train_model(
         logger.info("Training NMT model")
         with (
             progress_reporter.start_next_phase() as phase_progress,
-            self._nmt_model_factory.create_model_trainer(parallel_corpus) as model_trainer,
+            self._nmt_model_factory.create_model_trainer(parallel_corpus, parallel_terms_corpus) as model_trainer,
         ):
             model_trainer.train(progress=phase_progress, check_canceled=check_canceled)
             model_trainer.save()

machine/jobs/nmt_engine_build_job.py

@@ -22,7 +22,7 @@ def create_source_tokenizer_trainer(self, corpus: TextCorpus) -> Trainer: ...
     def create_target_tokenizer_trainer(self, corpus: TextCorpus) -> Trainer: ...
 
     @abstractmethod
-    def create_model_trainer(self, corpus: ParallelTextCorpus) -> Trainer: ...
+    def create_model_trainer(self, corpus: ParallelTextCorpus, terms_corpus: ParallelTextCorpus) -> Trainer: ...
 
     @abstractmethod
     def create_engine(self) -> TranslationEngine: ...

machine/jobs/nmt_model_factory.py

@@ -45,6 +45,7 @@ def _train_model(
         source_corpus: TextCorpus,
         target_corpus: TextCorpus,
         parallel_corpus: ParallelTextCorpus,
+        parallel_terms_corpus: ParallelTextCorpus,
         progress_reporter: PhasedProgressReporter,
         check_canceled: Optional[Callable[[], None]],
     ) -> Tuple[int, float]:

machine/jobs/smt_engine_build_job.py

@@ -28,14 +28,20 @@ def run(
         target_corpus = self._translation_file_service.create_target_corpus()
         parallel_corpus: ParallelTextCorpus = source_corpus.align_rows(target_corpus)
 
-        parallel_corpus_size = parallel_corpus.count(include_empty=False)
+        source_terms_corpus = self._translation_file_service.create_source_terms_corpus()
+        target_terms_corpus = self._translation_file_service.create_target_terms_corpus()
+        parallel_terms_corpus: ParallelTextCorpus = source_terms_corpus.align_rows(target_terms_corpus)
+
+        parallel_corpus_size = parallel_corpus.count(include_empty=False) + parallel_terms_corpus.count(
+            include_empty=False
+        )
         progress_reporter = self._get_progress_reporter(progress, parallel_corpus_size)
 
         if parallel_corpus_size == 0:
             train_corpus_size, confidence = self._respond_to_no_training_corpus()
         else:
             train_corpus_size, confidence = self._train_model(
-                source_corpus, target_corpus, parallel_corpus, progress_reporter, check_canceled
+                source_corpus, target_corpus, parallel_corpus, parallel_terms_corpus, progress_reporter, check_canceled
             )
 
         if check_canceled is not None:
@@ -63,6 +69,7 @@ def _train_model(
         source_corpus: TextCorpus,
         target_corpus: TextCorpus,
         parallel_corpus: ParallelTextCorpus,
+        parallel_terms_corpus: ParallelTextCorpus,
         progress_reporter: PhasedProgressReporter,
         check_canceled: Optional[Callable[[], None]],
     ) -> Tuple[int, float]: ...

machine/jobs/translation_engine_build_job.py

@@ -1,6 +1,6 @@
 from contextlib import contextmanager
 from pathlib import Path
-from typing import Any, Generator, Iterator, List, Optional, TypedDict, Union
+from typing import Any, Generator, Iterator, List, Optional, TypedDict
 
 import json_stream
 
@@ -26,46 +26,61 @@ def __init__(
         self,
         type: SharedFileServiceType,
         config: Any,
-        source_filenames: Optional[Union[str, List[str]]] = None,
-        target_filenames: Optional[Union[str, List[str]]] = None,
+        source_filename: Optional[str] = "train.src.txt",
+        target_filename: Optional[str] = "train.trg.txt",
+        source_terms_filename: Optional[str] = "train.key-terms.src.txt",
+        target_terms_filename: Optional[str] = "train.key-terms.trg.txt",
         source_pretranslation_filename: str = "pretranslate.src.json",
         target_pretranslation_filename: str = "pretranslate.trg.json",
     ) -> None:
 
-        if source_filenames is None:
-            source_filenames = ["train.src.txt", "train.key-terms.src.txt"]
-        if target_filenames is None:
-            target_filenames = ["train.trg.txt", "train.key-terms.trg.txt"]
-
-        self._source_filenames = [source_filenames] if isinstance(source_filenames, str) else list(source_filenames)
-        self._target_filenames = [target_filenames] if isinstance(target_filenames, str) else list(target_filenames)
+        self._source_filename = source_filename
+        self._target_filename = target_filename
+        self._source_terms_filename = source_terms_filename
+        self._target_terms_filename = target_terms_filename
         self._source_pretranslation_filename = source_pretranslation_filename
         self._target_pretranslation_filename = target_pretranslation_filename
 
         self.shared_file_service: SharedFileServiceBase = get_shared_file_service(type, config)
 
     def create_source_corpus(self) -> TextCorpus:
         return TextFileTextCorpus(
-            self.shared_file_service.download_file(f"{self.shared_file_service.build_path}/{source_filename}")
-            for source_filename in self._source_filenames
+            self.shared_file_service.download_file(f"{self.shared_file_service.build_path}/{self._source_filename}")
         )
 
     def create_target_corpus(self) -> TextCorpus:
         return TextFileTextCorpus(
-            self.shared_file_service.download_file(f"{self.shared_file_service.build_path}/{target_filename}")
-            for target_filename in self._target_filenames
+            self.shared_file_service.download_file(f"{self.shared_file_service.build_path}/{self._target_filename}")
         )
 
     def exists_source_corpus(self) -> bool:
-        return all(
-            self.shared_file_service._exists_file(f"{self.shared_file_service.build_path}/{source_filename}")
-            for source_filename in self._source_filenames
-        )
+        return self.shared_file_service._exists_file(f"{self.shared_file_service.build_path}/{self._source_filename}")
 
     def exists_target_corpus(self) -> bool:
-        return all(
-            self.shared_file_service._exists_file(f"{self.shared_file_service.build_path}/{target_filename}")
-            for target_filename in self._target_filenames
+        return self.shared_file_service._exists_file(f"{self.shared_file_service.build_path}/{self._target_filename}")
+
+    def create_source_terms_corpus(self) -> TextCorpus:
+        return TextFileTextCorpus(
+            self.shared_file_service.download_file(
+                f"{self.shared_file_service.build_path}/{self._source_terms_filename}"
+            )
+        )
+
+    def create_target_terms_corpus(self) -> TextCorpus:
+        return TextFileTextCorpus(
+            self.shared_file_service.download_file(
+                f"{self.shared_file_service.build_path}/{self._target_terms_filename}"
+            )
+        )
+
+    def exists_source_terms_corpus(self) -> bool:
+        return self.shared_file_service._exists_file(
+            f"{self.shared_file_service.build_path}/{self._source_terms_filename}"
+        )
+
+    def exists_target_terms_corpus(self) -> bool:
+        return self.shared_file_service._exists_file(
+            f"{self.shared_file_service.build_path}/{self._target_terms_filename}"
         )
 
     def get_source_pretranslations(self) -> ContextManagedGenerator[PretranslationInfo, None, None]:

machine/jobs/translation_file_service.py

@@ -9,6 +9,7 @@
 import torch  # pyright: ignore[reportMissingImports]
 from accelerate import Accelerator  # pyright: ignore[reportMissingImports]
 from accelerate.utils.memory import should_reduce_batch_size  # pyright: ignore[reportMissingImports]
+from datasets import concatenate_datasets
 from datasets.arrow_dataset import Dataset
 from sacremoses import MosesPunctNormalizer
 from torch import Tensor  # pyright: ignore[reportMissingImports]
@@ -36,9 +37,11 @@
     PreTrainedTokenizerFast,
     Seq2SeqTrainer,
     Seq2SeqTrainingArguments,
+    TensorType,
     TrainerCallback,
     set_seed,
 )
+from transformers.tokenization_utils import BatchEncoding
 from transformers.trainer_callback import TrainerControl, TrainerState
 from transformers.trainer_utils import get_last_checkpoint
 from transformers.training_args import TrainingArguments
@@ -88,6 +91,7 @@ def __init__(
         model: Union[PreTrainedModel, str],
         training_args: Seq2SeqTrainingArguments,
         corpus: Union[ParallelTextCorpus, Dataset],
+        terms_corpus: Optional[Union[ParallelTextCorpus, Dataset]] = None,
         src_lang: Optional[str] = None,
         tgt_lang: Optional[str] = None,
         max_src_length: Optional[int] = None,
@@ -98,6 +102,7 @@ def __init__(
         self._model = model
         self._training_args = training_args
         self._corpus = corpus
+        self._terms_corpus = terms_corpus
         self._src_lang = src_lang
         self._tgt_lang = tgt_lang
         self._trainer: Optional[Seq2SeqTrainer] = None
@@ -170,6 +175,13 @@ def train(
         else:
             train_dataset = self._corpus.filter_nonempty().to_hf_dataset(src_lang, tgt_lang)
 
+        train_terms_dataset = None
+        if self._terms_corpus is not None:
+            if isinstance(self._terms_corpus, Dataset):
+                train_terms_dataset = self._terms_corpus
+            else:
+                train_terms_dataset = self._terms_corpus.filter_nonempty().to_hf_dataset(src_lang, tgt_lang)
+
         def find_missing_characters(tokenizer: Any, train_dataset: Dataset, lang_codes: List[str]) -> List[str]:
             vocab = tokenizer.get_vocab().keys()
             charset = set()
@@ -222,7 +234,15 @@ def add_tokens(tokenizer: Any, missing_tokens: List[str]) -> Any:
                     lang_codes.append(src_lang)
                 if self._add_unk_tgt_tokens:
                     lang_codes.append(tgt_lang)
-                missing_tokens = find_missing_characters(tokenizer, train_dataset, lang_codes)
+                missing_tokens = find_missing_characters(
+                    tokenizer,
+                    (
+                        concatenate_datasets([train_dataset, train_terms_dataset])
+                        if train_terms_dataset is not None
+                        else train_dataset
+                    ),
+                    lang_codes,
+                )
                 if missing_tokens:
                     tokenizer = add_tokens(tokenizer, missing_tokens)
 
@@ -291,6 +311,22 @@ def add_tokens(tokenizer: Any, missing_tokens: List[str]) -> Any:
                 "memory"
             )
 
+        def batch_prepare_for_model(
+            tokenizer: Union[PreTrainedTokenizer, PreTrainedTokenizerFast],
+            batch_tokens: List[List[str]],
+            return_tensors: Optional[Union[str, TensorType]] = None,
+        ) -> BatchEncoding:
+            batch_outputs: Dict[str, Any] = {}
+            for tokens in batch_tokens:
+                ids = cast(List[int], tokenizer.convert_tokens_to_ids(tokens))
+                outputs = tokenizer.prepare_for_model(ids, add_special_tokens=False)
+
+                for key, value in outputs.items():
+                    if key not in batch_outputs:
+                        batch_outputs[key] = []
+                    batch_outputs[key].append(value)
+            return BatchEncoding(batch_outputs, tensor_type=return_tensors)
+
         def preprocess_function(examples):
             if isinstance(tokenizer, (NllbTokenizer, NllbTokenizerFast)):
                 inputs = [self._mpn.normalize(prefix + ex[src_lang]) for ex in examples["translation"]]
@@ -306,6 +342,42 @@ def preprocess_function(examples):
             model_inputs["labels"] = labels["input_ids"]
             return model_inputs
 
+        def preprocess_terms_function(examples):
+            if isinstance(tokenizer, (NllbTokenizer, NllbTokenizerFast)):
+                inputs = [self._mpn.normalize(ex[src_lang]) for ex in examples["translation"]]
+                targets = [self._mpn.normalize(ex[tgt_lang]) for ex in examples["translation"]]
+            else:
+                inputs = [ex[src_lang] for ex in examples["translation"]]
+                targets = [ex[tgt_lang] for ex in examples["translation"]]
+
+            src_term_tokens = tokenizer(
+                [prefix + i for i in inputs], max_length=max_src_length, truncation=True
+            ).tokens()
+            trg_term_tokens = tokenizer(text_target=targets, max_length=max_tgt_length, truncation=True).tokens()
+
+            src_term_partial_word_tokens = tokenizer(
+                [prefix + "\ufffc" + i for i in inputs], max_length=max_src_length + 2, truncation=True
+            ).tokens()
+            src_term_partial_word_tokens.remove("▁")
+            src_term_partial_word_tokens.remove("\ufffc")
+
+            trg_term_partial_word_tokens = tokenizer(
+                text_target=["\ufffc" + t for t in targets], max_length=max_tgt_length + 2, truncation=True
+            ).tokens()
+            trg_term_partial_word_tokens.remove("▁")
+            trg_term_partial_word_tokens.remove("\ufffc")
+
+            model_inputs = batch_prepare_for_model(
+                tokenizer, [[ex.strip() for ex in src_term_tokens + src_term_partial_word_tokens]]
+            )
+            # Tokenize targets with the `text_target` keyword argument
+            labels = batch_prepare_for_model(
+                tokenizer, [[ex.strip() for ex in trg_term_tokens + trg_term_partial_word_tokens]]
+            )
+
+            model_inputs["labels"] = labels["input_ids"]
+            return model_inputs
+
         logger.info("Run tokenizer")
         train_dataset = train_dataset.map(
             preprocess_function,
@@ -315,6 +387,22 @@ def preprocess_function(examples):
             desc="Running tokenizer on train dataset",
         )
 
+        if train_terms_dataset is not None:
+            if not isinstance(tokenizer, PreTrainedTokenizerFast):
+                logger.warning(
+                    f"Adding key terms as partial words is not possible when using the non-fast tokenizer '{type(tokenizer)}'."
+                )
+            train_terms_dataset = train_terms_dataset.map(
+                preprocess_terms_function if isinstance(tokenizer, PreTrainedTokenizerFast) else preprocess_function,
+                batched=True,
+                remove_columns=train_terms_dataset.column_names,
+                load_from_cache_file=True,
+                desc="Running tokenizer on train terms dataset",
+            )
+
+            # combine terms and non-terms datasets
+            train_dataset = concatenate_datasets([train_dataset, train_terms_dataset])
+
         data_collator = DataCollatorForSeq2Seq(
             tokenizer,
             model=model,