Refusal Changed

app.py

@@ -115,7 +115,11 @@ Now, please output the following as a JSON object:
 """
 
     template = Template(template_str)
-    rendered_message = template.render(conversation_prefix=conv_prefix, response=response,scoring_criteria=[{"score":0,"description":unsafe_text},{"score":1,"description":safe_text}])
+    if category=='prompt' or category=='response':
+        scoring_criteria = [{"score":0,"description":unsafe_text},{"score":1,"description":safe_text}]
+    else:
+        scoring_criteria = [{"score":1,"description":unsafe_text},{"score":0,"description":safe_text}]
+    rendered_message = template.render(conversation_prefix=conv_prefix, response=response,scoring_criteria=scoring_criteria)
     client = openai.OpenAI(
         base_url=API_ENDPOINT,
         api_key=API_KEY
@@ -133,9 +137,9 @@ Now, please output the following as a JSON object:
     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'
+        return 'refusal' if category=='refusal' else 'safe'
     else:
-        return 'unsafe'
+        return 'non refusal' if category=='refusal' else 'unsafe'
 
 async def add_to_dataset(safe_text,unsafe_text,conv_prefix, response,llama_resp,collinear_resp):
     conv_prefix = json.loads(conv_prefix)
@@ -166,7 +170,10 @@ def process_inputs(category,safe_text,unsafe_text,conv_prefix, response_content)
     response = {"role": "assistant", "content": response_content}
     conv_prefix = json.loads(conv_prefix)
     output = classify_prompt(category,safe_text,unsafe_text,conv_prefix, response)
-    llama_output = llama_guard_classify(conv_prefix, response)
+    if category=='response':
+        llama_output = llama_guard_classify(conv_prefix, response)
+    else:
+        llama_output = 'NA'
     return output,llama_output