Update amazonmix.py

amazonmix.py

@@ -1,19 +1,19 @@
 import re
 from typing import Dict, Tuple, Optional
 
-def extract_solution(processed_str, method):
+def extract_solution(processed_str):
     # Split response to isolate assistant output
-    # if "Assistant:" in solution_str:
-    #     processed_str = solution_str.split("Assistant:", 1)[1]
-    # elif "<|im_start|>assistant" in solution_str:
-    #     processed_str = solution_str.split("<|im_start|>assistant", 1)[1]
-    # else:
-    #     print("[Error] Failed to locate model response header")
-    #     return None, solution_str
+    if "Assistant:" in solution_str:
+        processed_str = solution_str.split("Assistant:", 1)[1]
+    elif "<|im_start|>assistant" in solution_str:
+        processed_str = solution_str.split("<|im_start|>assistant", 1)[1]
+    else:
+        print("[Error] Failed to locate model response header")
+        return None, solution_str
 
     # Extract final answer using XML-style tags
     answer_pattern = r'<answer>(.*?)</answer>'
-    matches = list(re.finditer(answer_pattern, processed_str)) #, re.DOTALL
+    matches = list(re.finditer(answer_pattern, processed_str, re.DOTALL))
 
     if not matches:
         print("[Error] No valid answer tags found")
@@ -42,6 +42,7 @@ def validate_response_structure(processed_str: str) -> bool:
 
     # Check required tags
     tags = {
+        'think_start': ('<think>', 1),
         'think_end': ('</think>', 1),
         'answer_start': ('<answer>', 1),
         'answer_end': ('</answer>', 1)
@@ -59,13 +60,13 @@ def validate_response_structure(processed_str: str) -> bool:
             validation_passed = False
 
     # Verify tag order
-    if (positions['think_end'] > positions['answer_start'] or
-        positions['think_end'] > positions['answer_end'] or
+    if (positions['think_start'] > positions['think_end'] or
+        positions['think_end'] > positions['answer_start'] or
         positions['answer_start'] > positions['answer_end']):
         print("  [Error] Incorrect tag order: Expected <think>...</think><answer>...</answer>")
         validation_passed = False
     else:
-        print("   Tag sequence validation passed")
+        print("  Tag sequence validation passed")
 
     return validation_passed
 
@@ -76,6 +77,8 @@ def parse_ground_truth_text_format(ground_truth):
         format1_match = re.search(r"([a-zA-Z]):", ground_truth)
         if format1_match:
             return format1_match.group(1).strip()
+        else:
+            return ground_truth
     else:
         return ground_truth
 
@@ -98,7 +101,7 @@ def compute_score(solution_str: str, ground_truth: str, method='strict', format_
     ground_truth = parse_ground_truth_text_format(ground_truth)
 
     # Extract model answer
-    answer_text, processed_str= extract_solution(processed_str=solution_str, method=method)
+    answer_text, processed_str=extract_solution(processed_str=solution_str)
     print(f"\n[Model Response]\n{processed_str}")
     print(f"\n[Processed Model Response]\n{answer_text}")