import gradio as gr
from transformers import AutoModelForCausalLM, AutoTokenizer
import torch

model_name = "Hitelcy/sarvix-multilingual-1"

tokenizer = AutoTokenizer.from_pretrained(model_name)
model = AutoModelForCausalLM.from_pretrained(
    model_name,
    torch_dtype=torch.float32,  # CPU doesn't support fp16 well, use float32
    device_map="cpu"
)

def clarify(user_input):
    prompt = f"<|im_start|>user\n{user_input}<|im_end|>\n<|im_start|>assistant\n"
    inputs = tokenizer(prompt, return_tensors="pt")
    out = model.generate(**inputs, max_new_tokens=60)
    response = tokenizer.decode(out[0], skip_special_tokens=True)
    # Extract just the assistant's reply, stripping the prompt echo
    response = response.split("assistant\n")[-1].strip()
    return response

demo = gr.Interface(
    fn=clarify,
    inputs=gr.Textbox(label="Your message"),
    outputs=gr.Textbox(label="Clarification"),
    title="Sarvix Clarify",
)

demo.launch()