latest

handler.py

@@ -1,30 +1,30 @@
-from transformers import Pipeline
+from transformers import AutoTokenizer, AutoModelForSequenceClassification
 import torch
-import joblib
 
-class CustomPipeline(Pipeline):
-    def __init__(self, model, tokenizer, device=-1, **kwargs):
-        super().__init__(model=model, tokenizer=tokenizer, device=device, **kwargs)
-        self.label_mapping = joblib.load("label_mapping.joblib")
+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"}
 
-    def _sanitize_parameters(self, **kwargs):
-        return {}, {}, {}
+    def __call__(self, inputs):
+        return self.predict(inputs)
 
-    def preprocess(self, inputs):
-        return self.tokenizer(inputs, return_tensors="pt", truncation=True, padding=True, max_length=512)
-
-    def _forward(self, model_inputs):
+    def predict(self, inputs):
+        # Tokenize the input
+        encoded_input = self.tokenizer(inputs, return_tensors='pt', truncation=True, padding=True)
+        
+        # Make prediction
         with torch.no_grad():
-            outputs = self.model(**model_inputs)
-        return outputs
-
-    def postprocess(self, model_outputs):
-        logits = model_outputs.logits
-        predicted_class = torch.argmax(logits, dim=1).item()
+            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]
-        confidence = torch.softmax(logits, dim=1)[0][predicted_class].item()
         
-        return {
-            "label": predicted_label,
-            "score": confidence
-        }
+        return {"label": predicted_label, "score": output.logits.softmax(dim=1).max().item()}
+
+def get_pipeline():
+    return EndpointHandler