kietnt0603's picture
Create app.py
a1cc112 verified
raw
history blame
1.2 kB
import torch
from peft import PeftModel, PeftConfig
from transformers import AutoModelForSeq2SeqLM, AutoTokenizer
config = PeftConfig.from_pretrained("kietnt0603/randeng-t5-vta-qa-lora")
model = AutoModelForSeq2SeqLM.from_pretrained("IDEA-CCNL/Randeng-T5-784M-QA-Chinese")
model = PeftModel.from_pretrained(model, "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
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()