Prevent model from being loaded multiple times if translate method is called multiple times (#801)

* cache model during translate step

* add clear_cache; add dataclass; rename variables to be more precise

* add clear_cache to NMTModel interface

* fix typo from rebase

Author: mshannon-sil

silnlp/nmt/config.py

@@ -91,6 +91,9 @@ def translate(
     @abstractmethod
     def get_checkpoint_path(self, ckpt: Union[CheckpointType, str, int]) -> Tuple[Path, int]: ...
 
+    @abstractmethod
+    def clear_cache(self) -> None: ...
+
 
 class Config(ABC):
     def __init__(self, exp_dir: Path, config: dict) -> None:

silnlp/nmt/hugging_face_config.py

@@ -6,6 +6,7 @@
 import shutil
 from contextlib import ExitStack
 from copy import deepcopy
+from dataclasses import dataclass
 from enum import Enum
 from itertools import repeat
 from math import exp, prod
@@ -832,6 +833,20 @@ def get_token_scores(self) -> List[float]:
     def get_sequence_score(self) -> List[float]:
         return [output["sequence_score"] for output in self.outputs]
 
+@dataclass
+class InferenceModelParams:
+    checkpoint: Union[CheckpointType, str, int]
+    src_lang: str
+    trg_lang: str
+
+    def __post_init__(self):
+        if not isinstance(self.checkpoint, (CheckpointType, str, int)):
+            raise ValueError("checkpoint must be a CheckpointType, string, or integer")
+        if not isinstance(self.src_lang, str):
+            raise ValueError("src_lang must be a string")
+        if not isinstance(self.trg_lang, str):
+            raise ValueError("trg_lang must be a string")
+
 
 class HuggingFaceNMTModel(NMTModel):
     def __init__(self, config: HuggingFaceConfig, mixed_precision: bool, num_devices: int) -> None:
@@ -841,6 +856,8 @@ def __init__(self, config: HuggingFaceConfig, mixed_precision: bool, num_devices
         self._dictionary: Optional[Dict[VerseRef, Set[str]]] = None
         self._is_t5 = self._config.model_prefix in SUPPORTED_T5_MODELS
         self._num_devices = num_devices
+        self._cached_inference_model: Optional[PreTrainedModel] = None
+        self._inference_model_params: Optional[InferenceModelParams] = None
 
     def train(self) -> None:
         training_args = self._create_training_arguments()
@@ -1218,8 +1235,13 @@ def translate(
     ) -> Iterable[TranslationGroup]:
         src_lang = self._config.data["lang_codes"].get(src_iso, src_iso)
         trg_lang = self._config.data["lang_codes"].get(trg_iso, trg_iso)
+        inference_model_params = InferenceModelParams(ckpt, src_lang, trg_lang)
         tokenizer = self._config.get_tokenizer()
-        model = self._create_inference_model(ckpt, tokenizer, src_lang, trg_lang)
+        if self._inference_model_params == inference_model_params and self._cached_inference_model is not None:
+            model = self._cached_inference_model
+        else:
+            model = self._cached_inference_model = self._create_inference_model(ckpt, tokenizer, src_lang, trg_lang)
+            self._inference_model_params = inference_model_params
         if model.config.max_length is not None and model.config.max_length < 512:
             model.config.max_length = 512
 
@@ -1304,6 +1326,10 @@ def get_checkpoint_path(self, ckpt: Union[CheckpointType, str, int]) -> Tuple[Pa
             raise ValueError(f"Unsupported checkpoint type: {ckpt}.")
         return ckpt_path, step
 
+    def clear_cache(self) -> None:
+        self._cached_inference_model = None
+        self._inference_model_params = None
+
     def _create_training_arguments(self) -> Seq2SeqTrainingArguments:
         parser = HfArgumentParser(Seq2SeqTrainingArguments)
         args: dict = {}