YashVardhan-coder
Initial clean deploy
b8c035c
Raw
History Blame Contribute Delete
1.95 kB
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
# Load model and tokenizer
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" # Force CPU since it's a lightweight laptop configuration
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()
# Tokenize and evaluate
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]
# Class labels: 0: SAFE, 1: INJECTION
safe_prob = float(probabilities[0])
injection_prob = float(probabilities[1])
# Return dict report
is_safe = safe_prob > 0.5
return {
"safe": is_safe,
"risk_score": round(injection_prob, 4),
"category": "safe" if is_safe else "prompt_injection"
}
# Global Singleton access instance
classifier = DebertaClassifier()