Merge pull request #11 from soumik12345/feat/llama-guard

guardrails_genie/guardrails/__init__.py

@@ -1,11 +1,15 @@
-from guardrails_genie.guardrails.entity_recognition import (
-    PresidioEntityRecognitionGuardrail,
-    RegexEntityRecognitionGuardrail,
-    RestrictedTermsJudge,
-    TransformersEntityRecognitionGuardrail,
-)
+try:
+    from guardrails_genie.guardrails.entity_recognition import (
+        PresidioEntityRecognitionGuardrail,
+        RegexEntityRecognitionGuardrail,
+        RestrictedTermsJudge,
+        TransformersEntityRecognitionGuardrail,
+    )
+except ImportError:
+    pass
 from guardrails_genie.guardrails.injection import (
     PromptInjectionClassifierGuardrail,
+    PromptInjectionLlamaGuardrail,
     PromptInjectionSurveyGuardrail,
 )
 from guardrails_genie.guardrails.secrets_detection import SecretsDetectionGuardrail
@@ -13,6 +17,7 @@ from guardrails_genie.guardrails.secrets_detection import SecretsDetectionGuardr
 from .manager import GuardrailManager
 
 __all__ = [
+    "PromptInjectionLlamaGuardrail",
     "PromptInjectionSurveyGuardrail",
     "PromptInjectionClassifierGuardrail",
     "PresidioEntityRecognitionGuardrail",

guardrails_genie/guardrails/entity_recognition/__init__.py

@@ -1,5 +1,16 @@
+import warnings
+
 from .llm_judge_entity_recognition_guardrail import RestrictedTermsJudge
-from .presidio_entity_recognition_guardrail import PresidioEntityRecognitionGuardrail
+
+try:
+    from .presidio_entity_recognition_guardrail import (
+        PresidioEntityRecognitionGuardrail,
+    )
+except ImportError:
+    warnings.warn(
+        "Presidio is not installed. If you want to use `PresidioEntityRecognitionGuardrail`, you can install the required packages using `pip install -e .[presidio]`"
+    )
+
 from .regex_entity_recognition_guardrail import RegexEntityRecognitionGuardrail
 from .transformers_entity_recognition_guardrail import (
     TransformersEntityRecognitionGuardrail,

guardrails_genie/guardrails/injection/__init__.py

@@ -1,4 +1,9 @@
 from .classifier_guardrail import PromptInjectionClassifierGuardrail
+from .llama_prompt_guardrail import PromptInjectionLlamaGuardrail
 from .survey_guardrail import PromptInjectionSurveyGuardrail
 
-__all__ = ["PromptInjectionSurveyGuardrail", "PromptInjectionClassifierGuardrail"]
+__all__ = [
+    "PromptInjectionLlamaGuardrail",
+    "PromptInjectionSurveyGuardrail",
+    "PromptInjectionClassifierGuardrail",
+]

guardrails_genie/guardrails/injection/llama_prompt_guardrail.py

@@ -0,0 +1,69 @@
+from typing import Optional
+
+import torch
+import torch.nn.functional as F
+import weave
+from transformers import AutoModelForSequenceClassification, AutoTokenizer
+
+from ..base import Guardrail
+
+
+class PromptInjectionLlamaGuardrail(Guardrail):
+    model_name: str = "meta-llama/Prompt-Guard-86M"
+    max_sequence_length: int = 512
+    temperature: float = 1.0
+    jailbreak_score_threshold: float = 0.5
+    indirect_injection_score_threshold: float = 0.5
+    _tokenizer: Optional[AutoTokenizer] = None
+    _model: Optional[AutoModelForSequenceClassification] = None
+
+    def model_post_init(self, __context):
+        self._tokenizer = AutoTokenizer.from_pretrained(self.model_name)
+        self._model = AutoModelForSequenceClassification.from_pretrained(
+            self.model_name
+        )
+
+    def get_class_probabilities(self, prompt):
+        inputs = self._tokenizer(
+            prompt,
+            return_tensors="pt",
+            padding=True,
+            truncation=True,
+            max_length=self.max_sequence_length,
+        )
+        with torch.no_grad():
+            logits = self._model(**inputs).logits
+        scaled_logits = logits / self.temperature
+        probabilities = F.softmax(scaled_logits, dim=-1)
+        return probabilities
+
+    @weave.op()
+    def get_score(self, prompt: str):
+        probabilities = self.get_class_probabilities(prompt)
+        return {
+            "jailbreak_score": probabilities[0, 2].item(),
+            "indirect_injection_score": (
+                probabilities[0, 1] + probabilities[0, 2]
+            ).item(),
+        }
+
+    @weave.op()
+    def guard(self, prompt: str):
+        score = self.get_score(prompt)
+        summary = ""
+        if score["jailbreak_score"] > self.jailbreak_score_threshold:
+            confidence = round(score["jailbreak_score"] * 100, 2)
+            summary += f"Prompt is deemed to be a jailbreak attempt with {confidence}% confidence."
+        if score["indirect_injection_score"] > self.indirect_injection_score_threshold:
+            confidence = round(score["indirect_injection_score"] * 100, 2)
+            summary += f" Prompt is deemed to be an indirect injection attempt with {confidence}% confidence."
+        return {
+            "safe": score["jailbreak_score"] < self.jailbreak_score_threshold
+            and score["indirect_injection_score"]
+            < self.indirect_injection_score_threshold,
+            "summary": summary.strip(),
+        }
+
+    @weave.op()
+    def predict(self, prompt: str):
+        return self.guard(prompt)

guardrails_genie/utils.py

@@ -1,18 +1,5 @@
-import os
-
 import pandas as pd
-import pymupdf4llm
 import weave
-import weave.trace
-from firerequests import FireRequests
-
-
-@weave.op()
-def get_markdown_from_pdf_url(url: str) -> str:
-    FireRequests().download(url, "temp.pdf", show_progress=False)
-    markdown = pymupdf4llm.to_markdown("temp.pdf", show_progress=False)
-    os.remove("temp.pdf")
-    return markdown
 
 
 class EvaluationCallManager:

pyproject.toml

@@ -9,25 +9,27 @@ dependencies = [
     "evaluate>=0.4.3",
     "google-generativeai>=0.8.3",
     "openai>=1.52.2",
-    "isort>=5.13.2",
-    "black>=24.10.0",
-    "ruff>=0.6.9",
-    "pip>=24.2",
-    "uv>=0.4.20",
     "weave @ git+https://github.com/wandb/weave@feat/eval-progressbar",
     "streamlit>=1.40.1",
     "python-dotenv>=1.0.1",
     "watchdog>=6.0.0",
-    "firerequests>=0.1.1",
-    "pymupdf4llm>=0.0.17",
     "transformers>=4.46.3",
     "torch>=2.5.1",
-    "presidio-analyzer>=2.2.355",
-    "presidio-anonymizer>=2.2.355",
     "instructor>=1.7.0",
 ]
 
 [project.optional-dependencies]
+presidio = [
+    "presidio-analyzer>=2.2.355",
+    "presidio-anonymizer>=2.2.355",
+]
+dev = [
+    "isort>=5.13.2",
+    "black>=24.10.0",
+    "ruff>=0.6.9",
+    "pip>=24.2",
+    "uv>=0.4.20",
+]
 docs = [
     "mkdocs>=1.6.1",
     "mkdocstrings>=0.26.1",