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

tokenizer = AutoTokenizer.from_pretrained("VietAI/vit5-base")
model = AutoModelForSeq2SeqLM.from_pretrained("quocanh944/viT5-med-qa")

def generate_answer(question):
    global model, tokenizer
    model.eval()
    input_text = "hỏi: " + question
    inputs = tokenizer(input_text, return_tensors="pt", max_length=512, truncation=True, padding="max_length")
    input_ids = inputs.input_ids
    attention_mask = inputs.attention_mask

    with torch.no_grad():
        outputs = model.generate(input_ids=input_ids, attention_mask=attention_mask, max_length=128, num_beams=4, early_stopping=True)

    answer = tokenizer.decode(outputs[0], skip_special_tokens=True)
    return answer

title = "Interactive demo: ViT5 with Medical Dataset"
description = "Demo for ViT5 with Medical Dataset. The model is fine-tuned on a Vietnamese medical dataset. The model is able to answer questions related to medical knowledge. Please input your question in the textbox and click submit to get the answer."
article = "This is a demo for ViT5 with Medical Dataset. The model is fine-tuned on a Vietnamese medical dataset. The model is able to answer questions related to medical knowledge. Please input your question in the textbox and click submit to get the answer."
examples = ["Tôi bị đau tay thì nên làm gì?", "Covid-19 là gì?", "Tôi nên làm gì khi bị sùi mào gà?", "Tôi nên ăn gì để tăng cân?"]

iface = gr.Interface(fn=generate_answer, 
                     inputs=gr.Textbox(),
                     outputs=gr.Textbox(),
                     title=title,
                     description=description,
                     article=article,
                     examples=examples)

iface.launch()