import gradio as gr
from transformers import pipeline

#Loading bert and roberta model for question - answer comparison.
bert_pipeline = pipeline("question-answering", model="bert-large-uncased-whole-word-masking-finetuned-squad")
roberta_pipeline = pipeline("question-answering", model="deepset/roberta-base-squad2")
#Function to get answers from both models. 
def answer_question(c,q):
    bert_answer=bert_pipeline(question=q, context=c)['answer']
    roberta_answer=roberta_pipeline(question=q, context=c)['answer']
    return bert_answer, roberta_answer

    
#Interface using gradio
"""with gr.Blocks() as demo:
    gr.Markdown("##Question Answering session with BERT and RoBERTa Model")
    context_input=gr.Textbox(label='Context', placeholder="Enter the context information here...")
    question_input=gr.Textbox(label='Question', placeholder="Enter the question here one by one...")
    bert_output=gr.Textbox(label='BERT Answer')
    roberta_output=gr.Textbox(label='Roberta Answer')
    submit_button=gr.Button("Get Answers")"""

with gr.Blocks(theme=gr.themes.Soft(primary_hue="blue", secondary_hue='orange')) as demo:
    gr.Markdown("##Question Answering session with BERT and RoBERTa Model")

    with gr.Row():
        with gr.Column():
            context_input=gr.Textbox(label='Context', placeholder="Enter the context information here...")
            question_input=gr.Textbox(label='Question', placeholder="Enter the question here one by one...")
            submit_button=gr.Button("Get Answers")

        with gr.Column():
            bert_output=gr.Textbox(label='BERT Answer')
            roberta_output=gr.Textbox(label='Roberta Answer')
    with gr.Row():
        gr.Markdown("<div style='color: blue;'>BERT Interpretation of the Context.  </div>")
        gr.Markdown("<div style='color: purple;'>ROBERTA Interpretation of the Context.  </div>")
    #Click Action
    submit_button.click(answer_question, inputs=[context_input, question_input], outputs=[bert_output, roberta_output])

demo.launch()