visualize senses in language modeling

app.py

@@ -1,13 +1,23 @@
 import torch
-import pandas as pd
 import transformers
+from transformers import AutoModelForCausalLM
+import pandas as pd
 import gradio as gr
 
+# Build model & get some layers
 tokenizer = transformers.AutoTokenizer.from_pretrained('gpt2')
-vecs = torch.load("senses/all_vecs_mtx.pt")
-lm_head = torch.load("senses/lm_head.pt")
+m = AutoModelForCausalLM.from_pretrained("lora-x/backpack-gpt2", trust_remote_code=True)
+m.eval()
 
+lm_head = m.get_lm_head() # (V, d)
+word_embeddings = m.backpack.get_word_embeddings() # (V, d)
+sense_network = m.backpack.get_sense_network() # (V, nv, d)
+num_senses = m.backpack.get_num_senses()
+sense_names = [i for i in range(num_senses)]
 
+"""
+Single token sense lookup
+"""
 def visualize_word(word, count=10, remove_space=False):
 
     if not remove_space:
@@ -19,14 +29,17 @@ def visualize_word(word, count=10, remove_space=False):
     tokens = ", ".join(tokens) # display tokenization for user
     print(f"Tokenized as: {tokens}")
     # look up sense vectors only for the first token
-    contents = vecs[token_ids[0]] # torch.Size([16, 768])
+    # contents = vecs[token_ids[0]] # torch.Size([16, 768])
+    sense_input_embeds = word_embeddings(torch.tensor([token_ids[0]]).long().unsqueeze(0)) # (bs=1, s=1, d), sense_network expects bs dim
+    senses = sense_network(sense_input_embeds) # -> (bs=1, nv, s=1, d) 
+    senses = torch.squeeze(senses) # (nv, s=1, d)
 
     # for pos and neg respectively, create a list (for each sense) of list (top k) of tuples (word, logit)
     pos_word_lists = [] 
     neg_word_lists = []
     sense_names = [] # column header
-    for i in range(contents.shape[0]):
-        logits = contents[i,:] @ lm_head.t() # (vocab,)    [768] @ [768, 50257] -> [50257]
+    for i in range(senses.shape[0]):
+        logits = lm_head(senses[i,:])
         sorted_logits, sorted_indices = torch.sort(logits, descending=True)
         sense_names.append('sense {}'.format(i))
 
@@ -54,30 +67,430 @@ def visualize_word(word, count=10, remove_space=False):
 
     return pos_df, neg_df, tokens
 
