load vecs and lm head only once

app.py

@@ -3,6 +3,10 @@ import pandas as pd
 import transformers
 import gradio as gr
 
+tokenizer = transformers.AutoTokenizer.from_pretrained('gpt2')
+vecs = torch.load("senses/all_vecs_mtx.pt")
+lm_head = torch.load("senses/lm_head.pt")
+
 
 def visualize_word(word, count=10, remove_space=False):
 
@@ -10,11 +14,6 @@ def visualize_word(word, count=10, remove_space=False):
         word = ' ' + word
     print(f"Looking up word '{word}'...")
 
-    # very dumb to have to load the tokenizer every time, trying to figure out how to pass a non-interface element into the function in gradio
-    tokenizer = transformers.AutoTokenizer.from_pretrained('gpt2')
-    vecs = torch.load("senses/all_vecs_mtx.pt")
-    lm_head = torch.load("senses/lm_head.pt")
-
     token_ids = tokenizer(word)['input_ids']
     tokens = [tokenizer.decode(token_id) for token_id in token_ids]
     tokens = ", ".join(tokens) # display tokenization for user
@@ -45,12 +44,14 @@ def visualize_word(word, count=10, remove_space=False):
                         columns=list(data.keys()))
         for prop, word_list in data.items():
             for i, word_pair in enumerate(word_list):
+                cell_value = "space ({:.2f})".format(word_pair[1])
                 cell_value = "{} ({:.2f})".format(word_pair[0], word_pair[1])
                 df.at[i, prop] = cell_value
         return df
     
     pos_df = create_dataframe(pos_word_lists, sense_names, count)
     neg_df = create_dataframe(neg_word_lists, sense_names, count)
+    print(pos_df)
 
     return pos_df, neg_df, tokens