Update JSON extraction in tinyllama_inference.py to select last JSON block

tinyllama_inference.py

@@ -15,39 +15,43 @@ def load_model():
     return tokenizer, model
 
 def evaluate_code(question, code):
-    # Refined prompt: instruct the model to output exactly one JSON object.
     prompt = f"""You are an expert code evaluator.
 Evaluate the following solution for the given problem.
 Respond with exactly one JSON object (with no extra text) that has exactly two keys:
   "stars": an integer between 0 and 5 (0 means completely incorrect, 5 means excellent),
   "feedback": a concise string message.
 The JSON must start with '{{' and end with '}}'.
-Do not output any text besides the JSON.
+Do not output anything else.
 Question: "{question}"
 Solution: "{code}"
 Your response:"""
+    
     tokenizer, model = load_model()
     inputs = tokenizer(prompt, return_tensors="pt")
     outputs = model.generate(
         **inputs,
         max_new_tokens=100,      # Allow enough tokens for a complete response
-        temperature=0.2,         # Small randomness for a bit of creativity
+        temperature=0.2,         # Small randomness for creativity
         pad_token_id=tokenizer.eos_token_id,
-        do_sample=True            # Enable sampling to encourage model generation
+        do_sample=True            # Enable sampling to encourage generation
     )
     response_text = tokenizer.decode(outputs[0], skip_special_tokens=True)
     print("Raw model response:", response_text)  # Debug output
-    
-    # Use findall to get all JSON objects and take the last one
+
+    # Use regex to extract all JSON objects (non-greedy)
     matches = re.findall(r'\{.*?\}', response_text)
-    if matches:
-        json_text = matches[-1]
+    result = None
+    for m in matches:
         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."}
+            temp = json.loads(m)
+            # Check that the parsed JSON contains both expected keys
+            if isinstance(temp, dict) and "stars" in temp and "feedback" in temp:
+                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