fix: mismatch between building vocabs

vejvarm · Nov 30, 2023 · 4d4b171 · 4d4b171
1 parent 0e6f387
commit 4d4b171
Showing 1 changed file with 2 additions and 2 deletions.
diff --git a/dataset.py b/dataset.py
@@ -220,7 +220,7 @@ def _build_vocabs_streaming(self):
                                              specials=[O, PAD_TOKEN],
                                              vocab_cache=self.vocab_cache.joinpath("ner_vocab.pkl"))
         vocabs[COREF] = self._counter_to_vocab(self.counters[COREF],
-                                               specials=['0', PAD_TOKEN],
+                                               specials=[NA_TOKEN, PAD_TOKEN],
                                                vocab_cache=self.vocab_cache.joinpath("coref_vocab.pkl"))
         vocabs[PREDICATE_POINTER] = self._counter_to_vocab(self.counters[PREDICATE_POINTER],
                                                            specials=[NA_TOKEN, PAD_TOKEN],
@@ -229,7 +229,7 @@ def _build_vocabs_streaming(self):
                                                       specials=[NA_TOKEN, PAD_TOKEN],
                                                       vocab_cache=self.vocab_cache.joinpath("type_vocab.pkl"))
         vocabs[ENTITY] = self._counter_to_vocab(self.counters[ENTITY],
-                                                specials=[PAD_TOKEN, NA_TOKEN],
+                                                specials=[NA_TOKEN, PAD_TOKEN],
                                                 vocab_cache=self.vocab_cache.joinpath("ent_vocab.pkl"))
 
         return vocabs