Skip to content

InjecGuard

any_guardrail.guardrails.injec_guard.injec_guard

InjecGuard

Bases: HuggingFace

Prompt injection detection encoder based model.

For more information, please see the model card:

Source code in src/any_guardrail/guardrails/injec_guard/injec_guard.py 
 9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
class InjecGuard(HuggingFace):
    """Prompt injection detection encoder based model.

    For more information, please see the model card:

    - [InjecGuard](https://huggingface.co/leolee99/InjecGuard).
    """

    SUPPORTED_MODELS: ClassVar = ["leolee99/InjecGuard"]

    def _load_model(self) -> None:
        from transformers import AutoModelForSequenceClassification, AutoTokenizer

        self.model = AutoModelForSequenceClassification.from_pretrained(self.model_id, trust_remote_code=True)
        self.tokenizer = AutoTokenizer.from_pretrained(self.model_id)  # type: ignore[no-untyped-call]

    def _post_processing(self, model_outputs: dict[str, Any]) -> GuardrailOutput:
        return _match_injection_label(model_outputs, INJECGUARD_LABEL, self.model.config.id2label)