✅ Adding or updating tests.

BrikerMan · Jul 5, 2020 · ce6860d · ce6860d
1 parent 9ee1e79
commit ce6860d
Show file tree

Hide file tree

Showing 4 changed files with 4 additions and 153 deletions.
diff --git a/kashgari/processors/sequence_processor.py b/kashgari/processors/sequence_processor.py
@@ -95,9 +95,9 @@ def build_vocab_generator(self,
             self.vocab2idx = vocab2idx
             self.idx2vocab = dict([(v, k) for k, v in self.vocab2idx.items()])
 
-            logger.debug(f"--- Build vocab dict finished, Total: {len(self.vocab2idx)} Top-10 ---")
-            for token, index in list(self.vocab2idx.items())[:10]:
-                logger.debug(f"Token: {token:8s} -> {index}")
+            top_k_vocab = [k for (k, v) in list(self.vocab2idx.items())[:10]]
+            logger.debug(f"--- Build vocab dict finished, Total: {len(self.vocab2idx)} ---")
+            logger.debug(f"Top-10: {top_k_vocab}")
 
     def transform(self,
                   samples: TextSamplesVar,

diff --git a/kashgari/tasks/labeling/bi_gru_crf_model.py b/kashgari/tasks/labeling/bi_gru_crf_model.py
diff --git a/kashgari/tasks/labeling/bi_lstm_crf_model.py b/kashgari/tasks/labeling/bi_lstm_crf_model.py
diff --git a/tests/test_generator.py b/tests/test_generator.py
@@ -40,8 +40,7 @@ def test_batch_generator(self):
                                       max_position=100,
                                       batch_size=12)
 
-        assert len(list(batch_dataset1.take())) == len(batch_dataset1)
-        duplicate_len = len(batch_dataset1) * 2
+        duplicate_len = len(batch_dataset1)
         assert len(list(batch_dataset1.take(duplicate_len))) == duplicate_len
         assert len(list(batch_dataset1.take(1))) == 1