Fix the attachment of the lemma classifier to the pipeline

stanza/models/lemma/trainer.py

@@ -292,7 +292,7 @@ def save(self, filename, skip_modules=True):
         torch.save(params, filename, _use_new_zipfile_serialization=False)
         logger.info("Model saved to {}".format(filename))
 
-    def load(self, filename, args, foundation_cache):
+    def load(self, filename, args, foundation_cache, lemma_classifier_args=None):
         try:
             checkpoint = torch.load(filename, lambda storage, loc: storage, weights_only=True)
         except BaseException:
@@ -313,4 +313,4 @@ def load(self, filename, args, foundation_cache):
         self.vocab = MultiVocab.load_state_dict(checkpoint['vocab'])
         self.contextual_lemmatizers = []
         for contextual in checkpoint.get('contextual', []):
-            self.contextual_lemmatizers.append(LemmaClassifier.from_checkpoint(contextual))
+            self.contextual_lemmatizers.append(LemmaClassifier.from_checkpoint(contextual, args=lemma_classifier_args))

stanza/pipeline/lemma_processor.py

@@ -48,11 +48,14 @@ def _set_up_model(self, config, pipeline, device):
             # since a long running program will remember everything
             # (unless we go back and make it smarter)
             # we make this an option, not the default
+            # TODO: need to update the cache to skip the contextual lemmatizer
             self.store_results = config.get('store_results', False)
             self._use_identity = False
             args = {'charlm_forward_file': config.get('forward_charlm_path', None),
                     'charlm_backward_file': config.get('backward_charlm_path', None)}
-            self._trainer = Trainer(args=args, model_file=config['model_path'], device=device, foundation_cache=pipeline.foundation_cache)
+            lemma_classifier_args = dict(args)
+            lemma_classifier_args['wordvec_pretrain_file'] = config.get('pretrain_path', None)
+            self._trainer = Trainer(args=args, model_file=config['model_path'], device=device, foundation_cache=pipeline.foundation_cache, lemma_classifier_args=lemma_classifier_args)
 
     def _set_up_requires(self):
         self._pretagged = self._config.get('pretagged', None)