Initial Commit

app.py

@@ -0,0 +1,110 @@
+import torch
+import pandas as pd
+import transformers
+import gradio as gr
+
+# def visualize_word(word, tokenizer, vecs, lm_head, count=5, contents=None):
+def visualize_word(word, count=10, remove_space=False):
+
+    if not remove_space:
+        word = ' ' + word
+    print(f"Looking up word ['{word}']")
+
+    # seems very dumb to have to load the tokenizer every time, but I don't know 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")
+    print("lm_head.shape = ", lm_head.shape)
+
+    token_ids = tokenizer(word)['input_ids']
+    tokens = [tokenizer.decode(token_id) for token_id in token_ids]
+    tokens = ", ".join(tokens)
+    # look up sense vectors only for the first token
+    contents = vecs[token_ids[0]] # torch.Size([16, 768])
+
+    sense_names = []
+    pos_sense_word_lists = []
+    neg_sense_word_lists = []
+
+    for i in range(contents.shape[0]):
+        logits = contents[i,:] @ lm_head.t() # (vocab,)    [768] @ [768, 50257] -> [50257]
+        sorted_logits, sorted_indices = torch.sort(logits, descending=True)
+        sense_names.append('sense {}'.format(i+1))
+
+        # currently a lot of repetition
+        pos_sorted_words = [tokenizer.decode(sorted_indices[j]) for j in range(count)]
+        pos_sorted_logits = [sorted_logits[j].item() for j in range(count)]
+        pos_word_list = list(zip(pos_sorted_words, pos_sorted_logits))
+        pos_sense_word_lists.append(pos_word_list)
+
+        neg_sorted_words = [tokenizer.decode(sorted_indices[-j-1]) for j in range(count)]
+        neg_sorted_logits = [sorted_logits[-j-1].item() for j in range(count)]
+        neg_word_list = list(zip(neg_sorted_words, neg_sorted_logits))
+        neg_sense_word_lists.append(neg_word_list)
+
+    pos_data = dict(zip(sense_names, pos_sense_word_lists))
+    pos_df = pd.DataFrame(index=[i for i in range(count)],
+                  columns=list(pos_data.keys()))
+    for prop, word_list in pos_data.items():
+        for i, word_pair in enumerate(word_list):
+            cell_value = "{} ({:.2f})".format(word_pair[0], word_pair[1])
+            pos_df.at[i, prop] = cell_value
+    
+    neg_data = dict(zip(sense_names, neg_sense_word_lists))
+    neg_df = pd.DataFrame(index=[i for i in range(count)],
+                    columns=list(neg_data.keys()))
+    for prop, word_list in neg_data.items():
+        for i, word_pair in enumerate(word_list):
+            cell_value = "{} ({:.2f})".format(word_pair[0], word_pair[1])
+            neg_df.at[i, prop] = cell_value
+
+    return pos_df, neg_df, tokens
+
+# argp = argparse.ArgumentParser()
+# argp.add_argument('vecs_path')
+# argp.add_argument('lm_head_path')
+# args = argp.parse_args()
+
+# Load tokenizer and parameters
+# tokenizer = transformers.AutoTokenizer.from_pretrained('gpt2')
+# vecs = torch.load(args.vecs_path)
+# lm_head = torch.load(args.lm_head_path)
+
+# visualize_word(input('Enter a word:'), tokenizer, vecs, lm_head, count=5)
+# visualize_word("fish", vecs, lm_head, count=COUNT)
+
+with gr.Blocks() as demo:
+    gr.Markdown("""
+    ## Backpack visualization: senses lookup
+    > Note: Backpack uses the GPT-2 tokenizer, which includes the space before a word as part of the token, so by default, a space character `' '` is added to the beginning of the word you look up. You can disable this by checking `Remove space before word`, but know this might cause strange behaviors like breaking `afraid` into `af` and `raid`, or `slight` into `s` and `light`.
+    """)
+    with gr.Row():
+        word = gr.Textbox(label="Word")
+        token_breakdown = gr.Textbox(label="Token Breakdown (senses are for the first token only)")
+        remove_space = gr.Checkbox(label="Remove space before word", default=False)
+        count = gr.Slider(minimum=1, maximum=20, value=10, label="Top K", step=1)
+    # sentence = gr.Textbox(label="Sentence")
+    pos_outputs = gr.Dataframe(label="Highest Scoring Senses")
+    neg_outputs = gr.Dataframe(label="Lowest Scoring Senses")
+    gr.Examples(
+    examples=["science", "afraid", "book", "slight"],
+    inputs=[word],
+    outputs=[pos_outputs, neg_outputs, token_breakdown],
+    fn=visualize_word,
+    # cache_examples=True,
+    )
+
+    gr.Button("Look up").click(
+        fn=visualize_word, 
+        inputs= [word, count, remove_space],
+        outputs= [pos_outputs, neg_outputs, token_breakdown],
+    )
+
+    # sentence.select(
+    #     fn=visualize_word,
+    #     inputs= [word, count],
+    #     outputs= [pos_outputs, neg_outputs],
+    # )
+
+demo.launch(share=False)
+