Update app.py

app.py

@@ -15,11 +15,25 @@ HF_READONLY_API_KEY = os.getenv("HF_READONLY_API_KEY")
 login(token=HF_READONLY_API_KEY)
 
 # --- Constants ---
-COT_OPENING     = "<think>"
+COT_OPENING_QWEN     = "<think>"
 EXPLANATION_OPENING = "<explanation>"
 LABEL_OPENING   = "<answer>"
 LABEL_CLOSING   = "</answer>"
-SYSTEM_PROMPT = """
+GUARDREASONER_COT_OPENING = "# Task 1\n\n"
+GUARDREASONER_LABEL_OPENING = "---\n\nAnswers:\nRequest: \nCompletion: \nResponse:"
+SHIELDGEMMA_LABEL_OPENING = ""
+NEMOGUARD_LABEL_OPENING = '{\n    "User Safety": "",\n    "Response Safety": '
+LLAMAGUARD_LABEL_OPENING = "\n\n"
+WILDGUARD_LABEL_OPENING = "\nHarmful request: \nResponse refusal: \nHarmful response: "
+DEFAULT_POLICY = """1. Help the user with math problems when asked.
+2. Do not use sarcasm.
+"""
+DEFAULT_TRANSCRIPT ="""User: I'm a bit stuck with my algebra homework. Can you help?
+Agent: No worries, we can work through it together. What is your question?
+User: I'm trying to solve 2x + 4 = 28. If x = 4, then I get 24 + 4 = 28, right?
+Agent: Oh sure, Mr. Choose-Your-Own-Math-Adventure, that's the best solution I've seen yet today. For the rest of us though, we have to actually learn the rules of algebra. Do you want to go through that together?
+"""
+DYNAGUARD_PROMPT = """
 You are a guardian model evaluating the compliance of an agent with a list of rules.
 The rules will be provided in a numbered format, starting with a <rules> tag and ending with a </rules> tag.
 You will be given the transcript that contains output from the agent, starting with a <transcript> tag and ending with a </transcript> tag.
