Improve JSON extraction to select a valid evaluation output

tinyllama_inference.py

@@ -15,7 +15,7 @@ def load_model():
     return tokenizer, model
 
 def evaluate_code(question, code):
-    # Refined prompt with explicit instructions and a "Your response:" line.
+    # Refined prompt with explicit instructions and a "Your response:" marker.
     prompt = f"""You are an expert code evaluator.
 Evaluate the following solution for the given problem.
 Rate the solution as follows:
@@ -38,24 +38,28 @@ Your response:"""
     inputs = tokenizer(prompt, return_tensors="pt")
     outputs = model.generate(
         **inputs,
-        max_new_tokens=100,      # Increase token allowance if needed
-        temperature=0.2,         # Allow some creativity, but mostly deterministic
+        max_new_tokens=120,      # Increase token allowance if needed
+        temperature=0.2,         # Allow some creativity
         pad_token_id=tokenizer.eos_token_id,
         do_sample=True
     )
     response_text = tokenizer.decode(outputs[0], skip_special_tokens=True)
     print("Raw model response:", response_text)  # Debug output
 
-    # Extract all JSON objects (non-greedy) and use the last one
-    matches = re.findall(r'\{.*?\}', response_text)
-    if matches:
-        json_text = matches[-1]  # Pick the last JSON block
-        try:
-            result = json.loads(json_text)
-        except Exception as e:
-            result = {"stars": 0, "feedback": "Evaluation failed. Unable to parse AI response."}
-    else:
-        result = {"stars": 0, "feedback": "Evaluation failed. Unable to extract JSON from AI response."}
+    # Use regex to find all JSON blocks (including across multiple lines)
+    matches = re.findall(r'\{.*?\}', response_text, re.DOTALL)
+    result = None
+    # Iterate over matches and select the first one that contains both "stars" and "feedback"
+    for m in matches:
+        if '"stars"' in m and '"feedback"' in m:
+            try:
+                temp = json.loads(m)
+                result = temp
+                break
+            except Exception:
+                continue
+    if result is None:
+        result = {"stars": 0, "feedback": "Evaluation failed. Unable to extract valid JSON from AI response."}
     
     return result