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

# 1. Point the script directly to your uploaded Hugging Face model
MODEL_ID = "Cydercoder/qwen2.5-coder-3b"

print("Loading tokenizer...")
tokenizer = AutoTokenizer.from_pretrained(MODEL_ID)

print("Loading model on CPU via standard clean mapping...")
# Forcing float32 without low_cpu_mem_usage prevents the Transformers v5 thread crash
model = AutoModelForCausalLM.from_pretrained(
    MODEL_ID,
    torch_dtype=torch.float32,
    device_map="cpu"
)

def chat_function(message, history):
    # Construct formatting conversation list matrices
    messages = [
        {"role": "system", "content": "You are an expert full-stack developer assistant."}
    ]
    
    # Re-insert existing browser chat history logs
    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})
    
    # Process tokens safely
    inputs = tokenizer.apply_chat_template(messages, add_generation_prompt=True, return_tensors="pt")
    
    # Set up a dynamic background streamer so answers appear word-by-word in browser
    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,
    )
    
    # Run text generation in a separate background processor thread
    thread = Thread(target=model.generate, kwargs=generation_kwargs)
    thread.start()
    
    partial_text = ""
    for new_text in streamer:
        partial_text += new_text
        yield partial_text

# 4. Initialize the native Gradio browser interface
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()