Update app.py

app.py

@@ -53,31 +53,24 @@ def format_rules(rules):
     return formatted_rules
 
 def format_output(text):
-    think_match = re.search(r"<think>(.*?)</think>", text, flags=re.DOTALL)
-    answers = re.findall(r"<answer>(.*?)</answer>", text, flags=re.DOTALL)
+    reasoning = re.search(r"<think>(.*?)</think>", text, flags=re.DOTALL)
+    answer = re.search(r"<answer>(.*?)</answer>", text, flags=re.DOTALL)
+    explanation = re.search(r"<explanation>(.*?)</explanation>", text, flags=re.DOTALL)
 
     display = ""
-    if think_match:
-        display += "Reasoning:\n" + think_match.group(1).strip() + "\n\n"
-    if answers:
-        display += "Answers:\n" + "\n".join([a.strip() for a in answers])
+    if think_match and len(reasoning.group(1).strip()) > 1:
+        display += "Reasoning:\n" + reasoning.group(1).strip() + "\n\n"
+    if answer:
+        display += "Answer:\n" + answer.group(1).strip() + "\n\n"
+    if explanation and len(explanation.group(1).strip()) > 1:
+        display += "Explanation:\n" + explanation.group(1).strip() + "\n\n"
     return display.strip() if display else text.strip()
 
     
-
 def format_transcript(transcript):
     formatted_transcript = f"<transcript>\n{transcript}\n</transcript>\n"
     return formatted_transcript
 
-def get_example(
-    dataset_path="tomg-group-umd/compliance_benchmark",
-    subset="compliance",
-    split="test_handcrafted",
-    example_idx=0,
-):
-    dataset = load_dataset(dataset_path, subset, split=split)
-    example = dataset[example_idx]
-    return example[INPUT_FIELD]
 
 def get_message(model, input, system_prompt=SYSTEM_PROMPT, enable_thinking=True):
     message = model.apply_chat_template(system_prompt, input, enable_thinking=enable_thinking)