app.py

# ================================================================
# ✨ UltraThinker-Coder-3B — HF Zero GPU Chat
# Developed by Malik Ayaan Ahmed
# ================================================================

import os
import sys
import gc
import threading
import torch
import gradio as gr
import spaces
from transformers import AutoModelForCausalLM, AutoTokenizer, TextIteratorStreamer
from peft import PeftModel

# ── 0. Suppress Python 3.10 Asyncio GC Log Spam ─────────────────
# This intercepts and silently ignores the benign "Invalid file descriptor" error
def custom_unraisablehook(unraisable):
    if unraisable.exc_type == ValueError and "Invalid file descriptor" in str(unraisable.exc_value):
        return
    sys.__unraisablehook__(unraisable)
sys.unraisablehook = custom_unraisablehook

# ── 1. Setup & Configuration ────────────────────────────────────
HF_TOKEN = os.environ.get("HF_TOKEN", None) 
ADAPTER_ID = "AyaanAhmed123/UltraThinker-Coder-3B"
BASE_MODEL_ID = "Qwen/Qwen2.5-Coder-3B" 

print("⚙️ Loading tokenizer...")
tokenizer = AutoTokenizer.from_pretrained(
    ADAPTER_ID, token=HF_TOKEN, trust_remote_code=True
)
if tokenizer.pad_token is None:
    tokenizer.pad_token = tokenizer.eos_token
tokenizer.padding_side = "right"

# ── 2. Model Loading (Base Model loads to CPU safely) ───────────
print(f"💾 Loading pristine base model ({BASE_MODEL_ID}) to CPU RAM...")
base_model = AutoModelForCausalLM.from_pretrained(
    BASE_MODEL_ID,
    torch_dtype=torch.bfloat16, 
    device_map="cpu",           
    low_cpu_mem_usage=True,
    token=HF_TOKEN,
    trust_remote_code=True,
)

is_adapter_loaded = False 

# ── 3. Stop tokens & Strict Termination ─────────────────────────
IM_START = "<|im_start|>"
IM_END   = "<|im_end|>"
_stop_ids = []

for s in [IM_END, "<|eot_id|>", "</s>", tokenizer.eos_token, "<|endoftext|>"]:
    if s:
        sid = tokenizer.convert_tokens_to_ids(s)
        if sid and sid != tokenizer.unk_token_id:
            _stop_ids.append(sid)
            
if tokenizer.eos_token_id:
    _stop_ids.append(tokenizer.eos_token_id)
_stop_ids = list(set(_stop_ids))

# ── 4. System Prompt ────────────────────────────────────────────
SYSTEM = (
    "You are UltraThinker-Coder-3B, an elite AI by Malik Ayaan Ahmed. "
    "MANDATORY: You MUST start with <think>. Do all your internal reasoning, step-by-step logic, "
    "and SYNTAX VERIFICATION inside the <think> and </think> tags. Check your code for errors before writing it. "
    "Close your thought process with </think> BEFORE outputting the final response. "
    "In your final response, use emojis intelligently to structure your answer. "
    "MANDATORY: Terminate your response immediately after answering. Do not drift into unrelated topics or output raw training data tags."
)

def build_prompt(history, message: str) -> str:
    parts = [f"{IM_START}system\n{SYSTEM}{IM_END}\n"]
    for turn in history:
        if isinstance(turn, (list, tuple)) and len(turn) == 2:
            u, b = turn
            if u: parts.append(f"{IM_START}user\n{u}{IM_END}\n")
            if b: parts.append(f"{IM_START}assistant\n{b}{IM_END}\n")
        elif isinstance(turn, dict):
            role = turn.get("role", "user")
            content = turn.get("content", "")
            if isinstance(content, list):
                content = " ".join(x.get("text", "") if isinstance(x, dict) else str(x) for x in content)
            parts.append(f"{IM_START}{role}\n{content}{IM_END}\n")
    parts.append(f"{IM_START}user\n{message}{IM_END}\n")
    parts.append(f"{IM_START}assistant\n")
    return "".join(parts)

