app.py

"""
Qwen Coder Battle — Fine-tuned LoRA vs Base Model
===================================================
Compares Md-Asif/qwen-coder-python-ft (a LoRA adapter) against its base
Qwen2.5-Coder-3B model, side by side, on Python coding instructions.

Memory-efficient design for HF ZeroGPU free tier:
  - Only ONE copy of the 3B model is ever loaded.
  - The LoRA adapter is attached via PEFT.
  - "Base model" output is generated with peft_model.disable_adapter(),
    which temporarily turns the LoRA weights off — no second model needed.

If you'd rather run this as two fully separate models, see the
`# ALTERNATE: two full models` comment block near the model loading code.

--------------------------------------------------------------------------
IMPORTANT — prompt format
--------------------------------------------------------------------------
Md-Asif/python-fine-tune trains on FOUR distinct (system_prompt, task_tag)
pairs, keyed by a `task` column: GENERATE, REFACTOR, DEBUG, CODE_REVIEW.
Every training example's user turn is prefixed with "<TASK:XXXX>\n\n".

The LoRA adapter has only ever seen instructions in that exact shape. Feed
it a plain instruction with a different system prompt and no tag, and it's
operating out-of-distribution — often producing WORSE output than the
untouched base model, which was never trained to expect the tag at all.

So: the fine-tuned model gets the task-specific system prompt + tagged
instruction. The base model gets a plain, generic system prompt and the
RAW instruction (no tag) — that's the fair, apples-to-apples comparison,
since neither format is "natural" to the base model anyway.
"""

import ast
import re
import time

import gradio as gr
import spaces
import torch
from peft import PeftModel
from transformers import AutoModelForCausalLM, AutoTokenizer

# ----------------------------------------------------------------------------
# CONFIG — edit these to match how you actually trained the adapter
# ----------------------------------------------------------------------------
# ⚠️ IMPORTANT: this MUST be the exact base model you fine-tuned on top of.
# If your adapter was trained on the Instruct checkpoint, use:
#   "Qwen/Qwen2.5-Coder-3B-Instruct"
# If it was trained on the plain base, use:
#   "Qwen/Qwen2.5-Coder-3B"
# Getting this wrong silently wrecks fine-tuned quality — the LoRA delta
# was computed relative to ONE specific set of base weights.
BASE_MODEL_ID = "Qwen/Qwen2.5-Coder-3B"
ADAPTER_ID = "Md-Asif/qwen-coder-python-ft"

# System prompts, EXACTLY matching what's in the training dataset's
# `messages[0]["content"]` for each task. Keep these byte-identical to the
# dataset — even small wording changes are a mild distribution shift.
TASK_SYSTEM_PROMPTS = {
    "GENERATE": (
        "You are a senior Python engineer.\n"
        "Write complete, correct, production-ready Python code.\n"
        "Include type hints, docstrings, and handle edge cases.\n"
        "No placeholders. No TODOs. Only working code."
    ),
    "REFACTOR": (
        "You are an expert Python refactoring assistant.\n"
        "Rewrite the code following:\n"
        "- PEP 8\n"
        "- Type hints\n"
        "- Docstrings\n"
        "- Best practices\n"
        "- SOLID principles\n\n"
        "Do not change behavior. Only improve structure and quality."
    ),
    "DEBUG": (
        "You are an expert Python debugging engineer.\n"
        "Locate the exact bug.\n"
        "Explain precisely why it happens.\n"
        "Produce corrected, complete Python code."
    ),
    "CODE_REVIEW": (
        "You are a senior Python code reviewer.\n"
        "Review code like a staff engineer.\n\n"
        "Find and fix:\n"
        "- Syntax bugs\n"
        "- Runtime bugs\n"
        "- Logic bugs\n"
        "- Performance issues\n"
        "- Security issues\n"
        "- Maintainability issues\n"
        "Produce corrected production-quality Python code and explain every change."
    ),
}

# What the BASE model gets instead — a neutral prompt with no dataset-specific
# tagging, since the base model was never trained on the <TASK:X> scheme.
BASE_SYSTEM_PROMPT = (
    "You are an expert Python programmer. Given an instruction, respond with "
    "clean, correct, well-documented Python code. Wrap code in a ```python "
    "fenced block. Keep explanations brief."
)

