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

model_name = "huihui-ai/Qwen2.5-Coder-3B-Instruct-abliterated"
print("Loading tokenizer...")
tokenizer = AutoTokenizer.from_pretrained(model_name)
print("Loading model...")
model = AutoModelForCausalLM.from_pretrained(
    model_name,
    torch_dtype=torch.float16,
    device_map="cpu",
    low_cpu_mem_usage=True,
)
print("Model ready!")

def chat(message, history):
    messages = [{"role": "user", "content": message}]
    text = tokenizer.apply_chat_template(messages, tokenize=False, add_generation_prompt=True)
    inputs = tokenizer(text, return_tensors="pt")
    with torch.no_grad():
        outputs = model.generate(**inputs, max_new_tokens=512, temperature=0.7)
    response = tokenizer.decode(outputs[0], skip_special_tokens=True)
    if response.startswith(text):
        response = response[len(text):]
    return response.strip()

demo = gr.ChatInterface(
    fn=chat,
    title="Uncensored Coder 3B",
    description="Qwen2.5-Coder-3B abliterated - uncensored code model",
)
demo.launch()