Update app.py

app.py

@@ -163,7 +163,7 @@ def gradient_cuff_reject(message,with_defense, sample_times,perturb_times,thresh
     labels=get_labels(responses)
     results=[None for _ in range(perturb_times+1)]
     results[0]=sum(labels)/len(labels)
-    if results[0]>0.5:
+    if results[0]>=0.5:
         return (True,1-results[0],None)
     
     # second-stage rejection
@@ -216,7 +216,8 @@ def chat(message, history, with_defense,threshold):
     messages = tok.apply_chat_template(chat, tokenize=False, add_generation_prompt=True)
     # Tokenize the messages string
     input_ids = tok([messages], return_tensors="pt")["input_ids"]
-    response= "[Gradient Cuff Checking: "+reject_information + "]\n"+ chat_engine(input_ids)
+    #response= "[Gradient Cuff Checking: "+reject_information + "]\n"+ chat_engine(input_ids)
+    response=chat_engine(input_ids)
     response=response.split(" ")
 
     # Initialize an empty string to store the generated text