add: off-the-shelf prompt injection guardrail

guardrails_genie/guardrails/__init__.py

@@ -1,4 +1,8 @@
-from .injection import PromptInjectionSurveyGuardrail
+from .injection import PromptInjectionProtectAIGuardrail, PromptInjectionSurveyGuardrail
 from .manager import GuardrailManager
 
-__all__ = ["PromptInjectionSurveyGuardrail", "GuardrailManager"]
+__all__ = [
+    "PromptInjectionSurveyGuardrail",
+    "PromptInjectionProtectAIGuardrail",
+    "GuardrailManager",
+]

guardrails_genie/guardrails/injection/__init__.py

@@ -1,3 +1,4 @@
+from .protectai_guardrail import PromptInjectionProtectAIGuardrail
 from .survey_guardrail import PromptInjectionSurveyGuardrail
 
-__all__ = ["PromptInjectionSurveyGuardrail"]
+__all__ = ["PromptInjectionSurveyGuardrail", "PromptInjectionProtectAIGuardrail"]

guardrails_genie/guardrails/injection/protectai_guardrail.py

@@ -0,0 +1,34 @@
+from typing import Optional
+
+import torch
+import weave
+from transformers import AutoModelForSequenceClassification, AutoTokenizer, pipeline
+from transformers.pipelines.base import Pipeline
+
+from ..base import Guardrail
+
+
+class PromptInjectionProtectAIGuardrail(Guardrail):
+    model_name: str = "ProtectAI/deberta-v3-base-prompt-injection-v2"
+    _classifier: Optional[Pipeline] = None
+
+    def model_post_init(self, __context):
+        tokenizer = AutoTokenizer.from_pretrained(self.model_name)
+        model = AutoModelForSequenceClassification.from_pretrained(self.model_name)
+        self._classifier = pipeline(
+            "text-classification",
+            model=model,
+            tokenizer=tokenizer,
+            truncation=True,
+            max_length=512,
+            device=torch.device("cuda" if torch.cuda.is_available() else "cpu"),
+        )
+
+    @weave.op()
+    def predict(self, prompt: str):
+        return self._classifier(prompt)
+
+    @weave.op()
+    def guard(self, prompt: str):
+        response = self.predict(prompt)
+        return {"safe": response[0]["label"] != "INJECTION"}