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'