fix index mapping for w2v example

garretthoffman · Apr 16, 2019 · cf7c3e7 · cf7c3e7
1 parent ed3ca12
commit cf7c3e7
Show file tree

Hide file tree

Showing 2 changed files with 15 additions and 1 deletion.
diff --git a/00_Word2Vec_Skipgram_TF.ipynb b/00_Word2Vec_Skipgram_TF.ipynb
@@ -167,7 +167,7 @@
    "outputs": [],
    "source": [
     "full_lexicon = \" \".join(messages).split()\n",
-    "vocab_to_int, int_to_vocab = utl.create_lookup_tables(full_lexicon)"
+    "vocab_to_int, int_to_vocab = utl.create_lookup_tables_w2v(full_lexicon)"
    ]
   },
   {

diff --git a/utils.py b/utils.py
@@ -42,6 +42,20 @@ def preprocess_ST_message(text):
 
     return words
 
+def create_lookup_tables_w2v(words):
+    """
+    Create lookup tables for vocabulary
+    :param words: Input list of words
+    :return: A tuple of dicts.  The first dict maps a vocab word to and integeter
+             The second maps an integer back to to the vocab word
+    """
+    word_counts = Counter(words)
+    sorted_vocab = sorted(word_counts, key=word_counts.get, reverse=True)
+    int_to_vocab = {ii: word for ii, word in enumerate(sorted_vocab)}
+    vocab_to_int = {word: ii for ii, word in int_to_vocab.items()}
+
+    return vocab_to_int, int_to_vocab
+
 def create_lookup_tables(words):
     """
     Create lookup tables for vocabulary