import gradio as gr
import torch
from peft import PeftModel, PeftConfig
from transformers import AutoTokenizer

ref_model = AutoModelForCausalLM.from_pretrained("EleutherAI/pythia-70m-deduped-v0", torch_dtype=torch.bfloat16)
peft_model_id = "w601sxs/pythia-70m-instruct-orca-chkpt-64000"

config = PeftConfig.from_pretrained(peft_model_id)
model = PeftModel.from_pretrained(ref_model, peft_model_id)
tokenizer = AutoTokenizer.from_pretrained(config.base_model_name_or_path)

model.eval()

def predict(text):
    inputs = tokenizer(prompt, return_tensors="pt")
    with torch.no_grad():
    outputs = model.generate(input_ids=inputs["input_ids"], max_new_tokens=10)
    out_text = tokenizer.batch_decode(outputs.detach().cpu().numpy(), skip_special_tokens=True)[0])
    
    return out_text
    

demo = gr.Interface(
    fn=predict, 
    inputs='text',
    outputs='text',
)

demo.launch()