import gradio as gr
import pickle
import pandas as pd

data = pickle.load(open("OOV_Train_2.pkl", "rb"))
data = pd.DataFrame(
    data, 
    columns=["Input_Seq", "Label", "Adj_Class", "Adj", "Nn", "Hypr", "Adj_NN"]
)

adjs = set(data["Adj"])
Nns  = set(list(data["Nn"]) + list(data["Hypr"]))
all_set = set(list(adjs) + list(Nns))

def test_input(words):
    word_dict = {}
    for w in words.split(","):
        if w in all_set:
            word_dict[w] = 1.
        else:
            word_dict[w] = 0.
    return word_dict

    
title = "BERT on a PLANE"

description = """
Did you know that logically speaking **A small cat is not a small animal**, and that **A fake smile is not a smile**? Learn more by testing our BERT model tuned to perform phrase-level adjective-noun entailment, via the [PLANE](https://aclanthology.org/2022.coling-1.359/) dataset.

Please note that the scope of the model is not to run lexical entailment or hypernym detection (e.g., *"A dog is an animal*"), but to perform a very specific subset of phrase-level entailment, based on adjective-nouns phrases. The type of question you can ask the model are limited, and should have one of three forms:

- An *Adjective-Noun* is a *Noun* (e.g. A red car is a car)

- An *Adjective-Noun* is a *Hypernym(Noun)* (e.g. A red car is a vehicle)

- An *Adjective-Noun* is a *Adjective-Hypernym(Noun)* (e.g. A red car is a vehicle)

Please note that, as in the examples above, the adjective should be the same for both phrases, and that the Hypernym(Noun) shoul be a true hypernym of the selected noun

The current model achieves an accuracy of 90% on out-of-distribution evaluation.
Use the next page to check if your test-items (i.e. adjective, noun and hypernyms) were part of the training data!"""

examples = [["A red car is a vehicle"], ["A fake smile is a smile"], ["A small cat is a small animal"]]
examples_w = [["red,car,vehicle"], ["fake,smile"], ["small,cat,animal"]]


interface_words = gr.Interface(
            fn=test_input,
            inputs=gr.Textbox(label="Input:word_1,word2,...,word_n"),
            outputs=gr.outputs.Label(label="In training-distribution?"),
            examples=examples_w,
)
# interface_words.launch()

interface_model = gr.Interface.load(
            "huggingface/lorenzoscottb/bert-base-cased-PLANE-ood-2",
            description=description,
            examples=examples,
            title=title,
)
# interface_model.launch()


gr.TabbedInterface(
    [interface_model, interface_words], ["Test Model", "Check if words in/out-distribution"]
).launch()