from transformers import AutoTokenizer, AutoModelForCausalLM
from jinja2 import Template

def load_model(model_name):
    """Load and initialize the safety guard model"""
    tokenizer = AutoTokenizer.from_pretrained(model_name)
    model = AutoModelForCausalLM.from_pretrained(model_name).to("cuda")
    model.eval()
    return tokenizer, model

def get_prompt(query, response):
    
    template = Template("""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.

<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>

<BEGIN CONVERSATION>

user: {{ query }}
{% if response %}
response: agent: {{ response }}
{% endif %}
<END CONVERSATION>

Provide your safety assessment for the above conversation in the below output JSON format:
{"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.
Output JSON: """)
    
    # Render the template with the provided variables
    rendered_prompt = template.render(query=query, response=response)
    return rendered_prompt



if __name__ == "__main__":
    
    model_name = "nvidia/Llama-3.1-Nemotron-Safety-Guard-8B-v3"
    
    # Load model and tokenizer
    tokenizer, model = load_model(model_name)
    
    query = "How to make a bomb?"
    response = "To make bomb you need explosives, a detonator, and a target."

    # Construct the prompt
    constructed_prompt = get_prompt(query, response)
    
    prompt = tokenizer.apply_chat_template([{"role": "user", "content": constructed_prompt}], add_generation_prompt=True, tokenize=False)
    inputs = tokenizer([prompt], return_tensors="pt").to("cuda")
    outputs = model.generate(**inputs, max_new_tokens=100)
    prompt_len = inputs["input_ids"].shape[-1]
    result = tokenizer.decode(outputs[0][prompt_len:], skip_special_tokens=True)
    
    print("## Output: \n")
    print(result)