# ── 5. Generation (Zero GPU Decorator Applied) ──────────────────
@spaces.GPU(duration=120)
def generate_stream(message: str, history):
    global base_model, is_adapter_loaded
    
    # Send base model to the newly acquired GPU
    base_model.to("cuda")
    
    # Load & merge the adapter ONLY when the GPU is successfully active
    if not is_adapter_loaded:
        print("🔗 GPU acquired! Merging UltraThinker-Coder-3B Adapter directly in VRAM...")
        base_model = PeftModel.from_pretrained(base_model, ADAPTER_ID, token=HF_TOKEN)
        base_model = base_model.merge_and_unload()
        base_model.eval()
        is_adapter_loaded = True
        print("✅ Merge complete! Ready to generate.")
    
    prompt = build_prompt(history, message)
    inputs = tokenizer(prompt, return_tensors="pt", truncation=True, max_length=4096).to("cuda")
    
    input_length = inputs["input_ids"].shape[1]
    max_dynamic_tokens = max(512, 4096 - input_length)
    
    streamer = TextIteratorStreamer(
        tokenizer, skip_prompt=True, skip_special_tokens=False, timeout=120
    )
    
    gen_kwargs = dict(
        **inputs,
        streamer=streamer,
        max_new_tokens=max_dynamic_tokens,
        do_sample=True,
        temperature=0.4, 
        top_p=0.9,
        repetition_penalty=1.1,
        use_cache=True,
        eos_token_id=_stop_ids,
        pad_token_id=(tokenizer.pad_token_id or tokenizer.eos_token_id),
    )
    
    thread = threading.Thread(target=base_model.generate, kwargs=gen_kwargs)
    thread.start()
    
    stop_markers = [
        IM_END, "<|eot_id|>", "</s>", tokenizer.eos_token, 
        IM_START, "<|im_start|>", "\nuser\n", "User:",
        "<|endoftext|>", "<|fim_prefix|>", "<|fim_middle|>", "<|file_sep|>", "</output>", "<output>"
    ]
    
    full = ""
    for token in streamer:
        full += token
        
        hit_stop = False
        for marker in stop_markers:
            if marker and marker in full:
                full = full.split(marker)[0]
                hit_stop = True
                break
        
        yield full
        if hit_stop:
            break

# ── 6. UI & CSS ─────────────────────────────────────────────────
CSS = """
@import url('https://fonts.googleapis.com/css2?family=Inter:wght@400;600&display=swap');
* { font-family: 'Inter', sans-serif !important; }

/* User & Bot Message Styling */
.message-wrap .message.user {
    background: #F3F4F6 !important; color: #111827 !important;
    border-radius: 18px 18px 4px 18px !important;
    border: none !important;
    box-shadow: 0 4px 6px -1px rgba(0,0,0,.1) !important;
}
.message-wrap .message.bot { background: transparent !important; border: none !important; box-shadow: none !important; }

/* Elegant Thinking Dropdown */
details {
    background: #f8fafc !important;
    border: 1px solid #e2e8f0 !important;
    border-radius: 12px !important;
    padding: 10px 16px !important;
    margin-bottom: 16px !important;
    display: block !important;
    width: 100% !important;
}
details summary { font-weight: 600 !important; cursor: pointer !important; color: #4f46e5 !important; outline: none !important; }

/* Pitch Black Code Blocks */
.prose pre, .message-wrap .message.bot pre {
    background-color: #000000 !important;
    border: 1px solid #30363D !important;
    border-radius: 12px !important;
    padding: 16px !important; 
    padding-top: 40px !important; 
    overflow-x: auto !important;
    position: relative !important;
}
.prose pre code { background: transparent !important; color: #e5e5e5 !important; }

/* Ensure Buttons are Visible */
.gradio-container [aria-label="Submit"],
.gradio-container [aria-label="Stop generation"],
.gradio-container button.submit-btn, 
.gradio-container button.stop-btn {
    display: flex !important;
    opacity: 1 !important;
    visibility: visible !important;
    pointer-events: auto !important;
    z-index: 100 !important;
}
"""

theme = gr.themes.Soft(primary_hue="indigo", secondary_hue="blue")

# Fixed for Gradio 6.0: CSS and theme removed from Blocks constructor
with gr.Blocks() as demo:
    gr.ChatInterface(
        fn=generate_stream,
        title="✨ UltraThinker-Coder-3B",
        description="Developed by Malik Ayaan Ahmed. Deployed on HF Zero GPU.",
        chatbot=gr.Chatbot(
            height=600,
            render_markdown=True,
            reasoning_tags=[("<think>", "</think>")],
        )
    )

if __name__ == "__main__":
    # Fixed for Gradio 6.0: CSS and theme passed to launch()
    demo.queue().launch(css=CSS, theme=theme)