Judge ID Added

app.py

@@ -33,13 +33,21 @@ def llama_guard_classify(conv_prefix, response):
 
 def classify_prompt(category,conv_prefix, response):
     url = f"{API_ENDPOINT}/api/v1/judge/safety"
+    judge_id=None
+    if category=='response':
+        judge_id='eaad6030-c269-4ce8-8322-454127c380b8'
+    elif category=='prompt':
+        judge_id='7750e114-db3d-422f-be54-9692eb07baec'
+    else:
+        judge_id='7fd02b72-655b-4992-9380-ba496eefe12a'
 
     payload = {
         "model_name": "collinear_guard_classifier_beta",
         "nano_model_type": category,
         "conversation": conv_prefix,
         "response": response,
-        "space_id":"8b560bf4-3a76-4f00-b378-b528d02445c0"
+        "space_id":"8b560bf4-3a76-4f00-b378-b528d02445c0",
+        "judge_id":judge_id
     }
     headers = {
         "Authorization": f"Bearer {API_KEY}",
@@ -47,6 +55,7 @@ def classify_prompt(category,conv_prefix, response):
     }
 
     response = requests.request("POST", url, json=payload, headers=headers)
+    print(response.json())
 
     result = response.json()
     judgement = result['judgement']