Just one call to batch.doc.get and reusing the results should be slig…

…htly faster
stanfordnlp · Jul 19, 2023 · 9f05e7b · 9f05e7b
1 parent d21a95c
commit 9f05e7b
Showing 1 changed file with 6 additions and 2 deletions.
diff --git a/stanza/pipeline/lemma_processor.py b/stanza/pipeline/lemma_processor.py
@@ -8,6 +8,8 @@
 from stanza.pipeline._constants import *
 from stanza.pipeline.processor import UDProcessor, register_processor
 
+WORD_TAGS = [doc.TEXT, doc.UPOS]
+
 @register_processor(name=LEMMA)
 class LemmaProcessor(UDProcessor):
 
@@ -74,7 +76,9 @@ def process(self, document):
                     edits += es
 
             if self.config.get('ensemble_dict', False):
-                preds = self.trainer.postprocess([x for x, y in zip(batch.doc.get([doc.TEXT]), skip) if not y], preds, edits=edits)
+                word_tags = batch.doc.get(WORD_TAGS)
+                words = [x[0] for x in word_tags]
+                preds = self.trainer.postprocess([x for x, y in zip(words, skip) if not y], preds, edits=edits)
                 # expand seq2seq predictions to the same size as all words
                 i = 0
                 preds1 = []
@@ -84,7 +88,7 @@ def process(self, document):
                     else:
                         preds1.append(preds[i])
                         i += 1
-                preds = self.trainer.ensemble(batch.doc.get([doc.TEXT, doc.UPOS]), preds1)
+                preds = self.trainer.ensemble(word_tags, preds1)
             else:
                 preds = self.trainer.postprocess(batch.doc.get([doc.TEXT]), preds, edits=edits)