latest

handler.py

@@ -36,14 +36,24 @@ class EndpointHandler:
 
         # Additional analysis
         entropy = -np.sum(probabilities * np.log(probabilities + 1e-9))
-        max_prob_ratio = np.max(probabilities) / np.sort(probabilities)[-2]
-
-        # Adjust confidence based on entropy and probability ratio
-        adjusted_confidence = confidence * (1 - entropy/np.log(len(probabilities))) * max_prob_ratio
+        
+        # Adjust confidence based on entropy
+        adjusted_confidence = confidence * (1 - entropy/np.log(len(probabilities)))
 
-        # Lower the confidence for very short inputs
-        if len(text.split()) < 4:
-            adjusted_confidence *= 0.5
+        # Post-processing to better distinguish between INJECTION and JAILBREAK
+        injection_keywords = ['ignore', 'previous', 'instructions', 'don\'t', 'matter']
+        jailbreak_keywords = ['bypass', 'restrictions', 'override', 'security']
+        
+        injection_score = sum(keyword in text.lower() for keyword in injection_keywords) / len(injection_keywords)
+        jailbreak_score = sum(keyword in text.lower() for keyword in jailbreak_keywords) / len(jailbreak_keywords)
+        
+        if predicted_label in ['INJECTION', 'JAILBREAK']:
+            if injection_score > jailbreak_score:
+                predicted_label = 'INJECTION'
+            elif jailbreak_score > injection_score:
+                predicted_label = 'JAILBREAK'
+            
+            adjusted_confidence = max(adjusted_confidence, injection_score, jailbreak_score)
 
         return {
             "label": predicted_label,