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

model_name = "Qwen/Qwen2.5-Coder-1.5B-Instruct"

print("Loading model...")
tokenizer = AutoTokenizer.from_pretrained(model_name)
model = AutoModelForCausalLM.from_pretrained(
    model_name,
    torch_dtype=torch.float32,
    device_map="cpu"
)

def chat(pesan, history):
    messages = [
        {"role": "system", "content": "Kamu adalah AI coding expert."},
        {"role": "user", "content": pesan}
    ]
    text = tokenizer.apply_chat_template(
        messages,
        tokenize=False,
        add_generation_prompt=True
    )
    inputs = tokenizer([text], return_tensors="pt")
    outputs = model.generate(
        **inputs,
        max_new_tokens=2048
    )
    response = tokenizer.decode(
        outputs[0][len(inputs.input_ids[0]):],
        skip_special_tokens=True
    )
    return response

gr.ChatInterface(
    fn=chat,
    title="🤖 AI Coding Assistant",
    description="Tanya apapun tentang coding!"
).launch()