app.py

import spaces
import gradio as gr
import torch
from transformers import (
    AutoModelForCausalLM,
    AutoTokenizer,
    BitsAndBytesConfig,
    TextIteratorStreamer,
)
from threading import Thread
from typing import Generator

# ---------------------------------------------------------------------------
# Module-scope model loading - ZeroGPU manages GPU offload transparently
# ---------------------------------------------------------------------------
MODEL_ID = "Qwen/Qwen3-Coder-30B-A3B-Instruct"

quant_config = BitsAndBytesConfig(
    load_in_4bit=True,
    bnb_4bit_quant_type="nf4",
    bnb_4bit_use_double_quant=True,
    bnb_4bit_compute_dtype=torch.bfloat16,
)

tokenizer = AutoTokenizer.from_pretrained(MODEL_ID, trust_remote_code=True)
model = AutoModelForCausalLM.from_pretrained(
    MODEL_ID,
    quantization_config=quant_config,
    device_map="auto",
    torch_dtype=torch.bfloat16,
    trust_remote_code=True,
)
model.eval()

DEFAULT_SYSTEM = "You are an expert coding assistant. Write clean, efficient, well-documented code."


# ---------------------------------------------------------------------------
# ZeroGPU-decorated generation - xlarge for 30B MoE model
# ---------------------------------------------------------------------------
@spaces.GPU(duration=300)
def generate(
    messages: list[dict],
    temperature: float,
    top_p: float,
    max_new_tokens: int,
) -> str:
    """Run model inference inside a ZeroGPU worker process.
    Args are pickled across the process boundary.
    Returns CPU text - safe for unpickling in the main process.
    """
    inputs = tokenizer.apply_chat_template(
        messages,
        tokenize=True,
        add_generation_prompt=True,
        return_tensors="pt",
    ).to(model.device)

    with torch.inference_mode():
        outputs = model.generate(
            inputs,
            max_new_tokens=max_new_tokens,
            temperature=temperature,
            top_p=top_p,
            do_sample=temperature > 0.0,
            pad_token_id=tokenizer.eos_token_id,
        )

    generated = outputs[0][inputs.shape[1]:]
    return tokenizer.decode(generated, skip_special_tokens=True)


# ---------------------------------------------------------------------------
# Streaming variant - yields tokens as they're generated
# ---------------------------------------------------------------------------
@spaces.GPU(duration=300)
def generate_stream(
    messages: list[dict],
    temperature: float,
    top_p: float,
    max_new_tokens: int,
) -> Generator[str, None, None]:
    """Stream tokens from the model one-by-one."""
    inputs = tokenizer.apply_chat_template(
        messages,
        tokenize=True,
        add_generation_prompt=True,
        return_tensors="pt",
    ).to(model.device)

    streamer = TextIteratorStreamer(
        tokenizer,
        skip_prompt=True,
        skip_special_tokens=True,
    )

    generation_kwargs = dict(
        inputs=inputs,
        max_new_tokens=max_new_tokens,
        temperature=temperature,
        top_p=top_p,
        do_sample=temperature > 0.0,
        pad_token_id=tokenizer.eos_token_id,
        streamer=streamer,
    )

    thread = Thread(target=model.generate, kwargs=generation_kwargs)
    thread.start()

    for token in streamer:
        yield token


# ---------------------------------------------------------------------------
# Non-streaming wrapper (for API endpoint)
# ---------------------------------------------------------------------------
def predict(
    message: str,
    history: list,
    system_prompt: str,
    temperature: float,
    top_p: float,
    max_tokens: int,
):
    """Chat function - called both from UI and the auto-generated Gradio API."""
    messages = [{"role": "system", "content": system_prompt}]
    for user_msg, asst_msg in history:
        messages.append({"role": "user", "content": user_msg})
        if asst_msg:
            messages.append({"role": "assistant", "content": asst_msg})
    messages.append({"role": "user", "content": message})

    output = generate(messages, temperature, top_p, max_tokens)
    return output


# ---------------------------------------------------------------------------
# Streaming chat handler
# ---------------------------------------------------------------------------
def chat_fn(
    message: str,
    history: list,
    system_prompt: str,
    temperature: float,
    top_p: float,
    max_tokens: int,
):
    """Generator that yields partial (message, history) tuples for streaming UI."""
    messages = [{"role": "system", "content": system_prompt}]
    for user_msg, asst_msg in history:
        messages.append({"role": "user", "content": user_msg})
        if asst_msg:
            messages.append({"role": "assistant", "content": asst_msg})
    messages.append({"role": "user", "content": message})

    partial = ""
    for token in generate_stream(messages, temperature, top_p, max_tokens):
        partial += token
        yield partial


