added better handler

handler.py

@@ -1,32 +1,34 @@
-from transformers import AutoTokenizer, AutoModelForSequenceClassification
+import os
+import joblib
 import torch
+import numpy as np
+from transformers import DistilBertTokenizerFast, DistilBertForSequenceClassification
+import torch.nn.functional as F
 
 class EndpointHandler:
     def __init__(self, model_dir):
-        self.tokenizer = AutoTokenizer.from_pretrained(model_dir)
-        self.model = AutoModelForSequenceClassification.from_pretrained(model_dir)
-        self.label_mapping = {0: "SAFE", 1: "JAILBREAK", 2: "INJECTION", 3: "PHISHING"}
+        self.model = DistilBertForSequenceClassification.from_pretrained(model_dir)
+        self.tokenizer = DistilBertTokenizerFast.from_pretrained(model_dir)
+        self.label_mapping = joblib.load(os.path.join(model_dir, "label_mapping.joblib"))
+        self.labels = {v: k for k, v in self.label_mapping.items()}
 
     def __call__(self, inputs):
         if isinstance(inputs, dict) and 'inputs' in inputs:
             return self.predict(inputs['inputs'])
         return self.predict(inputs)
 
-    def predict(self, inputs):
-        # Tokenize the input
-        encoded_input = self.tokenizer(inputs, return_tensors='pt', truncation=True, padding=True)
-        
-        # Make prediction
+    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)
         with torch.no_grad():
-            output = self.model(**encoded_input)
-        
-        # Get the predicted class
-        predicted_class = torch.argmax(output.logits, dim=1).item()
-        
-        # Map the predicted class to its label
-        predicted_label = self.label_mapping[predicted_class]
-        
-        return {"label": predicted_label, "score": output.logits.softmax(dim=1).max().item()}
+            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)}
 
 def get_pipeline():
     return EndpointHandler