import gradio as gr
import torch
from peft import AutoPeftModelForSeq2SeqLM
from transformers import AutoTokenizer

model = AutoPeftModelForSeq2SeqLM.from_pretrained("kietnt0603/randeng-t5-vta-qa-lora")
tokenizer = AutoTokenizer.from_pretrained("IDEA-CCNL/Randeng-T5-784M-QA-Chinese")

device = 'cuda' if torch.cuda.is_available() else 'cpu'

def predict(text):
    input_ids = tokenizer(text, max_length=156, return_tensors="pt", padding="max_length", truncation=True).input_ids.to(device)
    outputs = model.generate(input_ids=input_ids, max_new_tokens=528, do_sample=True)
    pred = tokenizer.batch_decode(outputs.detach().cpu().numpy(), skip_special_tokens=True)[0]
    return pred[len('<extra_id_0>'):]

title = 'VTA-QA Demo'
article = "Loaded model from https://huggingface.co/kietnt0603/randeng-t5-vta-qa-lora"
# Create the Gradio interface
iface = gr.Interface(fn=predict,
                     inputs="textbox",
                     outputs="textbox",
                     title=title,
                     article=article)

# Launch the interface
iface.launch()