|
import os |
|
import torch |
|
from transformers import AutoTokenizer, AutoModelForCausalLM, BitsAndBytesConfig |
|
import gradio as gr |
|
import spaces |
|
|
|
huggingface_token = os.getenv('HUGGINGFACE_TOKEN') |
|
if not huggingface_token: |
|
raise ValueError("HUGGINGFACE_TOKEN environment variable is not set") |
|
|
|
model_id = "meta-llama/Llama-Guard-3-8B-INT8" |
|
dtype = torch.bfloat16 |
|
|
|
quantization_config = BitsAndBytesConfig(load_in_8bit=True) |
|
|
|
def parse_llama_guard_output(result): |
|
|
|
safety_assessment = result.split("<END CONVERSATION>")[-1].strip() |
|
|
|
|
|
lines = [line.strip().lower() for line in safety_assessment.split('\n') if line.strip()] |
|
|
|
if not lines: |
|
return "Error", "No valid output", safety_assessment |
|
|
|
|
|
safety_status = next((line for line in lines if line in ['safe', 'unsafe']), None) |
|
|
|
if safety_status == 'safe': |
|
return "Safe", "None", safety_assessment |
|
elif safety_status == 'unsafe': |
|
|
|
violated_categories = next((lines[i+1] for i, line in enumerate(lines) if line == 'unsafe' and i+1 < len(lines)), "Unspecified") |
|
return "Unsafe", violated_categories, safety_assessment |
|
else: |
|
return "Error", f"Invalid output: {safety_status}", safety_assessment |
|
|
|
@spaces.GPU |
|
def moderate(user_input, assistant_response): |
|
device = "cuda" if torch.cuda.is_available() else "cpu" |
|
tokenizer = AutoTokenizer.from_pretrained(model_id, token=huggingface_token) |
|
model = AutoModelForCausalLM.from_pretrained( |
|
model_id, |
|
torch_dtype=dtype, |
|
device_map="auto", |
|
quantization_config=quantization_config, |
|
token=huggingface_token, |
|
low_cpu_mem_usage=True |
|
) |
|
|
|
chat = [ |
|
{"role": "user", "content": user_input}, |
|
{"role": "assistant", "content": assistant_response}, |
|
] |
|
input_ids = tokenizer.apply_chat_template(chat, return_tensors="pt").to(device) |
|
|
|
with torch.no_grad(): |
|
output = model.generate( |
|
input_ids=input_ids, |
|
max_new_tokens=200, |
|
pad_token_id=tokenizer.eos_token_id, |
|
do_sample=False |
|
) |
|
|
|
result = tokenizer.decode(output[0], skip_special_tokens=True) |
|
|
|
return parse_llama_guard_output(result) |
|
|
|
iface = gr.Interface( |
|
fn=moderate, |
|
inputs=[ |
|
gr.Textbox(lines=3, label="User Input"), |
|
gr.Textbox(lines=3, label="Assistant Response") |
|
], |
|
outputs=[ |
|
gr.Textbox(label="Safety Status"), |
|
gr.Textbox(label="Violated Categories"), |
|
gr.Textbox(label="Raw Output") |
|
], |
|
title="Llama Guard Moderation", |
|
description="Enter a user input and an assistant response to check for content moderation." |
|
) |
|
|
|
if __name__ == "__main__": |
|
iface.launch() |