Upload 9 files

config.json

@@ -0,0 +1,34 @@
+{
+  "attention_probs_dropout_prob": 0.1,
+  "bos_token_id": null,
+  "eos_token_id": null,
+  "hidden_act": "gelu",
+  "hidden_dropout_prob": 0.1,
+  "hidden_size": 768,
+  "initializer_range": 0.02,
+  "intermediate_size": 3072,
+  "layer_norm_eps": 1e-07,
+  "legacy": true,
+  "max_position_embeddings": 512,
+  "max_relative_positions": -1,
+  "model_type": "deberta-v2",
+  "norm_rel_ebd": "layer_norm",
+  "num_attention_heads": 12,
+  "num_hidden_layers": 6,
+  "pad_token_id": 0,
+  "pooler_dropout": 0.0,
+  "pooler_hidden_act": "gelu",
+  "pooler_hidden_size": 768,
+  "pos_att_type": [
+    "p2c",
+    "c2p"
+  ],
+  "position_biased_input": false,
+  "position_buckets": 256,
+  "relative_attention": true,
+  "share_att_key": true,
+  "tie_word_embeddings": true,
+  "transformers_version": "5.8.1",
+  "type_vocab_size": 0,
+  "vocab_size": 128100
+}

family_encoder.pkl

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:45f0ae7f709029f857550b7955aaf092392fb19426131d0705c63b4fcc0cb6a9
+size 564

fine_encoder.pkl

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:025bf42a8d9aa15d15250157c2e90e8709f1bfd676a5db40e3c5c9c1f54833cd
+size 706

handler.py

