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

model_id = "Itaking/itakura_v4-model"
tokenizer = AutoTokenizer.from_pretrained(model_id)
model = AutoModelForCausalLM.from_pretrained(model_id, torch_dtype=torch.bfloat16)

def chat(message, history):
    messages = [{"role": "user", "content": message}]
    input_text = tokenizer.apply_chat_template(messages, tokenize=False, add_generation_prompt=True)
    inputs = tokenizer(input_text, return_tensors="pt")
    inputs.pop("token_type_ids", None)
    with torch.inference_mode():
        output = model.generate(**inputs, max_new_tokens=200, temperature=0.7,
                                top_p=0.9, do_sample=True, repetition_penalty=1.3)
    return tokenizer.decode(output[0][inputs["input_ids"].shape[1]:], skip_special_tokens=True)

gr.ChatInterface(chat, title="demo_model").launch()