from transformers import T5Tokenizer, T5Model, T5ForConditionalGeneration, T5TokenizerFast, TFT5ForConditionalGeneration, FlaxT5ForConditionalGeneration
import evaluate
import torch
import torch.nn as nn
import pandas as pd
import gradio as gr
import requests

Q_LEN = 256

model_name = 'PRAli22/t5-base-question-answering-system'
tokenizer = T5TokenizerFast.from_pretrained(model_name)
model = T5ForConditionalGeneration.from_pretrained(model_name)

def predict_answer(context, question, ref_answer=None):
    inputs = tokenizer(question, context, max_length=Q_LEN, padding="max_length", truncation=True, add_special_tokens=True)

    input_ids = torch.tensor(inputs["input_ids"], dtype=torch.long).unsqueeze(0)
    attention_mask = torch.tensor(inputs["attention_mask"], dtype=torch.long).unsqueeze(0)

    outputs = model.generate(input_ids=input_ids, attention_mask=attention_mask)

    predicted_answer = tokenizer.decode(outputs.flatten(), skip_special_tokens=True)

    if ref_answer:
        # Load the Bleu metric
        bleu = evaluate.load("google_bleu")
        score = bleu.compute(predictions=[predicted_answer],
                            references=[ref_answer])

        print("Context: \n", context)
        print("\n")
        print("Question: \n", question)
        return {
            "Reference Answer: ": ref_answer,
            "Predicted Answer: ": predicted_answer,
            "BLEU Score: ": score
        }
    else:
        return predicted_answer

css_code='body{background-image:url("https://media.istockphoto.com/id/1256252051/vector/people-using-online-translation-app.jpg?s=612x612&w=0&k=20&c=aa6ykHXnSwqKu31fFR6r6Y1bYMS5FMAU9yHqwwylA94=");}'

demo = gr.Interface(
    fn=predict_answer, 
    inputs=[
        gr.Textbox(label="text", placeholder="Enter the text "),
        gr.Textbox(label="question", placeholder="Enter the question")
    ], 
    outputs=gr.Textbox(label="answer"), 
    title="Question Answering System",
    description= "This is Question Answering System, it takes a text and question in English as inputs and returns it's answer",
    css = css_code
)

demo.launch()