Update app.py

app.py

@@ -1,411 +1,71 @@
-"""
-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,
-)
+from transformers import AutoTokenizer, AutoModelForCausalLM
 
-# ─────────────────────────────────────────────────────────────
-# MODEL LOADING  (runs once at startup)
-# ─────────────────────────────────────────────────────────────
+# --------- Load Model ---------
+MODEL_NAME = "google/gemma-2b-it"  # instruction-tuned version
 
-MODEL_ID = "google/gemma-2b-it"
+tokenizer = AutoTokenizer.from_pretrained(MODEL_NAME)
 
-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_NAME,
+    torch_dtype=torch.float16 if torch.cuda.is_available() else torch.float32,
+    device_map="auto"
 )
-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)
 
+# --------- Core Function ---------
+def evaluate_code(requirement, code):
+    prompt = f"""
+You are a strict software evaluator.
 
-# ─────────────────────────────────────────────────────────────
-# PRE-PROCESSING MODULE
-# ─────────────────────────────────────────────────────────────
+Your task:
+Check if the given code satisfies the requirement.
 
-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"
+Rules:
+- Answer ONLY "YES" or "NO"
+- Do NOT explain
+- Be strict and logical
 
+Requirement:
+{requirement}
 
-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"
-    )
-
+Code:
+{code}
 
-# ─────────────────────────────────────────────────────────────
-# 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
+Answer:
+"""
 
-    streamer = TextIteratorStreamer(
-        tokenizer,
-        skip_prompt=True,
-        skip_special_tokens=True,
-    )
+    inputs = tokenizer(prompt, return_tensors="pt").to(model.device)
 
-    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
+    outputs = model.generate(
+        **inputs,
+        max_new_tokens=5,
+        temperature=0.2,  # low randomness for consistency
+        do_sample=False
     )
 
-    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, [])
+    response = tokenizer.decode(outputs[0], skip_special_tokens=True)
 
+    # Extract only YES/NO
+    if "YES" in response.upper():
+        return "✅ YES"
+    elif "NO" in response.upper():
+        return "❌ NO"
+    else:
+        return "⚠️ Unable to determine"
 
-# ─────────────────────────────────────────────────────────────
-# UI MODULE
-# ─────────────────────────────────────────────────────────────
+# --------- UI ---------
+with gr.Blocks() as app:
+    gr.Markdown("## 🧠 Code Requirement Validator (Gemma 2)")
+    gr.Markdown("Enter a requirement and code. The model will check if the requirement is met.")
 
-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",
-    ],
-]
+    requirement = gr.Textbox(label="Requirement", lines=4, placeholder="e.g. Function should return sum of two numbers")
+    code = gr.Textbox(label="Code", lines=10, placeholder="Paste your code here...")
 
-CSS = """
-@import url('https://fonts.googleapis.com/css2?family=DM+Mono:wght@300;400;500&family=Syne:wght@400;600;700;800&display=swap');
+    output = gr.Textbox(label="Result")
 
-:root {
-    --bg: #0d0d0f;
-    --surface: #141418;
-    --surface2: #1c1c22;
-    --border: #2a2a35;
-    --accent: #7ee8a2;
-    --accent2: #4f8ef7;
-    --text: #e8e8f0;
-    --muted: #7070a0;
-}
+    btn = gr.Button("Evaluate")
 
-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,
-                
-            )
-            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],
-    )
+    btn.click(fn=evaluate_code, inputs=[requirement, code], outputs=output)
 
-if __name__ == "__main__":
-    demo.launch()
+app.launch()