Revert "Removing the loaded model cache in MultiSourceSeq2Seq."

This reverts commit 8f6f261.

Author: pengyunie

roosterize/ml/naming/MultiSourceSeq2Seq.py

@@ -237,6 +237,8 @@ def __init__(self, model_dir: Path, model_spec: ModelSpec):
 
         # Cache for processing data
         self.data_cache: dict = dict()
+        # Cache for loaded model during translation
+        self.loaded_model_cache = None
         return
 
     def get_input(
@@ -848,7 +850,15 @@ def eval_impl(
             # translate.main
             ArgumentParser.validate_translate_opts(opt)
 
-            translator = MultiSourceTranslator.build_translator(self.config.get_src_types(), opt, report_score=False)
+            # Cached model loading
+            if self.loaded_model_cache is None:
+                self.loaded_model_cache = MultiSourceTranslator.load_model(self.config.get_src_types(), opt)
+            translator = MultiSourceTranslator.build_translator(
+                self.config.get_src_types(),
+                opt,
+                loaded_model=self.loaded_model_cache,
+                report_score=False,
+            )
 
             has_target = True
             raw_data_keys = [f"src.{src_type}" for src_type in self.config.get_src_types()] + (["tgt"] if has_target else [])

roosterize/ml/onmt/MultiSourceTranslator.py

@@ -20,15 +20,27 @@
 from seutil import LoggingUtils
 
 
+class LoadedModel(NamedTuple):
+    fields: any
+    model: any
+    model_opt: any
+
+
 class MultiSourceTranslator(CustomTranslator):
 
     logger = LoggingUtils.get_logger(__name__)
 
+    @classmethod
+    def load_model(cls, src_types, opt) -> LoadedModel:
+        fields, model, model_opt = MultiSourceModelBuilder.load_test_model(src_types, opt)
+        return LoadedModel(fields, model, model_opt)
+
     @classmethod
     def build_translator(
             cls,
             src_types,
             opt,
+            loaded_model: LoadedModel = None,
             report_score=True,
             logger=None,
             out_file=None,
@@ -38,9 +50,10 @@ def build_translator(
 
         assert len(opt.models) == 1, "ensemble model is not supported"
 
-        # load_test_model = onmt.decoders.ensemble.load_test_model \
-        #     if len(opt.models) > 1 else onmt.model_builder.load_test_model
-        fields, model, model_opt = MultiSourceModelBuilder.load_test_model(src_types, opt)
+        if loaded_model is None:
+            fields, model, model_opt = MultiSourceModelBuilder.load_test_model(src_types, opt)
+        else:
+            fields, model, model_opt = loaded_model
 
         scorer = onmt.translate.GNMTGlobalScorer.from_opt(opt)