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

model_id = "mistralai/Mistral-Small-Instruct-2409"
tokenizer = AutoTokenizer.from_pretrained(model_id)
model = AutoModelForCausalLM.from_pretrained(
    model_id,
    torch_dtype="auto",
    device_map="auto"
)

@spaces.GPU
def respond(message, history):
    inputs = tokenizer(message, return_tensors="pt").to("cuda")
    inputs_size = len(inputs.input_ids[0])
    response = model.generate(
        **inputs,
        max_new_tokens=256,
        do_sample=True,
        temperature=0.2,
        top_p=0.9,
        repetition_penalty=1.5,
        eos_token_id=tokenizer.eos_token_id,
    )
    output = tokenizer.decode(response[0][inputs_size:], skip_special_tokens=True)
    return output

app = gr.ChatInterface(fn=respond, title="Simple Chat")

if __name__ == "__main__":
    app.launch()