remove old json string cleaning function and use Patronus nemo guardrails code

app.py

@@ -6,6 +6,7 @@ import spaces
 import openai
 import json
 import re
+from typing import List, Optional, Tuple, Union
 
 HF_TOKEN = os.environ.get("HF_TOKEN", None)
 LEPTON_API_TOKEN = os.environ.get("LEPTON_API_TOKEN", None)
@@ -66,25 +67,31 @@ HEADER = """
 **Getting Started**: Provide a question and document or context given to your model in addition to the answer given by the model and then click submit. The output panel will indicate whether the reponse is a hallucination (Fail) or if it is faithful to the given document or context (Pass) through the score Pass or Fail and provide reasoning behind the score.
 """
 
-def clean_json_string(json_str):
-    # # Replace single quotes with double quotes, but not apostrophes within words
-    # json_str = re.sub(r"(?<!\\)'([^']*)'", r'"\1"', json_str)
-    # # Add quotes around PASS or FAIL if they're not already quoted
-    # json_str = re.sub(r'"SCORE":\s*(PASS|FAIL)', r'"SCORE": "\1"', json_str)
-
-    # # Escape double quotes
-    # json_str = re.sub(r'(?<=: ")([^"]*)"([^"]*)"([^"]*)"', lambda m: m.group(1) + m.group(2).replace('"', '\\"') + m.group(3), json_str)
-    
-    # return json_str
-    
-    json_str = re.sub(r'"(\w+)"\s*:', r'\1:', json_str)
-    json_str = re.sub(r':\s*"([^"]*)"', r':"\1"', json_str)
-
-    def unescape_quotes(match):
-        return match.group(0).replace('\\"', '"')
-    
-    json_str = re.sub(r'"([^"\\]*(?:\\.[^"\\]*)*)"', lambda m: '"' + unescape_quotes(m) + '"', json_str)
-    return json_str
+def parse_patronus_lynx_response(
+    response: str,
+) -> Tuple[bool, Union[List[str], None]]:
+    """
+    Parses the response from the Patronus Lynx LLM and returns a tuple of:
+    - Whether the response is hallucinated or not.
+    - A reasoning trace explaining the decision.
+    """
+    # Default to hallucinated
+    hallucination, reasoning = True, None
+    reasoning_pattern = r'"REASONING":\s*\[(.*?)\]'
+    score_pattern = r'"SCORE":\s*"?\b(PASS|FAIL)\b"?'
+
+    reasoning_match = re.search(reasoning_pattern, response, re.DOTALL)
+    score_match = re.search(score_pattern, response)
+
+    if score_match:
+        score = score_match.group(1)
+        if score == "PASS":
+            hallucination = False
+    if reasoning_match:
+        reasoning_content = reasoning_match.group(1)
+        reasoning = re.split(r"['\"],\s*['\"]", reasoning_content)
+
+    return hallucination, reasoning
 
 def model_call(question, document, answer):
     if question == "" or document == "" or answer == "":
@@ -96,13 +103,7 @@ def model_call(question, document, answer):
         prompt=NEW_FORMAT
     )
     print("RESPONSE FROM CLIENT:", response)
-    generated_text = clean_json_string(response.choices[0].text)
-    print("CLEAN json", generated_text)
-    generated_text = json.loads(generated_text)
-    print("GENERATED TEXT", generated_text)
-    print("type of GENERATED TEXT", type(generated_text))
-    reasoning = generated_text["REASONING"][0]
-    score = generated_text["SCORE"]
+    score, reasoning = parse_patronus_lynx_response(response.choices[0].text)
     return reasoning, score
 
 inputs = [