Update inference.py

inference.py

@@ -8,15 +8,65 @@ MODEL_NAME   = os.getenv("MODEL_NAME",   "Qwen/Qwen2.5-7B-Instruct")
 BENCHMARK    = os.getenv("BENCHMARK",    "code-fix-env")
 API_KEY      = os.getenv("API_KEY") or os.getenv("HF_TOKEN")
 
-# 🔥 FIX 3: Early crash protection — no silent 0.1 failures
 if not API_KEY:
     raise ValueError("Missing API_KEY or HF_TOKEN environment variable")
 
 client = OpenAI(base_url=API_BASE_URL, api_key=API_KEY)
 
 
-def solve(task_name: str, task_input: str) -> str:
-    # 🔥 FIX 1 & 2: Correct START format with env + model
+# ── Shared grader ─────────────────────────────────────────────────────────────
+
+def grader(task_id, output):
+    if not output:
+        return 0.0
+
+    try:
+        data  = json.loads(output)
+        score = 0.5
+
+        if data.get("fixed_code"):
+            score += 0.2
+        if data.get("explanation"):
+            score += 0.2
+        if data.get("language"):
+            score += 0.1
+
+        return round(min(score, 1.0), 2)
+
+    except Exception:
+        return 0.2
+
+
+# ── Tasks ─────────────────────────────────────────────────────────────────────
+
+TASKS = [
+    {
+        "id":       "task_1",
+        "input":    "def add(a,b): return a-b",
+        "expected": "fix subtraction bug — should return a+b",
+        "grader":   grader,
+    },
+    {
+        "id":       "task_2",
+        "input":    "function x() { return 1+ }",
+        "expected": "fix syntax error — incomplete expression",
+        "grader":   grader,
+    },
+    {
+        "id":       "task_3",
+        "input":    "async function f(){ fetchData() }",
+        "expected": "fix missing await before fetchData()",
+        "grader":   grader,
+    },
+]
+
+
+# ── Solver ────────────────────────────────────────────────────────────────────
+
+def solve(task: dict) -> str:
+    task_name  = task["id"]
+    task_input = task["input"]
+
     print(f"[START] task={task_name} env={BENCHMARK} model={MODEL_NAME}", flush=True)
 
     try:
@@ -28,32 +78,30 @@ def solve(task_name: str, task_input: str) -> str:
                     "role": "system",
                     "content": (
                         "You are an expert developer.\n"
-                        "Return ONLY valid JSON.\n"
+                        "Return ONLY valid JSON — no markdown, no preamble.\n"
                         "Explanation must be MAX 2 lines.\n"
                         "Fixed code must be SHORT and COMPLETE.\n"
                         "Preserve all newlines and indentation in fixed_code.\n"
-                        "Do NOT cut output.\n"
                         "Format strictly:\n"
                         "{\"explanation\":\"...\",\"fixed_code\":\"...\",\"language\":\"...\"}"
                     ),
                 },
                 {
                     "role": "user",
-                    "content": f"Fix this code and explain the errors:\n{task_input}",
+                    "content": (
+                        f"Expected fix: {task['expected']}\n"
+                        f"Fix this code:\n{task_input}"
+                    ),
                 },
             ],
         )
 
         output = response.choices[0].message.content
-        score  = grade(task_input, output)
+        score  = task["grader"](task_name, output)
 
-        # 🔥 FIX 4: Dynamic reward from actual score, not static 0.5
-        # 🔥 FIX 1 & 2: Correct STEP format with action + done + error
         print(f"[STEP] step=1 action=solve reward={round(score, 2)} done=false error=null", flush=True)
-        print(f"[STEP] step=2 action=grade reward={round(score, 2)} done=true error=null", flush=True)
-
-        # 🔥 FIX 1 & 2: Correct END format with success + rewards list
-        print(f"[END] success=true steps=2 score={score} rewards=0.30,{score}", flush=True)
+        print(f"[STEP] step=2 action=grade reward={round(score, 2)} done=true  error=null", flush=True)
+        print(f"[END] success=true steps=2 score={score} rewards={score},{score}", flush=True)
 
         return output
 
@@ -64,46 +112,6 @@ def solve(task_name: str, task_input: str) -> str:
         return f"Error: {err}"
 
 
-def grade(task_input: str, output: str) -> float:
-    if not output or len(output.strip()) < 5:
-        return 0.1
-
-    # 🔥 FIX 5: Start at 0.3 — more realistic baseline, avoids fake inflation
-    score = 0.3
-    try:
-        raw    = output.strip().replace("```json", "").replace("```", "").strip()
-        parsed = json.loads(raw)
-        if parsed.get("explanation"):
-            score += 0.15
-        if parsed.get("fixed_code"):
-            score += 0.15
-        if parsed.get("language"):
-            score += 0.1
-    except Exception:
-        score = 0.2
-
-    return round(min(max(score, 0.1), 0.9), 2)
-
-
-# ── Tasks ─────────────────────────────────────────────────────────────────────
-TASKS = [
-    {
-        "id": "task_1",
-        "description": "Fix syntax error in Python",
-        "input": "def hello(\n    print('hello world')",
-    },
-    {
-        "id": "task_2",
-        "description": "Fix logic bug in JavaScript",
-        "input": "function add(a, b) { return a - b; }",
-    },
-    {
-        "id": "task_3",
-        "description": "Fix type error and missing await in async function",
-        "input": "async function fetchData() { let data = fetchFromAPI(); return data.json; }",
-    },
-]
-
-# 🔥 FIX 6: Single task only — let framework handle the loop
+# ── Entry point ───────────────────────────────────────────────────────────────
 if __name__ == "__main__":
-    solve(TASKS[0]["id"], TASKS[0]["input"])
+    solve(TASKS[0])