Raises AssertionError on Incomplete Vocabulary

fixes issue #621 added a new function CheckVocab, to check for presence of a word in vocabulary
tensorflow · Dec 8, 2016 · 27a178a · 27a178a
1 parent 821d3da
commit 27a178a
Show file tree

Hide file tree

Showing 2 changed files with 10 additions and 4 deletions.
diff --git a/textsum/data.py b/textsum/data.py
@@ -56,6 +56,12 @@ def __init__(self, vocab_file, max_size):
         if self._count > max_size:
           raise ValueError('Too many words: >%d.' % max_size)
 
+  def CheckVocab(self, word):
+    if word not in self._word_to_id:
+      return None
+    return self._word_to_id[word]
+
+
   def WordToId(self, word):
     if word not in self._word_to_id:
       return self._word_to_id[UNKNOWN_TOKEN]

diff --git a/textsum/seq2seq_attention.py b/textsum/seq2seq_attention.py
@@ -160,10 +160,10 @@ def _Eval(model, data_batcher, vocab=None):
 def main(unused_argv):
   vocab = data.Vocab(FLAGS.vocab_path, 1000000)
   # Check for presence of required special tokens.
-  assert vocab.WordToId(data.PAD_TOKEN) > 0
-  assert vocab.WordToId(data.UNKNOWN_TOKEN) >= 0
-  assert vocab.WordToId(data.SENTENCE_START) > 0
-  assert vocab.WordToId(data.SENTENCE_END) > 0
+  assert vocab.CheckVocab(data.PAD_TOKEN) > 0
+  assert vocab.CheckVocab(data.UNKNOWN_TOKEN) >= 0
+  assert vocab.CheckVocab(data.SENTENCE_START) > 0
+  assert vocab.CheckVocab(data.SENTENCE_END) > 0
 
   batch_size = 4
   if FLAGS.mode == 'decode':