type hint

asyml · Apr 6, 2021 · 35bd1dc · 35bd1dc
1 parent a1e7c74
commit 35bd1dc
Showing 1 changed file with 4 additions and 4 deletions.
diff --git a/forte_wrapper/allennlp/allennlp_processors.py b/forte_wrapper/allennlp/allennlp_processors.py
@@ -140,16 +140,16 @@ def _process(self, input_pack: DataPack):
 
         batch_size = self.configs['infer_batch_size']
         if batch_size <= 0:
-            batches = [input_pack.get(Sentence)]
+            batches = iter([input_pack.get(Sentence)])
         else:
-            batches = more_itertools.chunked(input_pack.get(Sentence),
-                                             batch_size, strict=False)
+            batches = more_itertools.ichunked(input_pack.get(Sentence),
+                                              batch_size)
         for sentences in batches:
             inputs = [{"sentence": s.text} for s in sentences]
             results = {k: p.predict_batch_json(inputs)
                        for k, p in self.predictor.items()}
             for i, sentence in enumerate(sentences):
-                result = {}
+                result: Dict[str, List[str]] = {}
                 for key in self.predictor.keys():
                     if key == 'srl':
                         result.update(parse_allennlp_srl_results(