Change output

app.py

@@ -43,13 +43,18 @@ def classify_prompt(category,conv_prefix, response):
 
     response = requests.request("POST", url, json=payload, headers=headers)
 
-    print(response.text)
-    # val = output_value.group(1) if output_value else None
-    # if int(val)==1:
-    #     return 'refusal' if category=='refusal' else 'safe'
-    # else:
-    #     return 'non refusal' if category=='refusal' else 'unsafe'
-    return 'a'
+    result = response.json()
+    judgement = result['judgement']
+    if category =='refusal':
+        if judgement==1:
+            return 'Non Refusal'
+        else:
+            return 'Refusal'
+    else:
+        if judgement ==1:
+            return 'Safe'
+        else:
+            return 'Unsafe'
 
 async def add_to_dataset(conv_prefix, response,llama_resp,collinear_resp):
     conv_prefix = json.loads(conv_prefix)