import gradio as gr
import torch
from transformers import AutoModelForCausalLM, AutoTokenizer
torch.set_default_device("cuda")

model = AutoModelForCausalLM.from_pretrained("microsoft/phi-2", torch_dtype="auto", trust_remote_code=True)
tokenizer = AutoTokenizer.from_pretrained("microsoft/phi-2", trust_remote_code=True)


def generate_primes(text_input):
    inputs = tokenizer(text_input, return_tensors="pt", return_attention_mask=False)
    outputs = model.generate(**inputs, max_length=200)
    generated_text = tokenizer.batch_decode(outputs)[0]
    return generated_text

iface = gr.Interface(
    fn=generate_primes,
    inputs=gr.Textbox(),
    outputs=gr.Textbox(),
    live=True,
    title="Prime Generator",
    description="Generate primes based on input text."
)

iface.launch()