latest

README.md

@@ -0,0 +1,14 @@
+# DistilBERT for Text Classification
+
+This repository contains a fine-tuned DistilBERT model for text classification. The model is designed to classify text into four categories: SAFE, JAILBREAK, INJECTION, and PHISHING.
+
+## Model Details
+
+- Base model: DistilBERT (distilbert-base-uncased)
+- Task: Sequence Classification
+- Number of labels: 4
+- Labels: SAFE, JAILBREAK, INJECTION, PHISHING
+
+## Usage
+
+To use this model, you can leverage the Hugging Face Transformers library:

handler.py

@@ -0,0 +1,30 @@
+from transformers import Pipeline
+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")
+
+    def _sanitize_parameters(self, **kwargs):
+        return {}, {}, {}
+
+    def preprocess(self, inputs):
+        return self.tokenizer(inputs, return_tensors="pt", truncation=True, padding=True, max_length=512)
+
+    def _forward(self, model_inputs):
+        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()
+        predicted_label = self.label_mapping[predicted_class]
+        confidence = torch.softmax(logits, dim=1)[0][predicted_class].item()
+        
+        return {
+            "label": predicted_label,
+            "score": confidence
+        }