@@ -0,0 +1,333 @@
+import os
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+import joblib
+
+from transformers import AutoTokenizer, AutoModel
+from typing import Dict, List, Any
+
+
+# =========================================================
+# 1. Multi-Task Architecture
+# =========================================================
+class MultiTaskModel(nn.Module):
+
+    def __init__(self, model_name, num_fine, num_family):
+        super().__init__()
+
+        # Base Encoder
+        self.encoder = AutoModel.from_pretrained(model_name)
+
+        hidden = self.encoder.config.hidden_size
+
+        self.dropout = nn.Dropout(0.2)
+
+        # Binary Classification Head
+        self.binary_head = nn.Linear(hidden, 1)
+
+        # Fine-Grained Attack Type Head
+        self.multi_head = nn.Linear(hidden, num_fine)
+
+        # Attack Family Head
+        self.family_head = nn.Linear(hidden, num_family)
+
+    # =====================================================
+    # Mean Pooling
+    # =====================================================
+    def mean_pooling(self, hidden, attention_mask):
+
+        mask = attention_mask.unsqueeze(-1).float()
+
+        pooled = (
+            (hidden * mask).sum(1)
+            /
+            mask.sum(1).clamp(min=1e-9)
+        )
+
+        return pooled
+
+    # =====================================================
+    # Forward Pass
+    # =====================================================
+    def forward(self, input_ids, attention_mask):
+
+        outputs = self.encoder(
+            input_ids=input_ids,
+            attention_mask=attention_mask
+        )
+
+        pooled = self.mean_pooling(
+            outputs.last_hidden_state,
+            attention_mask
+        )
+
+        x = self.dropout(pooled)
+
+        binary_logits = self.binary_head(x)
+        multi_logits = self.multi_head(x)
+        family_logits = self.family_head(x)
+
+        return (
+            binary_logits,
+            multi_logits,
+            family_logits
+        )
+
+
+# =========================================================
+# 2. Hugging Face Endpoint Handler
+# =========================================================
+class EndpointHandler:
+
+    def __init__(self, path=""):
+
+        # =================================================
+        # Device
+        # =================================================
+        self.device = torch.device(
+            "cuda" if torch.cuda.is_available() else "cpu"
+        )
+
+        print(f"[INFO] Using device: {self.device}")
+
+        # =================================================
+        # Load Label Encoders
+        # =================================================
+        self.fine_le = joblib.load(
+            os.path.join(path, "fine_encoder.pkl")
+        )
+
+        self.family_le = joblib.load(
+            os.path.join(path, "family_encoder.pkl")
+        )
+
+        # =================================================
+        # Load Tokenizer
+        # =================================================
+        self.tokenizer = AutoTokenizer.from_pretrained(path)
+
+        # =================================================
+        # Initialize Model
+        # =================================================
+        self.model = MultiTaskModel(
+            model_name="microsoft/deberta-v3-small",
+            num_fine=len(self.fine_le.classes_),
+            num_family=len(self.family_le.classes_)
+        ).to(self.device)
+
+        # =================================================
+        # Load Weights
+        # =================================================
+        checkpoint_path = os.path.join(
+            path,
+            "multitask_model_FINAL.pt"
+        )
+
+        checkpoint = torch.load(
+            checkpoint_path,
+            map_location=self.device
+        )
+
+        state_dict = (
+            checkpoint["model_state"]
+            if "model_state" in checkpoint
+            else checkpoint
+        )
+
+        self.model.load_state_dict(state_dict)
+
+        self.model.eval()
+
+        print("[INFO] RedLockX loaded successfully")
+
+        # =================================================
+        # Detection Threshold
+        # =================================================
+        self.threshold = 0.75
+
+    # =====================================================
+    # Predict Single Input
+    # =====================================================
+    def predict_single(self, text: str):
+
+        # ================================================
+        # Tokenize
+        # ================================================
+        tokenized = self.tokenizer(
+            text,
+            return_tensors="pt",
+            truncation=True,
+            padding=True,
+            max_length=512
+        )
+
+        tokenized = {
+            k: v.to(self.device)
+            for k, v in tokenized.items()
+        }
+
+        # ================================================
+        # Inference
+        # ================================================
+        with torch.no_grad():
+
+            (
+                binary_logits,
+                multi_logits,
+                family_logits
+            ) = self.model(
+                tokenized["input_ids"],
+                tokenized["attention_mask"]
+            )
+
+            # ============================================
+            # Binary Probability
+            # ============================================
+            danger_prob = torch.sigmoid(
+                binary_logits
+            ).item()
+
+            is_dangerous = danger_prob > self.threshold
+
+            # ============================================
+            # Multi-Class Probabilities
+            # ============================================
+            multi_probs = F.softmax(
+                multi_logits,
+                dim=1
+            )
+
+            family_probs = F.softmax(
+                family_logits,
+                dim=1
+            )
+
+            fine_idx = torch.argmax(
+                multi_probs,
+                dim=1
+            ).item()
+
+            family_idx = torch.argmax(
+                family_probs,
+                dim=1
+            ).item()
+
+            fine_score = multi_probs[0][fine_idx].item()
+
+            family_score = family_probs[0][family_idx].item()
+
+            # ============================================
+            # SAFE Handling
+            # ============================================
+            if is_dangerous:
+
+                attack_type = self.fine_le.inverse_transform(
+                    [fine_idx]
+                )[0]
+
+                attack_family = self.family_le.inverse_transform(
+                    [family_idx]
+                )[0]
+
+            else:
+
+                attack_type = "none"
+                attack_family = "none"
+
+                fine_score = 0.0
+                family_score = 0.0
+
+        # ================================================
+        # Basic Explainability
+        # ================================================
+        suspicious_keywords = [
+            "ignore",
+            "override",
+            "reveal",
+            "system prompt",
+            "developer mode",
+            "bypass",
+            "disable",
+            "forget instructions",
+            "pretend",
+            "simulate",
+            "jailbreak"
+        ]
+
+        found_keywords = []
+
+        text_lower = text.lower()
+
+        for keyword in suspicious_keywords:
+
+            if keyword in text_lower:
+                found_keywords.append(keyword)
+
+        # ================================================
+        # Final Response
+        # ================================================
+        return {
+
+            "status": (
+                "DANGEROUS"
+                if is_dangerous
+                else "SAFE"
+            ),
+
+            "confidence": round(
+                danger_prob
+                if is_dangerous
+                else 1 - danger_prob,
+                4
+            ),
+
+            "attack_type": {
+                "label": attack_type,
+                "score": round(fine_score, 4)
+            },
+
+            "attack_family": {
+                "label": attack_family,
+                "score": round(family_score, 4)
+            },
+
+            "trigger_words": found_keywords
+        }
+
+    # =====================================================
+    # Main Endpoint Entry
+    # =====================================================
+    def __call__(
+        self,
+        data: Dict[str, Any]
+    ) -> List[Dict[str, Any]]:
+
+        # ================================================
+        # Extract Inputs
+        # ================================================
+        inputs = (
+            data["inputs"]
+            if isinstance(data, dict)
+            else data
+        )
+
+        # ================================================
+        # Single Input → Convert to List
+        # ================================================
+        if isinstance(inputs, str):
+            inputs = [inputs]
+
+        # ================================================
+        # Batch Inference
+        # ================================================
+        results = []
+
+        for text in inputs:
+
+            result = self.predict_single(text)
+
+            results.append(result)
+
+        return results
+    

multitask_model_FINAL.pt

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:7fa886a0b5a8d9e062d6a6a49b78e4250b33548c97f74c4536901cb4fcbc7ac9
+size 565316387

requirements.txt

@@ -0,0 +1,5 @@
+torch
+transformers
+sentencepiece
+joblib
+scikit-learn==1.6.1

tokenizer_config.json

@@ -0,0 +1,24 @@
+{
+  "add_prefix_space": true,
+  "backend": "tokenizers",
+  "bos_token": "[CLS]",
+  "cls_token": "[CLS]",
+  "do_lower_case": false,
+  "eos_token": "[SEP]",
+  "extra_special_tokens": [
+    "[PAD]",
+    "[CLS]",
+    "[SEP]"
+  ],
+  "is_local": false,
+  "local_files_only": false,
+  "mask_token": "[MASK]",
+  "model_max_length": 1000000000000000019884624838656,
+  "pad_token": "[PAD]",
+  "sep_token": "[SEP]",
+  "split_by_punct": false,
+  "tokenizer_class": "DebertaV2Tokenizer",
+  "unk_id": 3,
+  "unk_token": "[UNK]",
+  "vocab_type": "spm"
+}

tokenizer_meta.json

@@ -0,0 +1 @@
+{"model": "microsoft/deberta-v3-small", "max_len": 512}