Improve robustness of prediction extraction.

ansari-project · Apr 1, 2024 · e9c1462 · e9c1462
1 parent f44e09f
commit e9c1462
Showing 1 changed file with 12 additions and 9 deletions.
diff --git a/tests/test_answer_quality.py b/tests/test_answer_quality.py
@@ -43,15 +43,18 @@ def answer_question(question, q_temp, cache):
 
 
 def extract_prediction(row):
-    raw = row["json_prediction"]
-    raw = raw.replace("```", "")
-    raw = raw.replace("json", "")
-    raw = "{" + raw.split("{")[1]
-    raw = raw.split("}")[0] + "}"
-    raw = raw.strip()
-    LOGGER.info(f"raw is: {raw}")
-    raw_dict = json.loads(raw)
-    return str(raw_dict["answer"])
+    try: 
+        raw = row["json_prediction"]
+        raw = raw.replace("```", "")
+        raw = raw.replace("json", "")
+        raw = "{" + raw.split("{")[1]
+        raw = raw.split("}")[0] + "}"
+        raw = raw.strip()
+        LOGGER.info(f"raw is: {raw}")
+        raw_dict = json.loads(raw)
+        return str(raw_dict["answer"])
+    except IndexError:
+        return "OUTPUT_PARSE_ERROR"
 
 
 def is_correct(row):