| import torch |
| from transformers import AutoTokenizer, AutoModelForSequenceClassification |
| import threading |
| from typing import Dict, Any |
|
|
| class DebertaClassifier: |
| _instance = None |
| _lock = threading.Lock() |
|
|
| def __new__(cls): |
| with cls._lock: |
| if cls._instance is None: |
| cls._instance = super(DebertaClassifier, cls).__new__(cls) |
| cls._instance._initialized = False |
| return cls._instance |
|
|
| def initialize(self): |
| if self._initialized: |
| return |
| |
| |
| self.model_name = "deepset/deberta-v3-base-injection" |
| self.tokenizer = AutoTokenizer.from_pretrained(self.model_name) |
| self.model = AutoModelForSequenceClassification.from_pretrained(self.model_name) |
| self.device = "cpu" |
| self.model.to(self.device) |
| self._initialized = True |
|
|
| def predict(self, text: str) -> Dict[str, Any]: |
| """ |
| Classifies whether input text is a prompt injection. |
| Returns safety report with probability score. |
| """ |
| if not self._initialized: |
| self.initialize() |
|
|
| |
| inputs = self.tokenizer(text, return_tensors="pt", truncation=True, max_length=512) |
| with torch.no_grad(): |
| outputs = self.model(**inputs) |
| logits = outputs.logits |
| probabilities = torch.softmax(logits, dim=1)[0] |
| |
| |
| safe_prob = float(probabilities[0]) |
| injection_prob = float(probabilities[1]) |
|
|
| |
| is_safe = safe_prob > 0.5 |
| return { |
| "safe": is_safe, |
| "risk_score": round(injection_prob, 4), |
| "category": "safe" if is_safe else "prompt_injection" |
| } |
|
|
| |
| classifier = DebertaClassifier() |
|
|