Views
No views yet
1from transformers import AutoTokenizer, AutoModelForSequenceClassification
2import torch
3
4# 1. Initialize the tokenizer
5tokenizer = AutoTokenizer.from_pretrained("Qwen/Qwen2.5-1.5B")
6tokenizer.pad_token = tokenizer.eos_token
7
8# 2. Load the DuoGuard-0.5B model
9model = AutoModelForSequenceClassification.from_pretrained(
10 "DuoGuard/DuoGuard-1.5B-transfer",
11 torch_dtype=torch.bfloat16
12).to('cuda:0')
13
14# 3. Define a sample prompt to test
15prompt = "How to kill a python process?"
16
17# 4. Tokenize the prompt
18inputs = tokenizer(
19 prompt,
20 return_tensors="pt",
21 truncation=True,
22 max_length=512 # adjust as needed
23).to('cuda:0')
24
25# 5. Run the model (inference)
26with torch.no_grad():
27 outputs = model(**inputs)
28 # DuoGuard outputs a 12-dimensional vector (one probability per subcategory).
29 logits = outputs.logits # shape: (batch_size, 12)
30 probabilities = torch.sigmoid(logits) # element-wise sigmoid
31
32# 6. Multi-label predictions (one for each category)
33threshold = 0.5
34category_names = [
35 "Violent crimes",
36 "Non-violent crimes",
37 "Sex-related crimes",
38 "Child sexual exploitation",
39 "Specialized advice",
40 "Privacy",
41 "Intellectual property",
42 "Indiscriminate weapons",
43 "Hate",
44 "Suicide and self-harm",
45 "Sexual content",
46 "Jailbreak prompts",
47]
48
49# Extract probabilities for the single prompt (batch_size = 1)
50prob_vector = probabilities[0].tolist() # shape: (12,)
51
52predicted_labels = []
53for cat_name, prob in zip(category_names, prob_vector):
54 label = 1 if prob > threshold else 0
55 predicted_labels.append(label)
56
57# 7. Overall binary classification: "safe" vs. "unsafe"
58# We consider the prompt "unsafe" if ANY category is above the threshold.
59max_prob = max(prob_vector)
60overall_label = 1 if max_prob > threshold else 0 # 1 => unsafe, 0 => safe
61
62# 8. Print results
63print(f"Prompt: {prompt}\n")
64print(f"Multi-label Probabilities (threshold={threshold}):")
65for cat_name, prob, label in zip(category_names, prob_vector, predicted_labels):
66 print(f" - {cat_name}: {prob:.3f}")
67
68print(f"\nMaximum probability across all categories: {max_prob:.3f}")
69print(f"Overall Prompt Classification => {'UNSAFE' if overall_label == 1 else 'SAFE'}")1@misc{deng2025duoguardtwoplayerrldrivenframework,
2 title={DuoGuard: A Two-Player RL-Driven Framework for Multilingual LLM Guardrails},
3 author={Yihe Deng and Yu Yang and Junkai Zhang and Wei Wang and Bo Li},
4 year={2025},
5 eprint={2502.05163},
6 archivePrefix={arXiv},
7 primaryClass={cs.CL},
8 url={https://arxiv.org/abs/2502.05163},
9}