MAX_NEW_TOKENS_DEFAULT = 512
WARMUP_TOKENS = 8  # tiny generation to absorb CUDA cold-start before timing
DTYPE = torch.bfloat16

# ----------------------------------------------------------------------------
# MODEL LOADING (runs once, at Space startup, on the CPU host)
# ZeroGPU intercepts `.to("cuda")` / `device_map="auto"` so this is safe to
# write as if a GPU is always present — the real GPU is attached only inside
# functions decorated with @spaces.GPU.
# ----------------------------------------------------------------------------
print("Loading tokenizer...")
try:
    tokenizer = AutoTokenizer.from_pretrained(ADAPTER_ID)
except Exception:
    tokenizer = AutoTokenizer.from_pretrained(BASE_MODEL_ID)

if tokenizer.pad_token_id is None:
    tokenizer.pad_token = tokenizer.eos_token

print("Loading base model (on CPU first — do NOT use device_map='auto' here;")
print("it forces real CUDA writes before ZeroGPU has attached a GPU)...")
base_model = AutoModelForCausalLM.from_pretrained(
    BASE_MODEL_ID,
    dtype=DTYPE,
)

print("Attaching LoRA adapter (also on CPU)...")
# torch_device="cpu" is required here: PEFT's internal infer_device() checks
# torch.cuda.is_available(), which ZeroGPU's patching always reports as True
# (even outside a GPU-attached window). Without this, PEFT tries to load the
# adapter's safetensors directly onto a CUDA device that isn't physically
# attached yet at import time, which crashes with "No CUDA GPUs are available".
model = PeftModel.from_pretrained(base_model, ADAPTER_ID, torch_device="cpu")

print("Moving model to CUDA (ZeroGPU defers this safely until a request comes in)...")
model.to("cuda")
model.eval()

# ALTERNATE: two full models (uses ~2x VRAM, only do this if disable_adapter()
# gives you trouble, e.g. adapter changed the tokenizer/embedding size):
#
# base_model = AutoModelForCausalLM.from_pretrained(BASE_MODEL_ID, dtype=DTYPE).to("cuda").eval()
# ft_model = PeftModel.from_pretrained(
#     AutoModelForCausalLM.from_pretrained(BASE_MODEL_ID, dtype=DTYPE),
#     ADAPTER_ID,
# ).to("cuda").eval()
# then generate with base_model and ft_model directly, no disable_adapter() needed.


# ----------------------------------------------------------------------------
# GENERATION
# ----------------------------------------------------------------------------
def _build_inputs(system_prompt: str, user_content: str):
    messages = [
        {"role": "system", "content": system_prompt},
        {"role": "user", "content": user_content},
    ]
    text = tokenizer.apply_chat_template(
        messages, tokenize=False, add_generation_prompt=True
    )
    return tokenizer(text, return_tensors="pt").to(model.device)


def build_inputs_ft(instruction: str, task: str):
    """Fine-tuned path: task-specific system prompt + <TASK:X> tagged instruction,
    matching the training data format exactly."""
    system_prompt = TASK_SYSTEM_PROMPTS[task]
    tagged_instruction = f"<TASK:{task}>\n\n{instruction}"
    return _build_inputs(system_prompt, tagged_instruction)


def build_inputs_base(instruction: str):
    """Base path: neutral system prompt, raw instruction, no tag."""
    return _build_inputs(BASE_SYSTEM_PROMPT, instruction)


def _run_generate(inputs, max_new_tokens: int, temperature: float, use_sampling: bool):
    gen_kwargs = dict(
        **inputs,
        max_new_tokens=max_new_tokens,
        do_sample=use_sampling,
        pad_token_id=tokenizer.pad_token_id,
    )
    if use_sampling:
        gen_kwargs["temperature"] = max(temperature, 0.01)
        gen_kwargs["top_p"] = 0.9

    start = time.perf_counter()
    with torch.no_grad():
        output_ids = model.generate(**gen_kwargs)
    elapsed = time.perf_counter() - start

    new_tokens = output_ids[0][inputs["input_ids"].shape[1]:]
    text = tokenizer.decode(new_tokens, skip_special_tokens=True)
    return text, elapsed, len(new_tokens)


