import gradio as gr
import subprocess
import sys


def install(package):
    subprocess.check_call([sys.executable, "-m", "pip", "install", package])

install('transformers')
install('torch')

from transformers import AutoModelForCausalLM, AutoTokenizer
import torch
model = AutoModelForCausalLM.from_pretrained("Sahibsingh12/phi-1-5-finetuned-cazton_complete", trust_remote_code=True, torch_dtype=torch.float32)
tokenizer = AutoTokenizer.from_pretrained("microsoft/phi-1_5", trust_remote_code=True)

def infer(text):
    inputs = tokenizer(f'''question: {text} answer: ''', return_tensors="pt", return_attention_mask=False)
    outputs = model.generate(**inputs, max_length=156)
    text = tokenizer.batch_decode(outputs)[0]
    text = text.split('answer:',1)[1]
    return text



iface = gr.Interface(fn=infer, inputs="text", outputs="text")
iface.launch(share=True)