-with gr.Blocks() as demo:
+"""
+Returns:
+    - tokens: the tokenization of the input sentence, also used as options to choose from for get_token_contextual_weights
+    - top_k_words_df: a dataframe of the top k words predicted by the model
+    - length: of the input sentence, stored as a gr.State variable so other methods can find the 
+        contextualization weights for the *last* token that's needed
+    - contextualization_weights: gr.State variable, stores the contextualization weights for the input sentence
+"""
+def predict_next_word (sentence, top_k = 5, contextualization_weights = None):  
+
+    # For better tokenization, by default, adds a space at the beginning of the sentence if it doesn't already have one
+    # and remove trailing space
+    sentence = sentence.strip()
+    if sentence[0] != ' ':
+        sentence = ' ' + sentence
+    print(f"Sentence: '{sentence}'")
+
+    # Make input, keeping track of original length
+    token_ids = tokenizer(sentence)['input_ids']
+    tokens = [[tokenizer.decode(token_id) for token_id in token_ids]] # a list of a single list because used as dataframe
+    length = len(token_ids)
+    inp = torch.zeros((1,512)).long()
+    inp[0,:length] = torch.tensor(token_ids).long()
+
+    # Get output at correct index
+    if contextualization_weights is None:
+        print("contextualization_weights IS None, freshly computing contextualization_weights")
+        output = m(inp)
+        logits, contextualization_weights = output.logits[0,length-1,:], output.contextualization
+        # Store contextualization weights and return it as a gr.State var for use by get_token_contextual_weights
+    else:
+        print("contextualization_weights is NOT None, using passed in contextualization_weights")
+        output = m.run_with_custom_contextualization(inp, contextualization_weights)
+        logits = output.logits[0,length-1,:]
+    probs = logits.softmax(dim=-1) # probs over next word
+    probs, indices = torch.sort(probs, descending=True)
+    top_k_words = [(tokenizer.decode(indices[i]), round(probs[i].item(), 3)) for i in range(top_k)]
+    top_k_words_df = pd.DataFrame(top_k_words, columns=['word', 'probability'], index=range(1, top_k+1))
+
+    top_k_words_df = top_k_words_df.T
+
+    print(top_k_words_df)
+
+    return tokens, top_k_words_df, length, contextualization_weights
+
+
+"""
+Returns a dataframe of senses with weights for the selected token.
+
+Args:
+    contextualization_weights: a gr.State variable that stores the contextualization weights for the input sentence.
+    length: length of the input sentence, used to get the contextualization weights for the last token
+    token: the selected token
+    token_index: the index of the selected token in the input sentence
+    count: how many top words to display for each sense
+"""
+def get_token_contextual_weights (contextualization_weights, length, token, token_index, count = 7):
+    print(">>>>>in get_token_contextual_weights")
+    print(f"Selected {token_index}th token: {token}")
+    
+    # get contextualization weights for the selected token
+    # Only care about the weights for the last word, since that's what contributes to the output
+    token_contextualization_weights = contextualization_weights[0, :, length-1, token_index] 
+    token_contextualization_weights_list = [round(x, 3) for x in token_contextualization_weights.tolist()]
+
+    # get sense vectors of the selected token
+    token_ids = tokenizer(token)['input_ids'] # keep as a list bc sense_network expects s dim
+    sense_input_embeds = word_embeddings(torch.tensor(token_ids).long().unsqueeze(0)) # (bs=1, s=1, d), sense_network expects bs dim
+    senses = sense_network(sense_input_embeds) # -> (bs=1, nv, s=1, d) 
+    senses = torch.squeeze(senses) # (nv, s=1, d)
+
+    # build dataframe
+    neg_word_lists = []
+    pos_dfs, neg_dfs = [], []
+
+    for i in range(num_senses):
+        logits = lm_head(senses[i,:]) # (vocab,) [768, 50257] -> [50257]
+        sorted_logits, sorted_indices = torch.sort(logits, descending=True)
+
+        pos_sorted_words = [tokenizer.decode(sorted_indices[j]) for j in range(count)]
+        pos_df = pd.DataFrame(pos_sorted_words)
+        pos_dfs.append(pos_df)
+
+        neg_sorted_words = [tokenizer.decode(sorted_indices[-j-1]) for j in range(count)]
+        neg_df = pd.DataFrame(neg_sorted_words)
+        neg_dfs.append(neg_df)
+
+    sense0words, sense1words, sense2words, sense3words, sense4words, sense5words, \
+        sense6words, sense7words, sense8words, sense9words, sense10words, sense11words, \
+            sense12words, sense13words, sense14words, sense15words = pos_dfs
+    
+    sense0slider, sense1slider, sense2slider, sense3slider, sense4slider, sense5slider, \
+        sense6slider, sense7slider, sense8slider, sense9slider, sense10slider, sense11slider, \
+            sense12slider, sense13slider, sense14slider, sense15slider = token_contextualization_weights_list
+    
+    return token, token_index, sense0words, sense1words, sense2words, sense3words, sense4words, sense5words, sense6words, \
+        sense7words, sense8words, sense9words, sense10words, sense11words, sense12words, sense13words, sense14words, sense15words, \
+            sense0slider, sense1slider, sense2slider, sense3slider, sense4slider, sense5slider, sense6slider, sense7slider, \
+                sense8slider, sense9slider, sense10slider, sense11slider, sense12slider, sense13slider, sense14slider, sense15slider
+
+"""
+Wrapper for when the user selects a new token in the tokens dataframe. 
+Converts `evt` (the selected token) to `token` and `token_index` which are used by get_token_contextual_weights.
+"""
+def new_token_contextual_weights (contextualization_weights, length, evt: gr.SelectData, count = 7):
+    print(">>>>>in new_token_contextual_weights")
+    token_index = evt.index[1] # selected token is the token_index-th token in the sentence
+    token = evt.value
+    if not token:
+        return None, None, None, None, None, None, None, None, None, None, None, None, None, None, None, \
+            None, None, None, None, None, None, None, None, None, None, None, None, None, None, None, \
+                None, None, None, None, None, None, None, None, None, None, None, None, None, None, None
+    return get_token_contextual_weights (contextualization_weights, length, token, token_index, count)
+
+def change_sense0_weight(contextualization_weights, length, token_index, new_weight):
+    contextualization_weights[0, 0, length-1, token_index] = new_weight
+    return contextualization_weights
+def change_sense1_weight(contextualization_weights, length, token_index, new_weight):
+    contextualization_weights[0, 1, length-1, token_index] = new_weight
+    return contextualization_weights
+def change_sense2_weight(contextualization_weights, length, token_index, new_weight):
+    contextualization_weights[0, 2, length-1, token_index] = new_weight
+    return contextualization_weights
+def change_sense3_weight(contextualization_weights, length, token_index, new_weight):
+    contextualization_weights[0, 3, length-1, token_index] = new_weight
+    return contextualization_weights
+def change_sense4_weight(contextualization_weights, length, token_index, new_weight):
+    contextualization_weights[0, 4, length-1, token_index] = new_weight
+    return contextualization_weights
+def change_sense5_weight(contextualization_weights, length, token_index, new_weight):
+    contextualization_weights[0, 5, length-1, token_index] = new_weight
+    return contextualization_weights
+def change_sense6_weight(contextualization_weights, length, token_index, new_weight):
+    contextualization_weights[0, 6, length-1, token_index] = new_weight
+    return contextualization_weights
+def change_sense7_weight(contextualization_weights, length, token_index, new_weight):
+    contextualization_weights[0, 7, length-1, token_index] = new_weight
+    return contextualization_weights
+def change_sense8_weight(contextualization_weights, length, token_index, new_weight):
+    contextualization_weights[0, 8, length-1, token_index] = new_weight
+    return contextualization_weights
+def change_sense9_weight(contextualization_weights, length, token_index, new_weight):
+    contextualization_weights[0, 9, length-1, token_index] = new_weight
+    return contextualization_weights
+def change_sense10_weight(contextualization_weights, length, token_index, new_weight):
+    contextualization_weights[0, 10, length-1, token_index] = new_weight
+    return contextualization_weights
+def change_sense11_weight(contextualization_weights, length, token_index, new_weight):
+    contextualization_weights[0, 11, length-1, token_index] = new_weight
+    return contextualization_weights
+def change_sense12_weight(contextualization_weights, length, token_index, new_weight):
+    contextualization_weights[0, 12, length-1, token_index] = new_weight
+    return contextualization_weights
+def change_sense13_weight(contextualization_weights, length, token_index, new_weight):
+    contextualization_weights[0, 13, length-1, token_index] = new_weight
+    return contextualization_weights
+def change_sense14_weight(contextualization_weights, length, token_index, new_weight):
+    contextualization_weights[0, 14, length-1, token_index] = new_weight
+    return contextualization_weights
+def change_sense15_weight(contextualization_weights, length, token_index, new_weight):
+    contextualization_weights[0, 15, length-1, token_index] = new_weight
+    return contextualization_weights
+
+"""
+Clears all gr.State variables used to store info across methods when the input sentence changes.
+"""
+def clear_states(contextualization_weights, token_index, length):
+    contextualization_weights = None
+    token_index = None
+    length = 0
+    return contextualization_weights, token_index, length
+
+def reset_weights(contextualization_weights):
+    print("Resetting weights...")
+    contextualization_weights = None
+    return contextualization_weights
+
+with gr.Blocks( css = """#sense0slider, #sense1slider, #sense2slider, #sense3slider, #sense4slider, #sense5slider, #sense6slider, #sense7slider,
+                        #sense8slider, #sense9slider, #sense1slider0, #sense11slider, #sense12slider, #sense13slider, #sense14slider, #sense15slider 
+                        { height: 200px; width: 200px;  transform: rotate(270deg); }""" ) 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`.
+    ## Backpack Sense Visualization
     """)
-    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)
-    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],
-    )
-
-demo.launch(auth=("caesar", "wins"))
+
+    with gr.Tab("Language Modeling"):
+        contextualization_weights = gr.State(None) # store session data for sharing between functions
+        token_index = gr.State(None)
+        length = gr.State(0)
+        top_k = gr.State(10)
+        with gr.Row():
+            with gr.Column(scale=8):
+                input_sentence = gr.Textbox(label="Input Sentence", placeholder='Enter a sentence and click "Predict next word"')
+            with gr.Column(scale=1):
+                predict = gr.Button(value="Predict next word", variant="primary")
+                reset_weights_button = gr.Button("Reset weights")
+        top_k_words = gr.Dataframe(label="Next Word Predictions (top k)")
+        gr.Markdown("""### **Tokens:** click on a token to see its senses and contextualization weights""")
+        tokens = gr.DataFrame(label="")
+        with gr.Row():
+            with gr.Column(scale=1):
+                selected_token = gr.Textbox(label="Current Selected Token", interactive=False)
+            with gr.Column(scale=8):
+                gr.Markdown("""#####
+                Once a token is chosen, you can use the sliders below to change the weights of any senses for that token, \
+                    and then click "Predict next word" to see updated next-word predictions. \
+                    You can change the weights of *multiple senses of multiple tokens*; \
+                    changes will be preserved until you click "Reset weights".
+                """)
+        # sense sliders and top sense words dataframes
+        with gr.Row():
+            with gr.Column(scale=0, min_width=120):
+                sense0slider= gr.Slider(minimum=0, maximum=1, value=0, step=0.01, label="Sense 0", elem_id="sense0slider", interactive=True)
+            with gr.Column(scale=0, min_width=120):
+                sense1slider= gr.Slider(minimum=0, maximum=1, value=0, step=0.01, label="Sense 1", elem_id="sense1slider", interactive=True)
+            with gr.Column(scale=0, min_width=120):
+                sense2slider= gr.Slider(minimum=0, maximum=1, value=0, step=0.01, label="Sense 2", elem_id="sense2slider", interactive=True)
+            with gr.Column(scale=0, min_width=120):
+                sense3slider= gr.Slider(minimum=0, maximum=1, value=0, step=0.01, label="Sense 3", elem_id="sense3slider", interactive=True)
+            with gr.Column(scale=0, min_width=120):
+                sense4slider= gr.Slider(minimum=0, maximum=1, value=0, step=0.01, label="Sense 4", elem_id="sense4slider", interactive=True)
+            with gr.Column(scale=0, min_width=120):
+                sense5slider= gr.Slider(minimum=0, maximum=1, value=0, step=0.01, label="Sense 5", elem_id="sense5slider", interactive=True)
+            with gr.Column(scale=0, min_width=120):
+                sense6slider= gr.Slider(minimum=0, maximum=1, value=0, step=0.01, label="Sense 6", elem_id="sense6slider", interactive=True)
+            with gr.Column(scale=0, min_width=120):
+                sense7slider= gr.Slider(minimum=0, maximum=1, value=0, step=0.01, label="Sense 7", elem_id="sense7slider", interactive=True)
+        with gr.Row():
+            with gr.Column(scale=0, min_width=120):
+                sense0words = gr.DataFrame()
+            with gr.Column(scale=0, min_width=120):
+                sense1words = gr.DataFrame()
+            with gr.Column(scale=0, min_width=120):
+                sense2words = gr.DataFrame()
+            with gr.Column(scale=0, min_width=120):
+                sense3words = gr.DataFrame()
+            with gr.Column(scale=0, min_width=120):
+                sense4words = gr.DataFrame()
+            with gr.Column(scale=0, min_width=120):
+                sense5words = gr.DataFrame()
+            with gr.Column(scale=0, min_width=120):
+                sense6words = gr.DataFrame()
+            with gr.Column(scale=0, min_width=120):
+                sense7words = gr.DataFrame()
+        with gr.Row():
+            with gr.Column(scale=0, min_width=120):
+                sense8slider= gr.Slider(minimum=0, maximum=1, value=0, step=0.01, label="Sense 8", elem_id="sense8slider", interactive=True)
+            with gr.Column(scale=0, min_width=120):
+                sense9slider= gr.Slider(minimum=0, maximum=1, value=0, step=0.01, label="Sense 9", elem_id="sense9slider", interactive=True)
+            with gr.Column(scale=0, min_width=120):
+                sense10slider= gr.Slider(minimum=0, maximum=1, value=0, step=0.01, label="Sense 10", elem_id="sense1slider0", interactive=True)
+            with gr.Column(scale=0, min_width=120):
+                sense11slider= gr.Slider(minimum=0, maximum=1, value=0, step=0.01, label="Sense 11", elem_id="sense11slider", interactive=True)
+            with gr.Column(scale=0, min_width=120):
+                sense12slider= gr.Slider(minimum=0, maximum=1, value=0, step=0.01, label="Sense 12", elem_id="sense12slider", interactive=True)
+            with gr.Column(scale=0, min_width=120):
+                sense13slider= gr.Slider(minimum=0, maximum=1, value=0, step=0.01, label="Sense 13", elem_id="sense13slider", interactive=True)
+            with gr.Column(scale=0, min_width=120):
+                sense14slider= gr.Slider(minimum=0, maximum=1, value=0, step=0.01, label="Sense 14", elem_id="sense14slider", interactive=True)
+            with gr.Column(scale=0, min_width=120):
+                sense15slider= gr.Slider(minimum=0, maximum=1, value=0, step=0.01, label="Sense 15", elem_id="sense15slider", interactive=True)
+        with gr.Row():
+            with gr.Column(scale=0, min_width=120):
+                sense8words = gr.DataFrame()
+            with gr.Column(scale=0, min_width=120):
+                sense9words = gr.DataFrame()
+            with gr.Column(scale=0, min_width=120):
+                sense10words = gr.DataFrame()
+            with gr.Column(scale=0, min_width=120):
+                sense11words = gr.DataFrame()
+            with gr.Column(scale=0, min_width=120):
+                sense12words = gr.DataFrame()
+            with gr.Column(scale=0, min_width=120):
+                sense13words = gr.DataFrame()
+            with gr.Column(scale=0, min_width=120):
+                sense14words = gr.DataFrame()
+            with gr.Column(scale=0, min_width=120):
+                sense15words = gr.DataFrame()
+        
+        # gr.Examples(
+        #     examples=[["Messi plays for", top_k, None]],
+        #     inputs=[input_sentence, top_k, contextualization_weights],
+        #     outputs=[tokens, top_k_words, length, contextualization_weights],
+        #     fn=predict_next_word,
+        # )
+
+        sense0slider.change(fn=change_sense0_weight,
+                            inputs=[contextualization_weights, length, token_index, sense0slider],
+                            outputs=[contextualization_weights])
+        sense1slider.change(fn=change_sense1_weight,
+                            inputs=[contextualization_weights, length, token_index, sense1slider],
+                            outputs=[contextualization_weights])
+        sense2slider.change(fn=change_sense2_weight,
+                            inputs=[contextualization_weights, length, token_index, sense2slider],
+                            outputs=[contextualization_weights])
+        sense3slider.change(fn=change_sense3_weight,
+                            inputs=[contextualization_weights, length, token_index, sense3slider],
+                            outputs=[contextualization_weights])
+        sense4slider.change(fn=change_sense4_weight,
+                            inputs=[contextualization_weights, length, token_index, sense4slider],
+                            outputs=[contextualization_weights])
+        sense5slider.change(fn=change_sense5_weight,
+                            inputs=[contextualization_weights, length, token_index, sense5slider],
+                            outputs=[contextualization_weights])
+        sense6slider.change(fn=change_sense6_weight,
+                            inputs=[contextualization_weights, length, token_index, sense6slider],
+                            outputs=[contextualization_weights])
+        sense7slider.change(fn=change_sense7_weight,
+                            inputs=[contextualization_weights, length, token_index, sense7slider],
+                            outputs=[contextualization_weights])
+        sense8slider.change(fn=change_sense8_weight,
+                            inputs=[contextualization_weights, length, token_index, sense8slider],
+                            outputs=[contextualization_weights])
+        sense9slider.change(fn=change_sense9_weight,
+                            inputs=[contextualization_weights, length, token_index, sense9slider],
+                            outputs=[contextualization_weights])
+        sense10slider.change(fn=change_sense10_weight,
+                            inputs=[contextualization_weights, length, token_index, sense10slider],
+                            outputs=[contextualization_weights])
+        sense11slider.change(fn=change_sense11_weight,
+                            inputs=[contextualization_weights, length, token_index, sense11slider],
+                            outputs=[contextualization_weights])
+        sense12slider.change(fn=change_sense12_weight,
+                            inputs=[contextualization_weights, length, token_index, sense12slider],
+                            outputs=[contextualization_weights])
+        sense13slider.change(fn=change_sense13_weight,
+                            inputs=[contextualization_weights, length, token_index, sense13slider],
+                            outputs=[contextualization_weights])
+        sense14slider.change(fn=change_sense14_weight,
+                            inputs=[contextualization_weights, length, token_index, sense14slider],
+                            outputs=[contextualization_weights])
+        sense15slider.change(fn=change_sense15_weight,
+                            inputs=[contextualization_weights, length, token_index, sense15slider],
+                            outputs=[contextualization_weights])
+    
+        predict.click(
+            fn=predict_next_word,
+            inputs = [input_sentence, top_k, contextualization_weights], 
+            outputs= [tokens, top_k_words, length, contextualization_weights],
+        )
+
+        tokens.select(fn=new_token_contextual_weights, 
+                    inputs=[contextualization_weights, length], 
+                    outputs= [selected_token, token_index, 
+                                
+                                sense0words, sense1words, sense2words, sense3words, sense4words, sense5words, sense6words, sense7words, 
+                                sense8words, sense9words, sense10words, sense11words, sense12words, sense13words, sense14words, sense15words,
+
+                                sense0slider, sense1slider, sense2slider, sense3slider, sense4slider, sense5slider, sense6slider, sense7slider, 
+                                sense8slider, sense9slider, sense10slider, sense11slider, sense12slider, sense13slider, sense14slider, sense15slider]
+                                )
+
+        reset_weights_button.click(
+                fn=reset_weights,
+                inputs=[contextualization_weights],
+                outputs=[contextualization_weights]
+            ).success(
+                fn=predict_next_word,
+                inputs = [input_sentence, top_k, contextualization_weights],
+                outputs= [tokens, top_k_words, length, contextualization_weights],
+            ).success(
+                fn=get_token_contextual_weights,
+                inputs=[contextualization_weights, length, selected_token, token_index],
+                outputs= [selected_token, token_index, 
+                        
+                        sense0words, sense1words, sense2words, sense3words, sense4words, sense5words, sense6words, sense7words, 
+                        sense8words, sense9words, sense10words, sense11words, sense12words, sense13words, sense14words, sense15words,
+                        
+                        sense0slider, sense1slider, sense2slider, sense3slider, sense4slider, sense5slider, sense6slider, sense7slider, 
+                        sense8slider, sense9slider, sense10slider, sense11slider, sense12slider, sense13slider, sense14slider, sense15slider]
+                        )
+        
+        input_sentence.change(
+            fn=clear_states,
+            inputs=[contextualization_weights, token_index, length],
+            outputs=[contextualization_weights, token_index, length]
+        )
+
+    with gr.Tab("Individual Word Sense Look Up"):
+        gr.Markdown("""> Note on tokenization: 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", placeholder="e.g. science")
+            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)
+        look_up_button = gr.Button("Look up")
+        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,
+        )
+
+        look_up_button.click(
+            fn=visualize_word, 
+            inputs= [word, count, remove_space],
+            outputs= [pos_outputs, neg_outputs, token_breakdown],
+        )
+
+demo.launch(auth=("caesar", "wins"))
+
+
+# Code for generating slider functions & event listners
+
+# for i in range(16):
+#     print(
+#         f"""def change_sense{i}_weight(contextualization_weights, length, token_index, new_weight):
+#     print(f"Changing weight for the {i}th sense of the {{token_index}}th token.")
+#     print("new_weight to be assigned = ", new_weight)
+#     contextualization_weights[0, {i}, length-1, token_index] = new_weight
+#     print("contextualization_weights: ", contextualization_weights[0, :, length-1, token_index])
+#     return contextualization_weights"""
+#     )
+
+# for i in range(16):
+#     print(
+#         f"""  sense{i}slider.change(fn=change_sense{i}_weight,
+#                     inputs=[contextualization_weights, length, token_index, sense{i}slider],
+#                     outputs=[contextualization_weights])"""
+#     )