def _warmup():
    """Absorb CUDA context attach / kernel autotune cost with a throwaway
    generation, so neither the FT nor the base timing eats the cold-start
    penalty. Must run once per @spaces.GPU-decorated call, before any timed
    generation."""
    dummy = _build_inputs(BASE_SYSTEM_PROMPT, "print hello world")
    with torch.no_grad():
        model.generate(
            **dummy,
            max_new_tokens=WARMUP_TOKENS,
            do_sample=False,
            pad_token_id=tokenizer.pad_token_id,
        )


@spaces.GPU(duration=90)
def run_battle(instruction: str, task: str, max_new_tokens: int, temperature: float, use_sampling: bool):
    if not instruction or not instruction.strip():
        raise gr.Error("Please enter an instruction first.")

    _warmup()

    # Fine-tuned model (adapter active, task-specific prompt + tag)
    ft_inputs = build_inputs_ft(instruction, task)
    ft_text, ft_time, ft_tokens = _run_generate(ft_inputs, max_new_tokens, temperature, use_sampling)

    # Base model (adapter disabled — same weights, no reload — neutral prompt)
    base_inputs = build_inputs_base(instruction)
    with model.disable_adapter():
        base_text, base_time, base_tokens = _run_generate(base_inputs, max_new_tokens, temperature, use_sampling)

    ft_metrics = analyze_response(ft_text, ft_time, ft_tokens)
    base_metrics = analyze_response(base_text, base_time, base_tokens)
    return ft_text, base_text, ft_metrics, base_metrics


# ----------------------------------------------------------------------------
# HEURISTIC CODE QUALITY SCORING
# This is NOT a correctness guarantee. It's a fast, transparent proxy based
# on static analysis. For true accuracy, you'd need to execute the code
# against real test cases (e.g. HumanEval-style pass@1) — see the note in
# the README for how to extend this.
# ----------------------------------------------------------------------------
CODE_BLOCK_RE = re.compile(r"```(?:python)?\s*\n(.*?)```", re.DOTALL)


def extract_code(text: str) -> str:
    match = CODE_BLOCK_RE.search(text)
    if match:
        return match.group(1).strip()
    # No fenced block found — fall back to the raw text, it may still be code
    return text.strip()


def analyze_response(text: str, elapsed: float, n_tokens: int) -> dict:
    code = extract_code(text)
    metrics = {
        "generation_time_s": round(elapsed, 2),
        "tokens_generated": n_tokens,
        "tokens_per_sec": round(n_tokens / elapsed, 2) if elapsed > 0 else 0.0,
        "syntax_valid": False,
        "num_functions": 0,
        "num_classes": 0,
        "has_docstring": False,
        "lines_of_code": len([l for l in code.splitlines() if l.strip()]),
        "raw_text": text,
    }
    try:
        tree = ast.parse(code)
        metrics["syntax_valid"] = True
        metrics["num_functions"] = sum(isinstance(n, ast.FunctionDef) for n in ast.walk(tree))
        metrics["num_classes"] = sum(isinstance(n, ast.ClassDef) for n in ast.walk(tree))
        metrics["has_docstring"] = any(
            ast.get_docstring(n) for n in ast.walk(tree)
            if isinstance(n, (ast.Module, ast.FunctionDef, ast.ClassDef))
        )
    except SyntaxError:
        pass

    # Composite heuristic score out of 100. Weighting is deliberately
    # syntax-heavy since broken code is disqualifying no matter what else it has.
    score = 0
    score += 50 if metrics["syntax_valid"] else 0
    score += 15 if metrics["has_docstring"] else 0
    score += 15 if metrics["num_functions"] + metrics["num_classes"] > 0 else 0
    # Reward reasonable conciseness — penalize extreme verbosity beyond ~120 lines
    score += 10 if 1 <= metrics["lines_of_code"] <= 120 else 5 if metrics["lines_of_code"] > 0 else 0
    # Small speed bonus, capped, so a slow-but-correct answer isn't punished hard
    score += min(10, metrics["tokens_per_sec"] / 5)
    metrics["quality_score"] = round(min(score, 100), 1)
    return metrics


