Update utils/mistral.py

utils/mistral.py

@@ -20,6 +20,7 @@ if not HFT:
 client = InferenceClient(model="mistralai/Mistral-Nemo-Instruct-2407", token=HFT)
 
 # Function to clean model output
+'''
 def Data_Cleaner(text):
     pattern = r".*?format:"
     result = re.split(pattern, text, maxsplit=1)
@@ -36,6 +37,25 @@ def Data_Cleaner(text):
     except json.JSONDecodeError:
         logging.error("Data cleaning led to invalid JSON")
         return text  # Return the original text if cleaning goes wrong
+'''
+def Data_Cleaner(text):
+    # Use a regex pattern to extract everything between ```json and ```
+    pattern = r"```json\s*(\{.*?\})\s*```"  # Non-greedy matching inside braces
+    match = re.search(pattern, text, re.DOTALL)  # DOTALL to match newlines
+
+    if match:
+        json_str = match.group(1).strip()  # Extract JSON block
+    else:
+        logging.error("JSON block not found in the text")
+        return text  # Return the original text if no match is found
+
+    # Validate and return the cleaned JSON if it's valid
+    try:
+        json_obj = json.loads(json_str)  # Validate JSON
+        return json_obj  # Return the parsed JSON as a dictionary
+    except json.JSONDecodeError:
+        logging.error("Extracted text is not valid JSON")
+        return text  # Return the original text if JSON decoding fails
 
 
 # Function to call Mistral and process output