Improve JSON extraction with fallback methods

tinyllama_inference.py

@@ -14,8 +14,43 @@ def load_model():
         model = AutoModelForCausalLM.from_pretrained(model_name)
     return tokenizer, model
 
+def extract_json(response_text):
+    # First, try to extract all JSON blocks using regex with DOTALL.
+    matches = re.findall(r'\{.*?\}', response_text, 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
+
+    # 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.
+        try:
+            temp = json.loads(possible)
+            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.
+    return {"stars": 0, "feedback": "Evaluation failed. Unable to extract valid JSON from AI response."}
+
 def evaluate_code(question, code):
-    # 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,29 +73,15 @@ Your response:"""
     inputs = tokenizer(prompt, return_tensors="pt")
     outputs = model.generate(
         **inputs,
-        max_new_tokens=120,      # Increase token allowance if needed
-        temperature=0.2,         # Allow some creativity
+        max_new_tokens=120,      # Allow enough tokens for a complete response
+        temperature=0.2,         # Small randomness for 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
 
-    # 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."}
-    
+    result = extract_json(response_text)
     return result
 
 # For direct command-line testing.