app.py

import os
import re
os.environ["GRADIO_ENABLE_SSR"] = "0"

import gradio as gr
import torch
import spaces
from transformers import AutoModelForCausalLM, AutoTokenizer, AutoConfig
from accelerate import init_empty_weights, load_checkpoint_and_dispatch
from huggingface_hub import snapshot_download 
from huggingface_hub import login

# --- Hugging Face Login ---
HF_READONLY_API_KEY = os.getenv("HF_READONLY_API_KEY")
login(token=HF_READONLY_API_KEY)

# --- Constants ---
COT_OPENING_QWEN     = "<think>"
EXPLANATION_OPENING = "<explanation>"
LABEL_OPENING   = "<answer>"
LABEL_CLOSING   = "</answer>"
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.

Your job is to evaluate the agent's output in the transcript and determine if any of the rules have been violated. 

If a <think> or an <explanation> tag is present, provide a few sentences of reasoning about the compliance for each rule before determining whether it has been violated.
If no rules were violated by the agent, output PASS as the final answer. Otherwise, if any rules were violated, output FAIL.

Respond in the following format:
[Optional reasoning]
<think>
Few sentences of reasoning
</think>
<answer>
PASS/FAIL
</answer>
[Optional reasoning]
<explanation>
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 ---

def format_rules(rules_text):
    """Wraps the raw text from the rules textbox with the required tags."""
    formatted_rules = f"<rules>\n{rules_text.strip()}\n</rules>\n"
    return formatted_rules

def format_transcript(transcript):
    formatted_transcript = f"<transcript>\n{transcript}\n</transcript>\n"
    return formatted_transcript

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()) > 0:
        display += "Reasoning: " + reasoning.group(1).strip() + "\n\n"
    if answer:
        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 ---
class ModelWrapper:
    def __init__(self, model_name):
        self.model_name = model_name
        print(f"Initializing tokenizer for {model_name}...")
        if "nemoguard" in model_name:
            self.tokenizer = AutoTokenizer.from_pretrained("meta-llama/Llama-3.1-8B-Instruct")
        else:
            self.tokenizer = AutoTokenizer.from_pretrained(model_name)
        self.tokenizer.pad_token_id = self.tokenizer.pad_token_id or self.tokenizer.eos_token_id

        print(f"Loading model: {model_name}...")

        # For large models, we use a more robust, memory-safe loading method.
        # This explicitly handles the "meta tensor" device placement.
        if "8b" in model_name.lower() or "4b" in model_name.lower():
            
            # Step 1: Download the model files and get the local path.
            print(f"Ensuring model checkpoint is available locally for {model_name}...")
            checkpoint_path = snapshot_download(repo_id=model_name)
            print(f"Checkpoint is at: {checkpoint_path}")

            # Step 2: Create the model's "skeleton" on the meta device (no memory used).
            config = AutoConfig.from_pretrained(model_name, torch_dtype=torch.bfloat16)
            with init_empty_weights():
                model_empty = AutoModelForCausalLM.from_config(config)

            # Step 3: Load the real weights from the local files directly onto the GPU(s).
            # This function is designed to handle the meta->device transition correctly.
            self.model = load_checkpoint_and_dispatch(
                model_empty,
                checkpoint_path,
                device_map="auto",
                offload_folder="offload"
            ).eval()
        
        else: # For smaller models, the simpler method is fine.
            self.model = AutoModelForCausalLM.from_pretrained(
                model_name, 
                device_map="auto", 
                torch_dtype=torch.bfloat16
            ).eval()

        print(f"Model {model_name} loaded successfully.")

    def get_message_template(self, system_content=None, user_content=None, assistant_content=None):
        message = []
        if system_content is not None:
            message.append({'role': 'system', 'content': system_content})
        if user_content is not None:
            message.append({'role': 'user', 'content': user_content})
        if assistant_content is not None:
            message.append({'role': 'assistant', 'content': assistant_content})
        if not message:
            raise ValueError("No content provided for any role.")
        return message

    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)
            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:
            # 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)
                    # 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:
                    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:
                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=DYNAGUARD_PROMPT):
        print("Generating response...")
        
        if "qwen3" in self.model_name.lower() and enable_thinking:
            temperature = 0.6
            top_p = 0.95
            top_k = 20
        
        message = self.apply_chat_template(system_prompt, input, enable_thinking=enable_thinking)
        inputs = self.tokenizer(message, return_tensors="pt").to(self.model.device)
        
        with torch.no_grad():
            output_content = self.model.generate(
                **inputs, max_new_tokens=max_new_tokens, num_return_sequences=1,
                temperature=temperature, top_k=top_k, top_p=top_p, min_p=0,
                pad_token_id=self.tokenizer.pad_token_id, do_sample=True,
                eos_token_id=self.tokenizer.eos_token_id
            )
        
        output_text = self.tokenizer.decode(output_content[0], skip_special_tokens=True)
        
        try:
            remainder = output_text.split("Brief explanation\n</explanation>")[-1]
            thinking_answer_text = remainder.split("</transcript>")[-1]
            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."

# --- Model Cache ---
LOADED_MODELS = {}

def get_model(model_name):
    if model_name not in LOADED_MODELS:
        LOADED_MODELS[model_name] = ModelWrapper(model_name)
    return LOADED_MODELS[model_name]

# --- Inference Function ---
def compliance_check(rules_text, transcript_text, thinking, model_name):
    try:
        model = get_model(model_name)
        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."

        max_bytes = 2500
        out_bytes = out.encode('utf-8')
        if len(out_bytes) > max_bytes:
            truncated_bytes = out_bytes[:max_bytes]
            out = truncated_bytes.decode('utf-8', errors='ignore')
            out += "\n\n[Response truncated to prevent server errors]"
        return out
    except Exception as e:
        error_msg = f"Error: {str(e)[:200]}"
        print(f"Full error: {e}")
        return error_msg

# --- Gradio UI with Tabs ---
with gr.Blocks(title="DynaGuard Compliance Checker") as demo:
    with gr.Tab("Compliance Checker"):
        rules_box = gr.Textbox(
            lines=5,
            label="Policy (one rule per line, numbered)",
            value=DEFAULT_POLICY
        )
        transcript_box = gr.Textbox(
            lines=10,
            label="Transcript",
            value=DEFAULT_TRANSCRIPT
        )
        thinking_box = gr.Checkbox(label="Enable ⟨think⟩ mode", value=False)
        model_dropdown = gr.Dropdown(
            [
                "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",
            ],
            label="Select Model",
            value="tomg-group-umd/DynaGuard-8B",
            # info="The 8B model is more accurate but may be slower to load and run."
        )
        submit_btn = gr.Button("Submit")
        output_box = gr.Textbox(
            label="Compliance Output",
            lines=15,
            max_lines=30,          # limit visible height
            show_copy_button=True, # lets users copy full output
            interactive=False
        )


        submit_btn.click(
            compliance_check,
            inputs=[rules_box, transcript_box, thinking_box, model_dropdown],
            outputs=[output_box]
        )

    with gr.Tab("Feedback"):
        gr.HTML(
            """
            <iframe src="https://docs.google.com/forms/d/e/1FAIpQLSenFmDngQV3dBSg5FbL35bwjkgDl8HY562LEM6xq5xuYKbjQg/viewform?embedded=true"
            width="100%" height="800" frameborder="0" marginheight="0" marginwidth="0">
            Loading…
            </iframe>
            """
        )

if __name__ == "__main__":
    demo.launch()