entrophy method

handler.py

@@ -18,17 +18,38 @@ class EndpointHandler:
         return self.predict(inputs)
 
     def predict(self, text):
-        if len(text.split()) < 4:
-            return {"label": "SAFE", "score": 1.0}
-
-        encoded_input = self.tokenizer(text, return_tensors='pt', truncation=True, max_length=128)
+        # Tokenize and encode the input
+        encoded_input = self.tokenizer(text, return_tensors='pt', truncation=True, max_length=512)
+        
+        # Get model prediction
         with torch.no_grad():
             outputs = self.model(**encoded_input)
-            probabilities = F.softmax(outputs.logits, dim=-1).cpu().numpy()[0]
-        confidence = np.max(probabilities)
-        predicted_label_idx = int(np.argmax(probabilities))
-        predicted_label = self.labels[predicted_label_idx]
-        return {"label": predicted_label, "score": float(confidence)}
+            logits = outputs.logits
+        
+        # Get probabilities
+        probabilities = F.softmax(logits, dim=-1).squeeze().numpy()
+        
+        # Get predicted class and confidence
+        predicted_class_idx = np.argmax(probabilities)
+        predicted_label = self.labels[predicted_class_idx]
+        confidence = probabilities[predicted_class_idx]
+
+        # 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
+
+        # Lower the confidence for very short inputs
+        if len(text.split()) < 4:
+            adjusted_confidence *= 0.5
+
+        return {
+            "label": predicted_label,
+            "score": float(adjusted_confidence),
+            "raw_scores": {label: float(prob) for label, prob in zip(self.labels.values(), probabilities)}
+        }
 
 def get_pipeline():
     return EndpointHandler