import gradio as gr
import torch
from transformers import AutoModelForCausalLM, AutoTokenizer, TextIteratorStreamer
from peft import PeftModel
from threading import Thread

# 1. Map both coordinates: The base model engine and your custom adapter layer
BASE_MODEL = "Qwen/Qwen2.5-Coder-3B-Instruct"
ADAPTER_MODEL = "Cydercoder/qwen2.5-coder-3b"

print("Loading official base tokenizer...")
tokenizer = AutoTokenizer.from_pretrained(BASE_MODEL)

print("Loading public base model on CPU...")
base_model = AutoModelForCausalLM.from_pretrained(
    BASE_MODEL,
    torch_dtype=torch.float32,
    device_map="cpu"
)

print("Merging your custom fine-tuned engineering weights...")
# This layers your specialized tasks right over the active model architecture
model = PeftModel.from_pretrained(base_model, ADAPTER_MODEL)

def chat_function(message, history):
    messages = [
        {"role": "system", "content": "You are an expert full-stack developer assistant fine-tuned for frontend, backend, animations, and debugging."}
    ]
    
    for user_msg, bot_msg in history:
        messages.append({"role": "user", "content": user_msg})
        messages.append({"role": "assistant", "content": bot_msg})
        
    messages.append({"role": "user", "content": message})
    
    inputs = tokenizer.apply_chat_template(messages, add_generation_prompt=True, return_tensors="pt")
    streamer = TextIteratorStreamer(tokenizer, skip_prompt=True, skip_special_tokens=True)
    
    generation_kwargs = dict(
        input_ids=inputs,
        streamer=streamer,
        max_new_tokens=512,
        temperature=0.6,
    )
    
    thread = Thread(target=model.generate, kwargs=generation_kwargs)
    thread.start()
    
    partial_text = ""
    for new_text in streamer:
        partial_text += new_text
        yield partial_text

demo = gr.ChatInterface(
    fn=chat_function,
    title="🤖 Cydercoder Qwen 3B AI Chatbot",
    description="Your custom fine-tuned assistant running 24/7 in the cloud for free.",
    examples=["Write a login form using React and Tailwind.", "Fix this code error: Cannot read properties of undefined"]
)

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