import gradio as gr
from transformers import pipeline

summarizer = pipeline('summarization')
from transformers import AutoModelForQuestionAnswering, AutoTokenizer, pipeline

model_name = "deepset/roberta-base-squad2"
nlp = pipeline('question-answering', model=model_name, tokenizer=model_name)

examples = [
        [   'Question-Answer',
             '',
            'The option to convert models between FARM and transformers gives freedom to the user and let people easily switch between frameworks.',
            'Why is model conversion important?'
         ],
        [   'Question-Answer',
             '',
            "The Amazon rainforest is a moist broadleaf forest that covers most of the Amazon basin of South America", 
            "Which continent is the Amazon rainforest in?"
        ],
        [   'Question-Answer',
             '',
            'I am a Programmer.',
            'Who am I?' 
        ]
    ]

def summarize_text(text):
    summary = summarizer(text, max_length=130, min_length=30, do_sample=False)
    summary = summary[0]['summary_text']
    return summary

def question_answer(context, question):
    QA_input = {
        'context': context,
        'question': question
    }
    res = nlp(QA_input)
    return (res['answer'])

def home_func(model_choice, summ_text, qa_context, qa_question):
    if model_choice=="Text Summarizer":
        if summ_text == "":
            return "Input correct text to be summarized"
        return summarize_text(summ_text)
    elif model_choice=="Question-Answer":
        if qa_context == "" or qa_question == "":
            return "Choose correct Context and associated questions"
              
        return question_answer(qa_context, qa_question) 

iface = gr.Interface(fn = home_func, 
                     inputs = [gr.inputs.Dropdown(["Text Summarizer", "Question-Answer"], type="value"), 
                               gr.inputs.Textbox(lines=5, placeholder="Enter your text here...", label="Text to be summarized"),
                               gr.inputs.Textbox(lines=5, placeholder="Choose from examples", label="Context"),
                               gr.inputs.Textbox(lines=5, placeholder="Choose from examples", label="Question")],
                     outputs="text",
                     examples=examples)

iface.launch()