# ----------------------------------------------------------------------------
# UI RENDERING HELPERS
# ----------------------------------------------------------------------------
def metric_row(label, ft_val, base_val, higher_is_better=True, suffix=""):
    try:
        ft_num, base_num = float(ft_val), float(base_val)
        ft_win = ft_num > base_num if higher_is_better else ft_num < base_num
        base_win = base_num > ft_num if higher_is_better else base_num < ft_num
    except (TypeError, ValueError):
        ft_win = base_win = False
    ft_cls = "metric-win" if ft_win else ""
    base_cls = "metric-win" if base_win else ""
    return f"""
    <tr>
      <td class="metric-label">{label}</td>
      <td class="{ft_cls}">{ft_val}{suffix}</td>
      <td class="{base_cls}">{base_val}{suffix}</td>
    </tr>"""


def render_comparison(ft_metrics: dict, base_metrics: dict) -> str:
    winner = "Fine-tuned" if ft_metrics["quality_score"] >= base_metrics["quality_score"] else "Base"
    winner_badge = f'<div class="winner-badge">🏆 {winner} model scored higher</div>'
    rows = "".join([
        metric_row("Quality score", ft_metrics["quality_score"], base_metrics["quality_score"], suffix="/100"),
        metric_row("Syntax valid", "✅ Yes" if ft_metrics["syntax_valid"] else "❌ No",
                   "✅ Yes" if base_metrics["syntax_valid"] else "❌ No", higher_is_better=None),
        metric_row("Functions / classes defined", ft_metrics["num_functions"] + ft_metrics["num_classes"],
                   base_metrics["num_functions"] + base_metrics["num_classes"]),
        metric_row("Has docstring", "✅ Yes" if ft_metrics["has_docstring"] else "❌ No",
                   "✅ Yes" if base_metrics["has_docstring"] else "❌ No", higher_is_better=None),
        metric_row("Lines of code", ft_metrics["lines_of_code"], base_metrics["lines_of_code"], higher_is_better=None),
        metric_row("Generation time", ft_metrics["generation_time_s"], base_metrics["generation_time_s"],
                   higher_is_better=False, suffix="s"),
        metric_row("Tokens / sec", ft_metrics["tokens_per_sec"], base_metrics["tokens_per_sec"], suffix=" tok/s"),
    ])
    return f"""
    <div class="comparison-wrap">
      {winner_badge}
      <table class="comparison-table">
        <thead>
          <tr><th></th><th>🟣 Fine-tuned</th><th>⚪ Base</th></tr>
        </thead>
        <tbody>{rows}</tbody>
      </table>
      <div class="disclaimer">
        Quality score is a static-analysis heuristic (syntax validity, structure,
        docstrings, speed) — not a guarantee of runtime correctness. Fine-tuned
        gets its trained system prompt + &lt;TASK:{{task}}&gt; tag; base gets a
        neutral prompt with no tag — matching how each was actually trained.
      </div>
    </div>"""


# ----------------------------------------------------------------------------
# CUSTOM STYLING
# ----------------------------------------------------------------------------
CUSTOM_CSS = """
:root {
    --accent: #7c3aed;
    --accent-2: #06b6d4;
}
.gradio-container { max-width: 1200px !important; margin: auto; }
#hero {
    background: linear-gradient(135deg, #1e1b4b 0%, #4c1d95 45%, #0e7490 100%);
    border-radius: 18px;
    padding: 28px 32px;
    margin-bottom: 18px;
    color: white;
}
#hero h1 { margin: 0 0 6px 0; font-size: 1.7rem; }
#hero p { margin: 0; opacity: 0.85; font-size: 0.95rem; }
.model-card {
    border-radius: 14px;
    padding: 4px;
    border: 1px solid rgba(124,58,237,0.25);
}
.model-card-label {
    font-weight: 700;
    font-size: 0.85rem;
    letter-spacing: 0.02em;
    text-transform: uppercase;
    padding: 6px 12px;
    border-radius: 999px;
    display: inline-block;
    margin-bottom: 8px;
}
.label-ft { background: rgba(124,58,237,0.15); color: #a78bfa; }
.label-base { background: rgba(148,163,184,0.15); color: #94a3b8; }
.comparison-wrap { padding: 6px 4px; }
.winner-badge {
    text-align: center;
    font-weight: 700;
    padding: 10px;
    border-radius: 10px;
    background: linear-gradient(135deg, rgba(124,58,237,0.15), rgba(6,182,212,0.15));
    margin-bottom: 14px;
}
.comparison-table { width: 100%; border-collapse: collapse; font-size: 0.92rem; }
.comparison-table th { text-align: left; padding: 8px 10px; opacity: 0.7; font-weight: 600; }
.comparison-table td { padding: 8px 10px; border-top: 1px solid rgba(148,163,184,0.15); }
.metric-label { opacity: 0.75; }
.metric-win { font-weight: 700; color: #22c55e; }
.disclaimer { margin-top: 12px; font-size: 0.78rem; opacity: 0.55; font-style: italic; }
"""

