Update app.py

app.py

@@ -3,7 +3,7 @@ import torch
 from transformers import AutoTokenizer, AutoModelForCausalLM
 
 # --------- Load Model ---------
-MODEL_NAME = "google/gemma-2b-it"  # instruction-tuned version
+MODEL_NAME = "google/gemma-2b-it"
 
 tokenizer = AutoTokenizer.from_pretrained(MODEL_NAME)
 
@@ -13,54 +13,92 @@ model = AutoModelForCausalLM.from_pretrained(
     device_map="auto"
 )
 
-# --------- Core Function ---------
-def evaluate_code(requirement, code):
-    prompt = f"""
-You are a strict software evaluator.
-
-Your task:
-Check if the given code satisfies the requirement.
-
-Rules:
-- Answer ONLY "YES" or "NO"
-- Do NOT explain
-- Be strict and logical
-
+# --------- Improved Prompt Function ---------
+def pre_processing(requirement, code):
+    return f"""
+# ROLE
+You are a strict software testing system.
+You evaluate Python code against given requirements and do not make assumptions.
+
+# DESCRIPTION
+Your job is to compare the Python code against the requirement
+and ensure that the code strictly meets the requirement.
+
+# INSTRUCTIONS
+Step 1: Break the requirement into clear testable conditions.
+Step 2: Check the code against each condition.
+Step 3: Mark each condition as YES or NO.
+Step 4: If all conditions are YES → YES, else → NO.
+
+# RULES
+- Be strict
+- Do not assume missing functionality
+- Do not execute the code
+- Do not guess
+- If unsure, mark NO
+- Follow output format exactly
+
+# INPUT DATA
 Requirement:
 {requirement}
 
 Code:
 {code}
 
-Answer:
+# OUTPUT FORMAT
+
+Conditions:
+1. Condition
+2. Condition
+
+Evaluation of each condition:
+1. YES/NO
+2. YES/NO
+
+Final Answer: YES or NO
 """
 
+# --------- Core Function ---------
+def evaluate_code(requirement, code):
+    prompt = pre_processing(requirement, code)
+
     inputs = tokenizer(prompt, return_tensors="pt").to(model.device)
 
     outputs = model.generate(
         **inputs,
-        max_new_tokens=5,
-        temperature=0.2,  # low randomness for consistency
+        max_new_tokens=150,   # increased for structured output
+        temperature=0.1,      # more deterministic
         do_sample=False
     )
 
     response = tokenizer.decode(outputs[0], skip_special_tokens=True)
 
-    # Extract only YES/NO
-    if "YES" in response.upper():
+    # --------- Extract Final Answer ---------
+    response_upper = response.upper()
+
+    if "FINAL ANSWER: YES" in response_upper:
         return "✅ YES"
-    elif "NO" in response.upper():
+    elif "FINAL ANSWER: NO" in response_upper:
         return "❌ NO"
     else:
-        return "⚠️ Unable to determine"
+        return "⚠️ Unable to determine\n\nFull Response:\n" + response
 
 # --------- 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.")
+    gr.Markdown("## 🧠 Code Requirement Validator (Gemma)")
+    gr.Markdown("Checks if your Python code satisfies a requirement using structured evaluation.")
 
-    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...")
+    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 Python code here..."
+    )
 
     output = gr.Textbox(label="Result")