type hint

asyml · Apr 6, 2021 · 3eae958 · 3eae958
1 parent 5a89282
commit 3eae958
Showing 1 changed file with 4 additions and 3 deletions.
diff --git a/forte_wrapper/allennlp/allennlp_processors.py b/forte_wrapper/allennlp/allennlp_processors.py
@@ -15,7 +15,7 @@
 import itertools
 import logging
 import more_itertools
-from typing import Any, Dict, List
+from typing import Any, Dict, Iterable, Iterator, List
 
 from allennlp.predictors import Predictor
 
@@ -139,11 +139,12 @@ def _process(self, input_pack: DataPack):
         self._process_existing_entries(input_pack)
 
         batch_size: int = self.configs['infer_batch_size']
+        batches: Iterator[Iterable[Sentence]]
         if batch_size <= 0:
             batches = iter([input_pack.get(Sentence)])
         else:
-            batches = more_itertools.ichunked(input_pack.get(Sentence),
-                                              batch_size)
+            batches = more_itertools.chunked(
+                input_pack.get(Sentence), batch_size)
         for sentences in batches:
             inputs: List[Dict[str, str]] = [{"sentence": s.text}
                                             for s in sentences]