Create app.py

app.py

@@ -0,0 +1,411 @@
+"""
+Python Code Review Tool — Hugging Face Spaces
+LLM: google/gemma-2b-it (local, no token restrictions)
+Hardware: CPU-optimised with 4-bit quantisation
+Architecture: UI → Validation → Pre-processing → Generation → Output
+"""
+
+import ast
+import re
+import textwrap
+import threading
+from dataclasses import dataclass, field
+
+import gradio as gr
+import torch
+from transformers import (
+    AutoTokenizer,
+    AutoModelForCausalLM,
+    BitsAndBytesConfig,
+    TextIteratorStreamer,
+)
+
+# ─────────────────────────────────────────────────────────────
+# MODEL LOADING  (runs once at startup)
+# ─────────────────────────────────────────────────────────────
+
+MODEL_ID = "google/gemma-2b-it"
+
+print("Loading tokeniser…")
+tokenizer = AutoTokenizer.from_pretrained(MODEL_ID)
+
+# 4-bit quant keeps RAM under ~3 GB — safe on HF free CPU
+bnb_config = BitsAndBytesConfig(
+    load_in_4bit=True,
+    bnb_4bit_compute_dtype=torch.float32,   # float32 for CPU
+    bnb_4bit_use_double_quant=True,
+    bnb_4bit_quant_type="nf4",
+)
+
+print("Loading model (4-bit)…")
+model = AutoModelForCausalLM.from_pretrained(
+    MODEL_ID,
+    quantization_config=bnb_config,
+    device_map="cpu",
+    low_cpu_mem_usage=True,
+)
+model.eval()
+print("Model ready ✓")
+
+
+# ─────────────────────────────────────────────────────────────
+# VALIDATION & FLOW MANAGEMENT MODULE
+# ─────────────────────────────────────────────────────────────
+
+@dataclass
+class ValidationResult:
+    valid: bool
+    error: str = ""
+    warnings: list = field(default_factory=list)
+
+
+def validate_input(code: str, description: str) -> ValidationResult:
+    warnings = []
+
+    if not code.strip():
+        return ValidationResult(False, "❌ Code input is empty.")
+    if not description.strip():
+        return ValidationResult(False, "❌ Description is empty.")
+    if len(code.strip()) < 10:
+        return ValidationResult(False, "❌ Code too short to review.")
+    if len(description.strip()) < 5:
+        return ValidationResult(False, "❌ Description too brief.")
+    if len(code) > 30_000:
+        return ValidationResult(False, "❌ Code exceeds 30,000 characters. Split into smaller sections.")
+
+    python_hints = ["def ", "class ", "import ", "for ", "while ", "if ", "print", "return", "="]
+    if not any(kw in code for kw in python_hints):
+        warnings.append("⚠️ Doesn't look like Python — proceeding anyway.")
+
+    try:
+        ast.parse(code)
+    except SyntaxError as e:
+        warnings.append(f"⚠️ Syntax error at line {e.lineno}: {e.msg} — review will still run.")
+
+    return ValidationResult(True, warnings=warnings)
+
+
+# ─────────────────────────────────────────────────────────────
+# PRE-PROCESSING MODULE
+# ─────────────────────────────────────────────────────────────
+
+def estimate_complexity(code: str) -> str:
+    lines   = [l for l in code.splitlines() if l.strip() and not l.strip().startswith("#")]
+    funcs   = len(re.findall(r"^\s*def ",   code, re.MULTILINE))
+    classes = len(re.findall(r"^\s*class ", code, re.MULTILINE))
+    if len(lines) < 20 and funcs <= 1 and classes == 0:
+        return "beginner"
+    elif len(lines) < 100 and classes <= 2:
+        return "intermediate"
+    return "advanced"
+
+
+def extract_entities(code: str) -> str:
+    try:
+        tree = ast.parse(code)
+        entities = []
+        for node in ast.walk(tree):
+            if isinstance(node, (ast.FunctionDef, ast.AsyncFunctionDef)):
+                entities.append(f"function `{node.name}`")
+            elif isinstance(node, ast.ClassDef):
+                entities.append(f"class `{node.name}`")
+        return ", ".join(entities[:10]) if entities else "none detected"
+    except Exception:
+        return "parse skipped (syntax errors present)"
+
+
+def build_prompt(code: str, description: str, level: str) -> str:
+    """
+    Gemma-IT chat format:
+      <start_of_turn>user\\n…<end_of_turn>\\n<start_of_turn>model\\n
+    """
+    entities = extract_entities(code)
+    dedented = textwrap.dedent(code).strip()
+
+    user_msg = f"""You are an expert Python code reviewer. Review the following {level}-level Python code carefully.
+
+## What the code should do
+{description.strip()}
+
+## Detected entities
+{entities}
+
+## Code
+```python
+{dedented}
+```
+
+Write a detailed markdown review covering every section below.
+
+### 1. 🎯 Description Alignment
+Does the code do what the description says? List any gaps or mismatches.
+
+### 2. 🐛 Bugs & Correctness
+List every bug with line references and corrected code snippets.
+
+### 3. ⚡ Performance
+Identify inefficiencies. Suggest faster alternatives with complexity notes.
+
+### 4. 🏗️ Code Quality & Pythonic Style
+PEP 8 compliance, naming, readability, idiomatic Python patterns.
+
+### 5. 🔒 Edge Cases & Robustness
+What inputs could break this? Add guard clauses or error handling.
+
+### 6. ✅ Top 3 Quick Wins
+The three highest-impact improvements to make first."""
+
+    return (
+        f"<start_of_turn>user\n{user_msg}<end_of_turn>\n"
+        "<start_of_turn>model\n"
+    )
+
+
+# ─────────────────────────────────────────────────────────────
+# GENERATION MODULE  — no max_new_tokens cap
+# ─────────────────────────────────────────────────────────────
+
+def stream_review(prompt: str):
+    """Yields token chunks. Generation stops at natural EOS — no token limit."""
+    inputs   = tokenizer(prompt, return_tensors="pt")
+    input_ids = inputs["input_ids"]   # CPU tensor
+
+    streamer = TextIteratorStreamer(
+        tokenizer,
+        skip_prompt=True,
+        skip_special_tokens=True,
+    )
+
+    gen_kwargs = dict(
+        input_ids=input_ids,
+        streamer=streamer,
+        do_sample=True,
+        temperature=0.3,        # focused, consistent output
+        top_p=0.9,
+        repetition_penalty=1.15,
+        # ← No max_new_tokens: model runs until EOS
+    )
+
+    thread = threading.Thread(target=model.generate, kwargs=gen_kwargs)
+    thread.start()
+
+    for token in streamer:
+        yield token
+
+    thread.join()
+
+
+# ─────────────────────────────────────────────────────────────
+# OUTPUT MODULE
+# ─────────────────────────────────────────────────────────────
+
+def format_output(raw: str, warnings: list) -> str:
+    warning_block = ""
+    if warnings:
+        warning_block = "\n".join(f"> {w}" for w in warnings) + "\n\n---\n\n"
+    cleaned = raw.strip().removesuffix("<end_of_turn>").strip()
+    return warning_block + cleaned
+
+
+# ─────────────────────────────────────────────────────────────
+# MAIN PIPELINE
+# ─────────────────────────────────────────────────────────────
+
+def review_pipeline(code: str, description: str):
+    # 1. Validate
+    result = validate_input(code, description)
+    if not result.valid:
+        yield result.error
+        return
+
+    header = ""
+    if result.warnings:
+        header = "\n".join(result.warnings) + "\n\n---\n\n"
+
+    yield header + "*⏳ Generating review with Gemma 2B — ~1–3 min on CPU. Streaming token by token…*\n\n"
+
+    # 2. Pre-process
+    level  = estimate_complexity(code)
+    prompt = build_prompt(code, description, level)
+
+    # 3. Stream generation
+    accumulated = ""
+    try:
+        for chunk in stream_review(prompt):
+            accumulated += chunk
+            yield header + accumulated
+    except Exception as e:
+        yield header + accumulated + f"\n\n❌ Generation error: {e}"
+        return
+
+    # 4. Final clean output
+    yield format_output(header + accumulated, [])
+
+
+# ─────────────────────────────────────────────────────────────
+# UI MODULE
+# ─────────────────────────────────────────────────────────────
+
+EXAMPLES = [
+    [
+        "def find_duplicates(lst):\n    duplicates = []\n    for i in range(len(lst)):\n        for j in range(len(lst)):\n            if i != j and lst[i] == lst[j]:\n                if lst[i] not in duplicates:\n                    duplicates.append(lst[i])\n    return duplicates",
+        "Find all duplicate values in a list",
+    ],
+    [
+        "import requests\n\ndef fetch_user(user_id):\n    r = requests.get(f'https://api.example.com/users/{user_id}')\n    data = r.json()\n    return data['name'], data['email']",
+        "Fetch a user's name and email from a REST API",
+    ],
+    [
+        "class BankAccount:\n    def __init__(self, owner, balance=0):\n        self.owner = owner\n        self.balance = balance\n\n    def deposit(self, amount):\n        self.balance += amount\n\n    def withdraw(self, amount):\n        self.balance -= amount\n\n    def get_balance(self):\n        return self.balance",
+        "A simple bank account class with deposit and withdraw",
+    ],
+    [
+        "def quicksort(arr):\n    if len(arr) <= 1:\n        return arr\n    pivot = arr[0]\n    left = [x for x in arr[1:] if x <= pivot]\n    right = [x for x in arr[1:] if x > pivot]\n    return quicksort(left) + [pivot] + quicksort(right)",
+        "Quicksort implementation",
+    ],
+]
+
+CSS = """
+@import url('https://fonts.googleapis.com/css2?family=DM+Mono:wght@300;400;500&family=Syne:wght@400;600;700;800&display=swap');
+
+:root {
+    --bg: #0d0d0f;
+    --surface: #141418;
+    --surface2: #1c1c22;
+    --border: #2a2a35;
+    --accent: #7ee8a2;
+    --accent2: #4f8ef7;
+    --text: #e8e8f0;
+    --muted: #7070a0;
+}
+
+body, .gradio-container {
+    background: var(--bg) !important;
+    font-family: 'Syne', sans-serif !important;
+    color: var(--text) !important;
+}
+.gradio-container { max-width: 1300px !important; margin: 0 auto !important; }
+
+.app-header {
+    text-align: center; padding: 44px 24px 28px;
+    border-bottom: 1px solid var(--border); margin-bottom: 28px; position: relative;
+}
+.app-header::before {
+    content: ''; position: absolute; top: 0; left: 0; right: 0; height: 3px;
+    background: linear-gradient(90deg, var(--accent), var(--accent2));
+}
+.app-title {
+    font-size: 2.6rem; font-weight: 800; letter-spacing: -0.04em; margin: 0 0 6px;
+    background: linear-gradient(135deg, var(--accent) 30%, var(--accent2));
+    -webkit-background-clip: text; -webkit-text-fill-color: transparent; background-clip: text;
+}
+.app-subtitle { color: var(--muted); font-size: 0.95rem; }
+.badge {
+    display: inline-block; margin-top: 10px;
+    background: var(--surface2); border: 1px solid var(--border);
+    border-radius: 20px; padding: 4px 14px;
+    font-size: 0.75rem; color: var(--accent); letter-spacing: 0.08em;
+    font-family: 'DM Mono', monospace;
+}
+
+label, .label-wrap span {
+    font-family: 'Syne', sans-serif !important; font-size: 0.75rem !important;
+    font-weight: 600 !important; letter-spacing: 0.12em !important;
+    text-transform: uppercase !important; color: var(--muted) !important;
+}
+textarea, input[type="text"] {
+    background: var(--surface) !important; border: 1px solid var(--border) !important;
+    border-radius: 8px !important; color: var(--text) !important;
+    font-family: 'DM Mono', monospace !important; font-size: 0.875rem !important;
+    transition: border-color 0.2s !important;
+}
+textarea:focus, input:focus {
+    border-color: var(--accent) !important;
+    box-shadow: 0 0 0 3px rgba(126,232,162,0.1) !important;
+}
+
+button.primary {
+    background: var(--accent) !important; color: #0d0d0f !important;
+    font-family: 'Syne', sans-serif !important; font-weight: 700 !important;
+    font-size: 0.88rem !important; letter-spacing: 0.06em !important;
+    border: none !important; border-radius: 8px !important;
+    padding: 13px 26px !important; text-transform: uppercase !important;
+    transition: all 0.2s !important;
+}
+button.primary:hover {
+    filter: brightness(1.1) !important; transform: translateY(-1px) !important;
+    box-shadow: 0 8px 20px rgba(126,232,162,0.2) !important;
+}
+button.secondary {
+    background: transparent !important; border: 1px solid var(--border) !important;
+    color: var(--muted) !important; font-family: 'Syne', sans-serif !important;
+    font-weight: 600 !important; border-radius: 8px !important; transition: all 0.2s !important;
+}
+button.secondary:hover { border-color: var(--accent2) !important; color: var(--accent2) !important; }
+
+.output-panel {
+    background: var(--surface) !important; border: 1px solid var(--border) !important;
+    border-radius: 12px !important; padding: 20px !important;
+}
+"""
+
+HEADER = """
+<div class="app-header">
+    <div class="app-title">⬡ PyReview · Gemma</div>
+    <div class="app-subtitle">Local Gemma 2B-IT · No token limits · Runs fully offline</div>
+    <div class="badge">google/gemma-2b-it · 4-bit quantised · CPU</div>
+</div>
+"""
+
+with gr.Blocks(css=CSS, title="PyReview — Gemma 2B") as demo:
+    gr.HTML(HEADER)
+
+    with gr.Row(equal_height=False):
+        with gr.Column(scale=5):
+            code_input = gr.Code(
+                label="YOUR PYTHON CODE",
+                language="python",
+                lines=20,
+                placeholder="# Paste your Python code here…",
+            )
+            desc_input = gr.Textbox(
+                label="WHAT SHOULD THIS CODE DO?",
+                placeholder="e.g. 'Sort a list of dicts by a given key, handling missing keys'",
+                lines=3,
+            )
+            with gr.Row():
+                submit_btn = gr.Button("🔍  Review Code", variant="primary", scale=3)
+                clear_btn  = gr.Button("Clear", variant="secondary", scale=1)
+
+            gr.Markdown(
+                "> ⏱ **CPU note:** Gemma 2B takes ~1–3 min per review on free CPU. "
+                "Output streams token-by-token as it generates.",
+                elem_classes=["output-panel"],
+            )
+
+        with gr.Column(scale=6):
+            output = gr.Markdown(
+                value="*Paste code + description, then click **Review Code**.*",
+                label="REVIEW OUTPUT",
+                elem_classes=["output-panel"],
+            )
+
+    gr.Examples(
+        examples=EXAMPLES,
+        inputs=[code_input, desc_input],
+        label="EXAMPLE SNIPPETS — click any row to load",
+    )
+
+    submit_btn.click(
+        fn=review_pipeline,
+        inputs=[code_input, desc_input],
+        outputs=output,
+        show_progress=True,
+    )
+    clear_btn.click(
+        fn=lambda: ("", "", "*Paste code + description, then click **Review Code**.*"),
+        outputs=[code_input, desc_input, output],
+    )
+
+if __name__ == "__main__":
+    demo.launch()