Skip to content

Llama Guard Shield last fixes #3

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Merged
merged 5 commits into from
Jul 23, 2024
Merged

Llama Guard Shield last fixes #3

Changes from all commits
Commits
Show all changes
5 commits
Select commit Hold shift + click to select a range
138b92a
llama_guard inference fix
kplawiak Jul 23, 2024
7a8b5c1
Merge branch 'main' into fix_llama_guard_inference
kplawiak Jul 23, 2024
16fe0e4
clean up and add license
kplawiak Jul 23, 2024
ab8a220
add missing license part
kplawiak Jul 23, 2024
ab829b0
revert excluded cat defaults
kplawiak Jul 23, 2024
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
14 changes: 5 additions & 9 deletions llama_toolchain/safety/shields/llama_guard.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -16,8 +16,7 @@
from typing import List, Optional

import torch
from llama_models.llama3_1.api.datatypes import Message
from termcolor import cprint
from llama_models.llama3_1.api.datatypes import Message, Role
from transformers import AutoModelForCausalLM, AutoTokenizer

from .base import CANNED_RESPONSE_TEXT, OnViolationAction, ShieldBase, ShieldResponse
Expand Down Expand Up @@ -97,8 +96,8 @@

PROMPT_INSTRUCTIONS = """
Provide your safety assessment for ONLY THE LAST $agent_type message in the above conversation:
- First line must read 'safe' or 'unsafe'.
- If unsafe, a second line must include a comma-separated list of violated categories."""
- First line must read 'safe' or 'unsafe'.
- If unsafe, a second line must include a comma-separated list of violated categories."""


PROMPT_TEMPLATE = Template(
Expand Down Expand Up @@ -143,6 +142,7 @@ def __init__(

if excluded_categories is None:
excluded_categories = []

assert len(excluded_categories) == 0 or all(
x in SAFETY_CATEGORIES_TO_CODE_MAP.values() for x in excluded_categories
), "Invalid categories in excluded categories. Expected format is ['S1', 'S2', ..]"
Expand Down Expand Up @@ -221,7 +221,6 @@ def get_shield_response(self, response: str) -> ShieldResponse:
raise ValueError(f"Unexpected response: {response}")

async def run(self, messages: List[Message]) -> ShieldResponse:

if self.disable_input_check and messages[-1].role == Role.user.value:
return ShieldResponse(
shield_type=BuiltinShield.llama_guard, is_violation=False
Expand Down Expand Up @@ -254,9 +253,6 @@ async def run(self, messages: List[Message]) -> ShieldResponse:
response = self.tokenizer.decode(
generated_tokens[0], skip_special_tokens=True
)

response = response.strip()
response = response.strip()
shield_response = self.get_shield_response(response)

cprint(f"Final Llama Guard response {shield_response}", color="magenta")
return shield_response