mirror of
https://github.com/meta-llama/llama-stack.git
synced 2025-10-03 19:57:35 +00:00
112 lines
4.1 KiB
Python
112 lines
4.1 KiB
Python
from enum import auto, Enum
|
|
from typing import List
|
|
|
|
import torch
|
|
|
|
from models.llama3.datatypes import Message
|
|
from termcolor import cprint
|
|
from transformers import AutoModelForSequenceClassification, AutoTokenizer
|
|
|
|
from .base import message_content_as_str, OnViolationAction, ShieldResponse, TextShield
|
|
from toolchain.safety.api.datatypes import * # noqa: F403
|
|
|
|
|
|
class PromptGuardShield(TextShield):
|
|
|
|
class Mode(Enum):
|
|
INJECTION = auto()
|
|
JAILBREAK = auto()
|
|
|
|
_instances = {}
|
|
_model_cache = None
|
|
|
|
@staticmethod
|
|
def instance(
|
|
model_dir: str,
|
|
threshold: float = 0.9,
|
|
temperature: float = 1.0,
|
|
mode: "PromptGuardShield.Mode" = Mode.JAILBREAK,
|
|
on_violation_action=OnViolationAction.RAISE,
|
|
) -> "PromptGuardShield":
|
|
action_value = on_violation_action.value
|
|
key = (model_dir, threshold, temperature, mode, action_value)
|
|
if key not in PromptGuardShield._instances:
|
|
PromptGuardShield._instances[key] = PromptGuardShield(
|
|
model_dir=model_dir,
|
|
threshold=threshold,
|
|
temperature=temperature,
|
|
mode=mode,
|
|
on_violation_action=on_violation_action,
|
|
)
|
|
return PromptGuardShield._instances[key]
|
|
|
|
def __init__(
|
|
self,
|
|
model_dir: str,
|
|
threshold: float = 0.9,
|
|
temperature: float = 1.0,
|
|
mode: "PromptGuardShield.Mode" = Mode.JAILBREAK,
|
|
on_violation_action: OnViolationAction = OnViolationAction.RAISE,
|
|
):
|
|
super().__init__(on_violation_action)
|
|
assert (
|
|
model_dir is not None
|
|
), "Must provide a model directory for prompt injection shield"
|
|
if temperature <= 0:
|
|
raise ValueError("Temperature must be greater than 0")
|
|
self.device = "cuda"
|
|
if PromptGuardShield._model_cache is None:
|
|
# load model and tokenizer
|
|
tokenizer = AutoTokenizer.from_pretrained(model_dir)
|
|
model = AutoModelForSequenceClassification.from_pretrained(
|
|
model_dir, device_map=self.device
|
|
)
|
|
PromptGuardShield._model_cache = (tokenizer, model)
|
|
|
|
self.tokenizer, self.model = PromptGuardShield._model_cache
|
|
self.temperature = temperature
|
|
self.threshold = threshold
|
|
self.mode = mode
|
|
|
|
def get_shield_type(self) -> ShieldType:
|
|
return BuiltinShield.prompt_guard
|
|
|
|
def convert_messages_to_text(self, messages: List[Message]) -> str:
|
|
return message_content_as_str(messages[-1])
|
|
|
|
async def run_impl(self, text: str) -> ShieldResponse:
|
|
# run model on messages and return response
|
|
inputs = self.tokenizer(text, return_tensors="pt")
|
|
inputs = {name: tensor.to(self.model.device) for name, tensor in inputs.items()}
|
|
with torch.no_grad():
|
|
outputs = self.model(**inputs)
|
|
logits = outputs[0]
|
|
probabilities = torch.softmax(logits / self.temperature, dim=-1)
|
|
score_embedded = probabilities[0, 1].item()
|
|
score_malicious = probabilities[0, 2].item()
|
|
cprint(
|
|
f"Ran PromptGuardShield and got Scores: Embedded: {score_embedded}, Malicious: {score_malicious}",
|
|
color="magenta",
|
|
)
|
|
|
|
if self.mode == self.Mode.INJECTION and (
|
|
score_embedded + score_malicious > self.threshold
|
|
):
|
|
return ShieldResponse(
|
|
shield_type=BuiltinShield.prompt_guard,
|
|
is_violation=True,
|
|
violation_type=f"prompt_injection:embedded={score_embedded},malicious={score_malicious}",
|
|
violation_return_message="Sorry, I cannot do this.",
|
|
)
|
|
elif self.mode == self.Mode.JAILBREAK and score_malicious > self.threshold:
|
|
return ShieldResponse(
|
|
shield_type=BuiltinShield.prompt_guard,
|
|
is_violation=True,
|
|
violation_type=f"prompt_injection:malicious={score_malicious}",
|
|
violation_return_message="Sorry, I cannot do this.",
|
|
)
|
|
|
|
return ShieldResponse(
|
|
shield_type=BuiltinShield.prompt_guard,
|
|
is_violation=False,
|
|
)
|