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

model_name = "ICEPVP8977/Uncensored_Qwen2.5_Coder_3B"
tokenizer_name = "Qwen/Qwen2.5-Coder-3B-Instruct"  # ambil tokenizer dari sini

tokenizer = AutoTokenizer.from_pretrained(
    tokenizer_name,
    trust_remote_code=True
)
model = AutoModelForCausalLM.from_pretrained(
    model_name,
    trust_remote_code=True,
    torch_dtype=torch.float32,
    low_cpu_mem_usage=True
)
model.eval()

def generate(prompt):
    messages = [{"role": "user", "content": prompt}]
    text = tokenizer.apply_chat_template(
        messages,
        tokenize=False,
        add_generation_prompt=True
    )
    inputs = tokenizer(text, return_tensors="pt", truncation=True, max_length=512)
    with torch.no_grad():
        outputs = model.generate(
            **inputs,
            max_new_tokens=256,
            do_sample=True,
            temperature=0.7,
            pad_token_id=tokenizer.eos_token_id
        )
    result = tokenizer.decode(
        outputs[0][inputs["input_ids"].shape[-1]:],
        skip_special_tokens=True
    )
    return result.strip()

gr.Interface(
    fn=generate,
    inputs=gr.Textbox(lines=4, label="Tanya apa saja"),
    outputs=gr.Textbox(label="Jawaban"),
    title="AI Assistant",
).launch()