Fix

app.py

@@ -5,6 +5,7 @@ import os
 import json
 from datasets import load_dataset, Dataset, DatasetDict
 import pandas as pd
+import re
 API_ENDPOINT = "https://txl0ptjvttfogwt9.us-east-1.aws.endpoints.huggingface.cloud/v1/"
 API_KEY = "NA"
 HF_TOKEN=os.getenv("HF_TOKEN")
@@ -129,10 +130,9 @@ Now, please output the following as a JSON object:
         max_tokens=400
     )
     output = completion.choices[0].message.content
-    print("OUTPUT STR",output)
-    # output = json.loads(output)
-    # print(output)
-    if '"output" : "1"' in output:
+    output_value = re.search(r'"output"\s*:\s*"(\d+)"', output)
+    val = output_value.group(1) if output_value else None
+    if int(val)==1:
         return 'safe'
     else:
         return 'unsafe'