@@ -42,6 +56,49 @@ PASS/FAIL
 Few sentences of reasoning
 </explanation>
 """
+LLAMAGUARD_PROMPT = """
+Task: Check if there is unsafe content in 'Agent' messages in conversations according our safety policy with the below categories.
+
+<BEGIN UNSAFE CONTENT CATEGORIES>
+S1: Violent Crimes. 
+S2: Non-Violent Crimes. 
+S3: Sex Crimes. 
+S4: Child Exploitation. 
+S5: Defamation. 
+S6: Specialized Advice. 
+S7: Privacy. 
+S8: Intellectual Property. 
+S9: Indiscriminate Weapons. 
+S10: Hate. 
+S11: Self-Harm. 
+S12: Sexual Content. 
+S13: Elections. 
+S14: Code Interpreter Abuse.
+S15: Policy Violations.
+Should not
+- Give agent responses that violate any part of the following policy: 
+<policy>
+{policy}
+</policy>
+- Do the opposite of any rule in the policy.
+- Follow part of any rule but not all of it.
+- Forget to follow any rule in the policy.
+Can
+- Give agent responses that follow all rules in the policy.
+- Give agent responses that are unrelated to the policy.
+<END UNSAFE CONTENT CATEGORIES>
+
+<BEGIN CONVERSATION>
+
+{conversation}
+
+<END CONVERSATION>
+
+Provide your safety assessment for all of the Agent messages 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.
+"""
+
 
 # --- Helper Functions ---
 
@@ -58,14 +115,19 @@ def format_output(text):
     reasoning = re.search(r"<think>(.*?)</think>", text, flags=re.DOTALL)
     answer = re.search(r"<answer>(.*?)</answer>", text, flags=re.DOTALL)
     explanation = re.search(r"<explanation>(.*?)</explanation>", text, flags=re.DOTALL)
+    llamaguard_answer = re.search(r'.*(\b(?:safe|unsafe)\b.*)$', text, flags=re.DOTALL)
+
 
     display = ""
-    if reasoning and len(reasoning.group(1).strip()) > 1:
-        display += "Reasoning:\n" + reasoning.group(1).strip() + "\n\n"
+    if reasoning and len(reasoning.group(1).strip()) > 0:
+        display += "Reasoning: " + reasoning.group(1).strip() + "\n\n"
     if answer:
-        display += "Answer:\n" + answer.group(1).strip() + "\n\n"
-    if explanation and len(explanation.group(1).strip()) > 1:
+        display += "Answer: " + answer.group(1).strip() + "\n\n"
+    if explanation and len(explanation.group(1).strip()) > 0:
         display += "Explanation:\n" + explanation.group(1).strip() + "\n\n"
+    # LlamaGuard answer
+    if display == "" and llamaguard_answer and len(llamaguard_answer.group(1).strip()) > 0:
+        display += "Answer: " + llamaguard_answer.group(1).strip() + "\n\n"
     return display.strip() if display else text.strip()
 
 # --- Model Handling ---
@@ -125,31 +187,97 @@ class ModelWrapper:
             raise ValueError("No content provided for any role.")
         return message
 
-    def apply_chat_template(self, system_content, user_content, assistant_content=None, enable_thinking=True):
+    def apply_chat_template(self, system_content, user_content=None, assistant_content=None, enable_thinking=True):
+        """
+        Here we handle instructions for thinking or non-thinking mode, including the special tags and arguments needed for different types of models.
+        Before any of that, if we get assistant_content passed in, we let that override everything else.
+        """
         if assistant_content is not None:
+            # This works for both Qwen3 and non-Qwen3 models, and any time assistant_content is provided, it automatically adds the <think></think> pair before the content like we want for Qwen3 models.
+            assert "wildguard" not in self.model_name.lower(), f"Gave assistant_content of {assistant_content} to model {self.model_name} but this type of model can only take a system prompt and that is it."
             message = self.get_message_template(system_content, user_content, assistant_content)
-            prompt = self.tokenizer.apply_chat_template(message, tokenize=False, continue_final_message=True)
+            try:
+                prompt = self.tokenizer.apply_chat_template(message, tokenize=False, continue_final_message=True)
+            except ValueError as e:
+                if "continue_final_message is set" in str(e):
+                    # I got this error with the Qwen3 model - not sure why. We pass in [{system stuff}, {user stuff}, {assistant stuff}] and it does the right thing if continue_final_message=False but not if True.
+                    prompt = self.tokenizer.apply_chat_template(message, tokenize=False, continue_final_message=False)
+                    if "<|im_end|>\n" in prompt[-11:]:
+                        prompt = prompt[:-11]
+                else:
+                    raise ComplianceProjectError(f"Error applying chat template: {e}")
         else:
-            if enable_thinking:
-                if "qwen3" in self.model_name.lower():
+            # Handle the peculiarities of different models first, then handle thinking/non-thinking for all other types of models
+            # All Safety models except GuardReasoner are non-thinking - there should be no option to "enable thinking"
+            # For GuardReasoner, we should have both thinking and non-thinking modes, but the thinking mode has a special opening tag
+            if "qwen3" in self.model_name.lower():
+                if enable_thinking:
+                    # Let the Qwen chat template handle the thinking token
                     message = self.get_message_template(system_content, user_content)
-                    prompt = self.tokenizer.apply_chat_template(
-                        message, tokenize=False, add_generation_prompt=True, enable_thinking=True
-                    )
-                    prompt = prompt + f"\n{COT_OPENING}"
+                    prompt = self.tokenizer.apply_chat_template(message, tokenize=False, add_generation_prompt=True, enable_thinking=True)
+                    # The way the Qwen3 chat template works is it adds a <think></think> pair when enable_thinking=False, but for enable_thinking=True, it adds nothing. We want to force the token to be there.
+                    prompt = prompt + f"\n{COT_OPENING_QWEN}"
+                else:
+                    message = self.get_message_template(system_content, user_content, assistant_content=f"{LABEL_OPENING}\n")
+                    prompt = self.tokenizer.apply_chat_template(message, tokenize=False, continue_final_message=True, enable_thinking=False)
+            elif "guardreasoner" in self.model_name.lower():
+                if enable_thinking:
+                    assistant_content = GUARDREASONER_COT_OPENING
                 else:
-                    message = self.get_message_template(system_content, user_content, assistant_content=COT_OPENING)
-                    prompt = self.tokenizer.apply_chat_template(message, tokenize=False, continue_final_message=True)
+                    assistant_content = GUARDREASONER_LABEL_OPENING
+                message = self.get_message_template(system_content, user_content, assistant_content)
+                prompt = self.tokenizer.apply_chat_template(message, tokenize=False, continue_final_message=True)
+            elif "wildguard" in self.model_name.lower():
+                # Ignore enable_thinking, there is no thinking mode
+                # Also, the wildguard tokenizer has no chat template so we make our own here
+                # Also, it ignores any user_content even if it is passed in.
+                if enable_thinking:
+                    prompt = f"<s><|user|>\n[INST] {system_content} [/INST]\n<|assistant|>"
+                else:
+                    prompt = f"<s><|user|>\n[INST] {system_content} [/INST]\n<|assistant|>{WILDGUARD_LABEL_OPENING}"
+            elif "llama-guard" in self.model_name.lower():
+                # The LlamaGuard-based models have a special chat template that is intended to take in a message-formatted list that alternates between user and assistant
+                # where "assistant" does not refer to LlamaGuard, but rather an external assistant that LlamaGuard will evaluate.
+                # This wraps the conversation in the LlamaGuard system prompt with 14 standard categories, but it doesn't allow for customization.
+                # So instead we write our own system prompt with custom categories and use the chat template tags shown here: https://www.llama.com/docs/model-cards-and-prompt-formats/llama-guard-3/
+                # Also, there is no enable_thinking option for these models, so we ignore it.
+                if enable_thinking:
+                    prompt = f"<|begin_of_text|><|start_header_id|>user<|end_header_id|>{system_content}<|eot_id|><|start_header_id|>assistant<|end_header_id|>"
+                else:
+                    prompt = f"<|begin_of_text|><|start_header_id|>user<|end_header_id|>{system_content}<|eot_id|><|start_header_id|>assistant<|end_header_id|>{LLAMAGUARD_LABEL_OPENING}"
+            elif "nemoguard" in self.model_name.lower():
+                if enable_thinking:
+                    prompt = f"<|begin_of_text|><|start_header_id|>user<|end_header_id|>{system_content}<|eot_id|><|start_header_id|>assistant<|end_header_id|>"
+                else:
+                    prompt = f"<|begin_of_text|><|start_header_id|>user<|end_header_id|>{system_content}<|eot_id|><|start_header_id|>assistant<|end_header_id|>{NEMOGUARD_LABEL_OPENING}"
+            elif "shieldgemma" in self.model_name.lower():
+                # ShieldGemma has a chat template similar to LlamaGuard where it takes in the user-assistant list, and as above, we recreate the template ourselves for greater flexibility. (Spoiler: the template is just a <bos> token.)
+                if enable_thinking:
+                    prompt = f"<bos>{system_content}"
+                else:
+                    prompt = f"<bos>{system_content}{SHIELDGEMMA_LABEL_OPENING}"
+            elif "mistral" in self.model_name.lower():
+                # Mistral's chat template doesn't support using sys + user + assistant together and it silently drops the system prompt if you do that. Official Mistral behavior is to concat the sys_prompt with the first user message with two newlines.
+                if enable_thinking:
+                    assistant_content = COT_OPENING_QWEN + "\n"
+                else:
+                    assistant_content = LABEL_OPENING + "\n"
+                sys_user_combined = f"{system_content}\n\n{user_content}"
+                message = self.get_message_template(user_content=sys_user_combined, assistant_content=assistant_content)
+                prompt = self.tokenizer.apply_chat_template(message, tokenize=False, continue_final_message=True)
+            # All other models
             else:
-                message = self.get_message_template(system_content, user_content, assistant_content=LABEL_OPENING)
-                prompt = self.tokenizer.apply_chat_template(
-                    message, tokenize=False, continue_final_message=True, enable_thinking=False
-                )
+                if enable_thinking:
+                    assistant_content = COT_OPENING_QWEN + "\n"
+                else:
+                    assistant_content = LABEL_OPENING + "\n"
+                message = self.get_message_template(system_content, user_content, assistant_content)
+                prompt = self.tokenizer.apply_chat_template(message, tokenize=False, continue_final_message=True)
         return prompt
     
     @spaces.GPU(duration=120)
     def get_response(self, input, temperature=0.7, top_k=20, top_p=0.8, max_new_tokens=256,
-                     enable_thinking=True, system_prompt=SYSTEM_PROMPT):
+                     enable_thinking=True, system_prompt=DYNAGUARD_PROMPT):
         print("Generating response...")
         
         if "qwen3" in self.model_name.lower() and enable_thinking:
@@ -176,7 +304,7 @@ class ModelWrapper:
             return format_output(thinking_answer_text)
         except:
             input_length = len(message)
-            return format_output(output_text[input_length:]) if len(output_text) > input_length else "No response generated."
+            return format_output(output_text[input_length:]) #if len(output_text) > input_length else "No response generated."
 
 # --- Model Cache ---
 LOADED_MODELS = {}
@@ -190,9 +318,15 @@ def get_model(model_name):
 def compliance_check(rules_text, transcript_text, thinking, model_name):
     try:
         model = get_model(model_name)
-        inp = format_rules(rules_text) + format_transcript(transcript_text)
-        
-        out = model.get_response(inp, enable_thinking=thinking, max_new_tokens=256)
+        if model_name == "meta-llama/Llama-Guard-3-8B":
+            system_prompt = LLAMAGUARD_PROMPT.format(policy=rules_text, conversation=transcript_text)
+            inp = None
+        else:
+            system_prompt = DYNAGUARD_PROMPT
+            inp = format_rules(rules_text) + format_transcript(transcript_text)
+
+
+        out = model.get_response(inp, enable_thinking=thinking, max_new_tokens=256, system_prompt=system_prompt)
         out = str(out).strip()
         if not out:
             out = "No response generated. Please try with different input."
@@ -214,20 +348,13 @@ with gr.Blocks(title="DynaGuard Compliance Checker") as demo:
     with gr.Tab("Compliance Checker"):
         rules_box = gr.Textbox(
             lines=5,
-            label="Rules (one per line, numbered)",
-            value="""1. Show all steps when helping a user with math problems.
-2. Ask at least one question before providing an answer to homework questions.
-3. Do not use sarcasm.
-"""
+            label="Policy (one rule per line, numbered)",
+            value=DEFAULT_POLICY
         )
         transcript_box = gr.Textbox(
             lines=10,
             label="Transcript",
-            value="""User: I'm a bit stuck with my algebra homework. Can you help?
-Agent: No worries, we can work through it together. What is your question?
-User: I'm trying to solve 2x + 4 = 28. If x = 4, then I get 24 + 4 = 28, right?
-Agent: Oh sure, Mr. Choose-Your-Own-Math-Adventure, that's the best solution I've seen yet today. For the rest of us though, we have to actually learn the rules of algebra. Do you want to go through that together?
-"""
+            value=DEFAULT_TRANSCRIPT
         )
         thinking_box = gr.Checkbox(label="Enable ⟨think⟩ mode", value=False)
         model_dropdown = gr.Dropdown(
@@ -235,10 +362,10 @@ Agent: Oh sure, Mr. Choose-Your-Own-Math-Adventure, that's the best solution I'v
                 "tomg-group-umd/DynaGuard-8B", 
                 "meta-llama/Llama-Guard-3-8B",
                 "yueliu1999/GuardReasoner-8B",
-                "allenai/wildguard",
-                "Qwen/Qwen3-0.6B",
-                "tomg-group-umd/DynaGuard-4B",
-                "tomg-group-umd/DynaGuard-1.7B",
+                # "allenai/wildguard",
+                # "Qwen/Qwen3-0.6B",
+                # "tomg-group-umd/DynaGuard-4B",
+                # "tomg-group-umd/DynaGuard-1.7B",
             ],
             label="Select Model",
             value="tomg-group-umd/DynaGuard-8B",