# ----------------------------------------------------------------------------
# GRADIO APP
# ----------------------------------------------------------------------------
EXAMPLES = [
    ["GENERATE", "Write a function that checks if a string is a valid palindrome, ignoring punctuation and case."],
    ["GENERATE", "Implement an LRU cache from scratch without using functools.lru_cache."],
    ["GENERATE", "Write a script that reads a CSV file and returns the top 5 rows by a given column, using pandas."],
    ["GENERATE", "Create a decorator that retries a function up to N times with exponential backoff."],
    ["DEBUG", "def add(a, b):\n    return a + b\n\nprint(add(2))\n\n# This raises a TypeError. Find and fix the bug."],
    ["REFACTOR", "def f(x):\n    y=[]\n    for i in range(len(x)):\n        if x[i]%2==0:\n            y.append(x[i])\n    return y"],
]

with gr.Blocks(css=CUSTOM_CSS, theme=gr.themes.Soft(primary_hue="violet", secondary_hue="cyan")) as demo:
    gr.HTML("""
    <div id="hero">
      <h1>⚡ Qwen Coder Battle</h1>
      <p>Md-Asif/qwen-coder-python-ft (LoRA fine-tune) vs Qwen2.5-Coder-3B (base) — each prompted the way it was actually trained. Judge for yourself.</p>
    </div>
    """)

    with gr.Row():
        instruction = gr.Textbox(
            label="Coding instruction",
            placeholder="e.g. Write a function that merges two sorted lists in O(n) time.",
            lines=3,
            scale=3,
        )
        task = gr.Dropdown(
            choices=["GENERATE", "REFACTOR", "DEBUG", "CODE_REVIEW"],
            value="GENERATE",
            label="Task type",
            info="Must match the dataset's <TASK:X> categories for a fair fine-tuned comparison.",
            scale=1,
        )

    with gr.Row():
        run_btn = gr.Button("⚔️  Run the battle", variant="primary", scale=3)
        with gr.Column(scale=2):
            with gr.Accordion("Generation settings", open=False):
                max_tokens = gr.Slider(64, 1024, value=MAX_NEW_TOKENS_DEFAULT, step=32, label="Max new tokens")
                use_sampling = gr.Checkbox(value=False, label="Use sampling (off = greedy, fairest for comparison)")
                temperature = gr.Slider(0.1, 1.5, value=0.7, step=0.1, label="Temperature (only if sampling is on)")

    gr.Examples(examples=EXAMPLES, inputs=[task, instruction])

    with gr.Row():
        with gr.Column():
            gr.HTML('<span class="model-card-label label-ft">🟣 Fine-tuned</span>')
            ft_output = gr.Code(language="python", label=None, lines=18)
        with gr.Column():
            gr.HTML('<span class="model-card-label label-base">⚪ Base model</span>')
            base_output = gr.Code(language="python", label=None, lines=18)

    comparison_html = gr.HTML()
    ft_metrics_state = gr.State()
    base_metrics_state = gr.State()

    def on_run(instr, task_choice, max_tok, temp, sample):
        ft_text, base_text, ft_m, base_m = run_battle(instr, task_choice, max_tok, temp, sample)
        return ft_text, base_text, render_comparison(ft_m, base_m), ft_m, base_m

    run_btn.click(
        on_run,
        inputs=[instruction, task, max_tokens, temperature, use_sampling],
        outputs=[ft_output, base_output, comparison_html, ft_metrics_state, base_metrics_state],
    )

if __name__ == "__main__":
    demo.queue(max_size=20).launch()