Improve JSON extraction with fallback methods

tinyllama_inference.py

@@ -15,39 +15,25 @@ def load_model():
     return tokenizer, model
 
 def extract_json(response_text):
-    # First, try to extract all JSON blocks using regex with DOTALL.
+    # First attempt: Use regex (non-greedy, with DOTALL) to find JSON blocks
     matches = re.findall(r'\{.*?\}', response_text, re.DOTALL)
-    for m in matches:
-        json_text = m.strip()
+    # Check the matches in reverse order (last one might be the evaluation output)
+    for m in reversed(matches):
         try:
-            temp = json.loads(json_text)
+            temp = json.loads(m)
             if isinstance(temp, dict) and "stars" in temp and "feedback" in temp:
                 return temp
         except Exception:
             continue
-
-    # Fallback: try splitting the text on "Your response:" and then extract JSON.
-    parts = response_text.split("Your response:")
-    if len(parts) > 1:
-        possible = parts[-1].strip()
-        # Try to extract JSON from this part.
+    # Fallback: try extracting JSON from each line that looks like a JSON object
+    json_lines = [line.strip() for line in response_text.splitlines() if line.strip().startswith('{') and line.strip().endswith('}')]
+    for line in reversed(json_lines):
         try:
-            temp = json.loads(possible)
+            temp = json.loads(line)
             if isinstance(temp, dict) and "stars" in temp and "feedback" in temp:
                 return temp
         except Exception:
-            # If it fails, try regex on this part.
-            matches = re.findall(r'\{.*?\}', possible, re.DOTALL)
-            for m in matches:
-                json_text = m.strip()
-                try:
-                    temp = json.loads(json_text)
-                    if isinstance(temp, dict) and "stars" in temp and "feedback" in temp:
-                        return temp
-                except Exception:
-                    continue
-
-    # If all methods fail, return a fallback result.
+            continue
     return {"stars": 0, "feedback": "Evaluation failed. Unable to extract valid JSON from AI response."}
 
 def evaluate_code(question, code):
@@ -73,13 +59,13 @@ Your response:"""
     inputs = tokenizer(prompt, return_tensors="pt")
     outputs = model.generate(
         **inputs,
-        max_new_tokens=120,      # Allow enough tokens for a complete response
-        temperature=0.2,         # Small randomness for creativity
+        max_new_tokens=120,      # Increase token allowance for a complete response
+        temperature=0.2,         # Low randomness for deterministic output
         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
+    print("Raw model response:", response_text)  # Debug: Inspect raw output
 
     result = extract_json(response_text)
     return result