model_guard.py

"""
model_guard.py

Wraps the protectai/deberta-v3-base-prompt-injection-v2 model behind the
same simple interface as pattern_guard.check(), so the two approaches can
be compared apples-to-apples in the UI.

Model card: https://huggingface.co/protectai/deberta-v3-base-prompt-injection-v2
License: Apache 2.0 (underlying training data may carry additional terms —
see model card "Limitations" section)

Known limitations of this model (stated on its model card):
  - Does not reliably detect jailbreak-style attacks (only injection)
  - English only
  - Not recommended for scanning system prompts (produces false positives)

These limitations are intentionally surfaced in the UI — they are the
basis for the "defense in depth" argument: pattern matching catches
jailbreak phrasing that this model misses, and vice versa for novel
phrasings of instruction override that no static pattern anticipated.
"""

import time
from dataclasses import dataclass
from functools import lru_cache

from transformers import pipeline

MODEL_ID = "protectai/deberta-v3-base-prompt-injection-v2"


@dataclass
class ModelGuardResult:
    blocked: bool
    label: str
    confidence: float
    latency_ms: float


@lru_cache(maxsize=1)
def _get_classifier():
    """
    Lazily load the model once per process and cache it.
    Cold start will be slow (model download + load); subsequent calls
    are fast. This is loaded on CPU — no GPU required for this model size.
    """
    return pipeline(
        task="text-classification",
        model=MODEL_ID,
        truncation=True,
        max_length=512,
    )


def check(text: str) -> ModelGuardResult:
    """
    Run the DeBERTa-v3 prompt injection classifier on `text`.

    Returns label "INJECTION" (1) or "SAFE" (0) along with the model's
    confidence score and wall-clock latency for this call.
    """
    if not text or not text.strip():
        return ModelGuardResult(
            blocked=False, label="SAFE", confidence=1.0, latency_ms=0.0
        )

    classifier = _get_classifier()

    start = time.perf_counter()
    result = classifier(text)[0]
    elapsed = (time.perf_counter() - start) * 1000

    # protectai model outputs LABEL_0 (safe) / LABEL_1 (injection) or
    # "SAFE"/"INJECTION" depending on model version — normalize both.
    raw_label = result["label"].upper()
    is_injection = raw_label in ("LABEL_1", "INJECTION", "1")

    return ModelGuardResult(
        blocked=is_injection,
        label="INJECTION" if is_injection else "SAFE",
        confidence=round(result["score"], 4),
        latency_ms=round(elapsed, 2),
    )


def warm_up():
    """Call once at Space startup to pay the model-load cost up front,
    rather than on the first user request."""
    _get_classifier()
    check("warm up")