Loading text tutorial: fixed OOV handling

tensorflow · Aug 27, 2023 · f2ea5cf · f2ea5cf
1 parent e4ab8e0
commit f2ea5cf
Showing 1 changed file with 5 additions and 2 deletions.
diff --git a/site/en/tutorials/load_data/text.ipynb b/site/en/tutorials/load_data/text.ipynb
@@ -1140,8 +1140,9 @@
       },
       "outputs": [],
       "source": [
-        "keys = vocab\n",
-        "values = range(2, len(vocab) + 2)  # Reserve `0` for padding, `1` for OOV tokens.\n",
+        "# Reserve `0` for padding, `1` for OOV tokens.\n",
+        "keys = ['', '[UNK]'] + vocab\n",
+        "values = range(len(keys))\n",
         "\n",
         "init = tf.lookup.KeyValueTensorInitializer(\n",
         "    keys, values, key_dtype=tf.string, value_dtype=tf.int64)\n",
@@ -1171,6 +1172,8 @@
         "  standardized = tf_text.case_fold_utf8(text)\n",
         "  tokenized = tokenizer.tokenize(standardized)\n",
         "  vectorized = vocab_table.lookup(tokenized)\n",
+        "  # StaticVocabularyTable returns the OOV token as vocab_size + 2. We overwrite it to be 1.\n",
+        "  vectorized = tf.where(vectorized == len(keys), tf.constant(1, dtype=tf.int64), vectorized)\n",
         "  return vectorized, label"
       ]
     },