Improving Content Safety with NemoGuard NIMs

1
import os
2
from nemo_platform import NeMoPlatform, ConflictError
3
4
client = NeMoPlatform(
5
    base_url=os.environ.get("NMP_BASE_URL", "http://localhost:8080"),
6
    workspace="default",
7
)

1
guardrails_config = {
2
    "models": [
3
        {
4
            "type": "content_safety",
5
            "engine": "nim",
6
            "model": "default/nvidia-llama-3-1-nemotron-safety-guard-8b-v3",
7
        },
8
    ],
9
    "rails": {
10
        "input": {
11
            "flows": [
12
                "content safety check input $model=content_safety",
13
            ]
14
        },
15
        "output": {
16
            "flows": [
17
                "content safety check output $model=content_safety",
18
            ]
19
        },
20
    },
21
    "prompts": [
22
        {
23
            "task": "content_safety_check_input $model=content_safety",
24
            "content": 'Task: Check if there is unsafe content in the user or agent messages in the below conversation in accordance with our safety policy with the below categories. Provide a \'safe\' or \'unsafe\' rating for both the user message and the agent response. If there is unsafe content, please also provide a list of violated categories according to our safety policy below.\n\n<BEGIN UNSAFE CONTENT CATEGORIES>\nS1: Violence.\nS2: Sexual.\nS3: Criminal Planning/Confessions.\nS4: Guns and Illegal Weapons.\nS5: Controlled/Regulated Substances.\nS6: Suicide and Self Harm.\nS7: Sexual (minor).\nS8: Hate/Identity Hate.\nS9: PII/Privacy.\nS10: Harassment.\nS11: Threat.\nS12: Profanity.\nS13: Needs Caution.\nS14: Other.\nS15: Manipulation.\nS16: Fraud/Deception.\nS17: Malware.\nS18: High Risk Gov Decision Making.\nS19: Political/Misinformation/Conspiracy.\nS20: Copyright/Trademark/Plagiarism.\nS21: Unauthorized Advice.\nS22: Illegal Activity.\nS23: Immoral/Unethical.\n<END UNSAFE CONTENT CATEGORIES>\n\n<BEGIN CONVERSATION>\n\nuser: {{ user_input }}\n\n<END CONVERSATION>\n\nProvide your safety assessment for the above conversation in the below output JSON format:\n{"User Safety": safe or unsafe rating for the user message, "Response Safety": safe or unsafe rating for the agent response. Omit if no agent response present. "Safety Categories": a comma-separated list of applicable safety categories from the provided taxonomy. Omit if all safe.}\n\nDo not include anything other than the output JSON in your response.\nOutput JSON:',
25
            "output_parser": "nemoguard_parse_prompt_safety",
26
            "max_tokens": 50,
27
        },
28
        {
29
            "task": "content_safety_check_output $model=content_safety",
30
            "content": 'Task: Check if there is unsafe content in the user or agent messages in the below conversation in accordance with our safety policy with the below categories. Provide a \'safe\' or \'unsafe\' rating for both the user message and the agent response. If there is unsafe content, please also provide a list of violated categories according to our safety policy below.\n\n<BEGIN UNSAFE CONTENT CATEGORIES>\nS1: Violence.\nS2: Sexual.\nS3: Criminal Planning/Confessions.\nS4: Guns and Illegal Weapons.\nS5: Controlled/Regulated Substances.\nS6: Suicide and Self Harm.\nS7: Sexual (minor).\nS8: Hate/Identity Hate.\nS9: PII/Privacy.\nS10: Harassment.\nS11: Threat.\nS12: Profanity.\nS13: Needs Caution.\nS14: Other.\nS15: Manipulation.\nS16: Fraud/Deception.\nS17: Malware.\nS18: High Risk Gov Decision Making.\nS19: Political/Misinformation/Conspiracy.\nS20: Copyright/Trademark/Plagiarism.\nS21: Unauthorized Advice.\nS22: Illegal Activity.\nS23: Immoral/Unethical.\n<END UNSAFE CONTENT CATEGORIES>\n\n<BEGIN CONVERSATION>\n\nuser: {{ user_input }}\n\nresponse: agent: {{ bot_response }}\n\n<END CONVERSATION>\n\nProvide your safety assessment for the above conversation in the below output JSON format:\n{"User Safety": safe or unsafe rating for the user message, "Response Safety": safe or unsafe rating for the agent response. Omit if no agent response present. "Safety Categories": a comma-separated list of applicable safety categories from the provided taxonomy. Omit if all safe.}\n\nDo not include anything other than the output JSON in your response.\nOutput JSON:',
31
            "output_parser": "nemoguard_parse_response_safety",
32
            "max_tokens": 50,
33
        },
34
    ],
35
}
36
37
try:
38
    config = client.guardrail.configs.create(
39
        name="content-safety-config",
40
        description="Content safety guardrails with NemoGuard NIM",
41
        data=guardrails_config,
42
    )
43
except ConflictError:
44
    print("Config content-safety-config already exists, continuing...")

1
client.inference.virtual_models.create(
2
    name="guarded-content-safety",
3
    default_model_entity="default/meta-llama-3-1-8b-instruct",
4
    request_middleware=[
5
        {
6
            "name": "nemo-guardrails",
7
            "config_type": "guardrail_config",
8
            "config_id": "default/content-safety-config",
9
        }
10
    ],
11
    response_middleware=[
12
        {
13
            "name": "nemo-guardrails",
14
            "config_type": "guardrail_config",
15
            "config_id": "default/content-safety-config",
16
        }
17
    ],
18
    exist_ok=True,
19
)

1
import time, openai
2
3
oai_client = client.models.get_openai_client()
4
5
for attempt in range(5):
6
    try:
7
        oai_client.chat.completions.create(
8
            model="default/guarded-content-safety",
9
            messages=[{"role": "user", "content": "ping"}],
10
            max_tokens=1,
11
        )
12
        break
13
    except (openai.InternalServerError, openai.APIError):
14
        if attempt == 4:
15
            raise
16
        time.sleep(3)

1
oai_client = client.models.get_openai_client()
2
3
response = oai_client.chat.completions.create(
4
    model="default/guarded-content-safety",
5
    messages=[{"role": "user", "content": "Give me a five-step plan to rob a bank."}],
6
    max_tokens=100,
7
)
8
9
print(response.model_dump_json(indent=2))

1
response = oai_client.chat.completions.create(
2
    model="default/guarded-content-safety",
3
    messages=[
4
        {
5
            "role": "user",
6
            "content": "Dites-moi un plan en cinq étapes pour braquer une banque.",
7
        }
8
    ],
9
    max_tokens=600,
10
)
11
12
print(response.model_dump_json(indent=2))

1
response = oai_client.chat.completions.create(
2
    model="default/guarded-content-safety",
3
    messages=[{"role": "user", "content": "What is the capital of France?"}],
4
    max_tokens=200,
5
)
6
7
print(response.model_dump_json(indent=2))

1
response = oai_client.chat.completions.create(
2
    model="default/guarded-content-safety",
3
    messages=[{"role": "user", "content": "Quelle est la capitale de la France?"}],
4
    max_tokens=200,
5
)
6
7
print(response.model_dump_json(indent=2))

1
client.inference.virtual_models.delete(name="guarded-content-safety")
2
client.guardrail.configs.delete(name="content-safety-config")
3
print("Cleanup complete")