import gradio as gr
from transformers import AutoTokenizer, AutoModelForSeq2SeqLM, Text2TextGenerationPipeline

pipe = Text2TextGenerationPipeline(model = AutoModelForSeq2SeqLM.from_pretrained("jpwahle/t5-word-sense-disambiguation"),
tokenizer = AutoTokenizer.from_pretrained("jpwahle/t5-word-sense-disambiguation"))

def wsd_gen(word, context, d1, d2, d3):
            question = 'question: question: which description describes the word' + ' " ' + word + ' " '
            descriptions_context = 'best in the following context? \descriptions:[  " ' + d1 + '" , " ' + d2 + ' " , or " '+ d3 + ' " ] context: ' + context + "'"
            raw_input = question + descriptions_context
            output = pipe(raw_input)[0]['generated_text']
            return output
word_mask = gr.inputs.Textbox(lines=1, placeholder= "Enter word to disambiguate", default="", label = "Based on the context, which description best matches this word: ")
input_context = gr.inputs.Textbox(lines=1, placeholder="Enter context", default="", label = "context: ")
input_desc1 = gr.inputs.Textbox(lines=1, placeholder="Enter description", default="", label = "description 1: ")
input_desc2 = gr.inputs.Textbox(lines=1, placeholder="Enter description", default="", label = "description 2: ")
input_desc3 = gr.inputs.Textbox(lines=1, placeholder="Enter description", default="", label = "description 3: ")
input_desc4 = gr.inputs.Textbox(lines=1, placeholder="Enter description", default="", label = "description 4: ")
gr.Interface(wsd_gen,
            inputs = [word_mask , input_context, input_desc1, input_desc2, input_desc3],
            outputs= "textbox",
            title = "T5-Word Sense Disambiguation",
            theme = "seafoam",
            allow_flagging="never").launch(inbrowser=True,share=True)