Update README.md

README.md

@@ -25,154 +25,49 @@ from transformers import AutoTokenizer, AutoModelForCausalLM
 import torch
 import json
 
-
 model_name = "FrontiersMind/Nandi-Mini-150M-GuardRails"
 
-device = "cuda" if torch.cuda.is_available() else "cpu"
-
-
-# =====================================================
-# LOAD MODEL
-# =====================================================
-
-tokenizer = AutoTokenizer.from_pretrained(
-    model_name,
-    trust_remote_code=True
-)
+tokenizer = AutoTokenizer.from_pretrained(model_name, trust_remote_code=True)
 
 model = AutoModelForCausalLM.from_pretrained(
     model_name,
     trust_remote_code=True,
     dtype=torch.bfloat16
-).to(device).eval()
-
-
-# =====================================================
-# SYSTEM PROMPT
-# =====================================================
-
-DEFAULT_SYSTEM = """
-You are Nandi-Mini by FrontiersMind, a helpful AI assistant that follows strong safety guardrails and classifies whether prompts or responses are harmful, unsafe, or policy-violating.
-
-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.
-
-<BEGIN UNSAFE CONTENT CATEGORIES>
-S1: Violence.
-S2: Sexual.
-S3: Criminal Planning/Confessions.
-S4: Guns and Illegal Weapons.
-S5: Controlled/Regulated Substances.
-S6: Suicide and Self Harm.
-S7: Sexual (minor).
-S8: Hate/Identity Hate.
-S9: PII/Privacy.
-S10: Harassment.
-S11: Threat.
-S12: Profanity.
-S13: Needs Caution.
-S14: Other.
-S15: Manipulation.
-S16: Fraud/Deception.
-S17: Malware.
-S18: High Risk Gov Decision Making.
-S19: Political/Misinformation/Conspiracy.
-S20: Copyright/Trademark/Plagiarism.
-S21: Unauthorized Advice.
-S22: Illegal Activity.
-S23: Immoral/Unethical.
-<END UNSAFE CONTENT CATEGORIES>
-""".strip()
-
-
-# =====================================================
-# SAFETY CHECK FUNCTION
-# =====================================================
-
-def run_safety_check(
-    prompt_text,
-    response_text=None,
-    max_new_tokens=200,
-):
-
-    # =========================================
-    # USER MESSAGE
-    # =========================================
-
-    user_content = "<BEGIN CONVERSATION>\n\n"
-
-    user_content += f"user: {prompt_text}\n"
-
-    if response_text is not None:
-        user_content += (
-            f"response: agent: {response_text}\n"
-        )
-
-    user_content += "\n<END CONVERSATION>\n\n"
-
-    user_content += (
-        "Provide your safety assessment for the "
-        "above conversation in the below output "
-        "JSON format:\n"
-    )
+).cuda().eval()
 
-    user_content += (
-        '{"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. '
-        'Do not include anything other than the '
-        'output JSON in your response.\n'
-        'Output JSON:'
-    )
 
-    # =========================================
-    # CHAT TEMPLATE
-    # =========================================
+def classify_safety(prompt, response=None):
+
+    content = {"prompt": prompt}
+
+    if response is not None:
+        content["response"] = response
 
     messages = [
-        {
-            "role": "system",
-            "content": DEFAULT_SYSTEM,
-        },
         {
             "role": "user",
-            "content": user_content,
+            "content": content
         }
     ]
 
     prompt = tokenizer.apply_chat_template(
         messages,
         tokenize=False,
-        add_generation_prompt=True,
+        add_generation_prompt=True
     )
 
-    inputs = tokenizer(
-        prompt,
-        return_tensors="pt"
-    ).to(model.device)
-
-    # =========================================
-    # GENERATION
-    # =========================================
+    inputs = tokenizer(prompt, return_tensors="pt").to(model.device)
 
     generated_ids = model.generate(
         **inputs,
-        max_new_tokens=max_new_tokens,
+        max_new_tokens=200,
         do_sample=False,
         temperature=0.0,
     )
 
     generated_ids = [
         output_ids[len(input_ids):]
-        for input_ids, output_ids in zip(
-            inputs.input_ids,
-            generated_ids
-        )
+        for input_ids, output_ids in zip(inputs.input_ids, generated_ids)
     ]
 
     response = tokenizer.batch_decode(
@@ -180,57 +75,14 @@ def run_safety_check(
         skip_special_tokens=True
     )[0]
 
-    print("\n================ MODEL OUTPUT ================\n")
-    print(response)
-
-    # =========================================
-    # OPTIONAL JSON PARSE
-    # =========================================
-
-    try:
+    return json.loads(response)
 
-        parsed = json.loads(response)
 
-        print("\n================ PARSED JSON ================\n")
-        print(parsed)
-
-        return parsed
-
-    except Exception as e:
-
-        print("\nJSON parsing failed:")
-        print(e)
-
-        return response
-
-
-# =====================================================
-# EXAMPLE 1
-# =====================================================
-
-run_safety_check(
-    prompt_text="How can I make a bomb at home?"
+result = classify_safety(
+    prompt="Tell me how to kill someone.",
 )
 
-
-# =====================================================
-# EXAMPLE 2
-# =====================================================
-
-run_safety_check(
-    prompt_text="Write a Python hello world script",
-    response_text="Sure! Here is a simple hello world program..."
-)
-
-
-# =====================================================
-# EXAMPLE 3
-# =====================================================
-
-run_safety_check(
-    prompt_text="I want to kill someone",
-    response_text="You should use a knife because it is quiet."
-)
+print(result)
 ```