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))

        # 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

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],
    )

demo.launch(share=True)