# ---------------------------------------------------------------------------
# Helpers
# ---------------------------------------------------------------------------
LANGUAGES = ["python", "javascript", "typescript", "rust", "go", "java", "cpp",
             "csharp", "ruby", "php", "sql", "bash", "html", "css", "json", "yaml"]


def build_examples():
    return [
        ["Write a Python async function that downloads a URL and retries 3 times on failure."],
        ["Create a Rust function that reads a CSV file and returns the row count."],
        ["Explain the difference between an interface and a type in TypeScript with examples."],
        ["Write a Go HTTP server that serves static files on port 8080 with CORS support."],
        ["Refactor this Python class to use dependency injection: class Database: ..."],
    ]


# ---------------------------------------------------------------------------
# Gradio UI
# ---------------------------------------------------------------------------
def create_ui():
    with gr.Blocks(
        title="CodeCraft - AI Coding Assistant",
        theme=gr.themes.Soft(
            primary_hue="indigo",
            neutral_hue="slate",
        ),
        fill_width=True,
    ) as demo:
        gr.Markdown(
            "# CodeCraft - AI Coding Assistant\n"
            "Powered by **Qwen3-Coder-30B-A3B-Instruct** (MoE, 3B active) - ZeroGPU xlarge"
        )

        chatbot = gr.Chatbot(
            label="Conversation",
            placeholder="Ask me anything about code...",
            render_markdown=True,
            show_copy_button=True,
            height=500,
        )

        with gr.Row():
            msg = gr.Textbox(
                label="Your message",
                placeholder="Write a Python async function that downloads a URL...",
                scale=8,
                container=False,
            )
            submit_btn = gr.Button("Send", variant="primary", scale=1, min_width=80)
            clear_btn = gr.Button("Clear", scale=1, min_width=80)

        with gr.Accordion("Settings", open=False):
            with gr.Row():
                system_prompt = gr.Textbox(
                    label="System Prompt",
                    value=DEFAULT_SYSTEM,
                    lines=2,
                    scale=3,
                )
                with gr.Column(scale=1):
                    temperature = gr.Slider(
                        label="Temperature", minimum=0.0, maximum=1.5,
                        value=0.3, step=0.05,
                    )
            with gr.Row():
                top_p = gr.Slider(
                    label="Top-P", minimum=0.6, maximum=1.0,
                    value=0.9, step=0.05,
                )
                max_tokens = gr.Slider(
                    label="Max Tokens", minimum=128, maximum=8192,
                    value=2048, step=128,
                )

        gr.Examples(
            examples=build_examples(),
            inputs=[msg],
            label="Try these prompts",
        )

        # -- State: chat history --
        history_state = gr.State([])

        # -- Event wiring --
        def respond(message, history, system, temp, top_p_val, max_tok):
            if not message.strip():
                return "", history, history
            history = history + [(message, None)]
            yield "", history, []
            for partial in chat_fn(message, history[:-1], system, temp, top_p_val, max_tok):
                history[-1] = (message, partial)
                yield "", history, []
            yield "", history, [message]

        msg.submit(
            respond,
            inputs=[msg, history_state, system_prompt, temperature, top_p, max_tokens],
            outputs=[msg, chatbot, history_state],
            concurrency_limit=4,
            api_name="predict",
        )
        submit_btn.click(
            respond,
            inputs=[msg, history_state, system_prompt, temperature, top_p, max_tokens],
            outputs=[msg, chatbot, history_state],
            concurrency_limit=4,
            api_name=False,
        )

        def clear_conversation():
            return [], "", []

        clear_btn.click(
            clear_conversation,
            outputs=[history_state, chatbot, msg],
            concurrency_limit=4,
        )

        gr.Markdown(
            """
            ### API
            This Space exposes a REST API at `/gradio_api/call/predict`.
            See the [Gradio docs](https://www.gradio.app/guides/sharing-your-app#api) for usage.
            """
        )

        return demo


if __name__ == "__main__":
    demo = create_ui()
    demo.queue(default_concurrency_limit=4)
    demo.launch()