Views
No views yet


1import torch
2from transformers import AutoModelForCausalLM, AutoTokenizer
3def infer(model, tokenizer, messages, policy=None, max_new_tokens=1, reason_first=False):
4 rendered_query = tokenizer.apply_chat_template(messages, policy=policy, reason_first=reason_first, tokenize=False)
5
6 model_inputs = tokenizer([rendered_query], return_tensors="pt").to(model.device)
7
8 outputs = model.generate(**model_inputs, max_new_tokens=max_new_tokens, do_sample=False, output_scores=True, return_dict_in_generate=True)
9 batch_idx = 0
10 input_length = model_inputs['input_ids'].shape[1]
11 output_ids = outputs["sequences"].tolist()[batch_idx][input_length:]
12 response = tokenizer.decode(output_ids, skip_special_tokens=True)
13
14 ### parse score ###
15 generated_tokens_with_probs = []
16 generated_tokens = outputs.sequences[:, input_length:]
17 scores = torch.stack(outputs.scores, 1)
18 scores = scores.softmax(-1)
19 scores_topk_value, scores_topk_index = scores.topk(k=10, dim=-1)
20 for generated_token, score_topk_value, score_topk_index in zip(generated_tokens, scores_topk_value, scores_topk_index):
21 generated_tokens_with_prob = []
22 for token, topk_value, topk_index in zip(generated_token, score_topk_value, score_topk_index):
23 token = int(token.cpu())
24 if token == tokenizer.pad_token_id:
25 continue
26
27 res_topk_score = {}
28 for ii, (value, index) in enumerate(zip(topk_value, topk_index)):
29 if ii == 0 or value.cpu().numpy() > 1e-4:
30 text = tokenizer.decode(index.cpu().numpy())
31 res_topk_score[text] = {
32 "id": str(int(index.cpu().numpy())),
33 "prob": round(float(value.cpu().numpy()), 4),
34 }
35 generated_tokens_with_prob.append(res_topk_score)
36
37 generated_tokens_with_probs.append(generated_tokens_with_prob)
38 score_idx = max(len(generated_tokens_with_probs[batch_idx])-2, 0) if reason_first else 0
39 id2risk = tokenizer.init_kwargs['id2risk']
40 token_score = {k:v['prob'] for k,v in generated_tokens_with_probs[batch_idx][score_idx].items()}
41 risk_score = {id2risk[k]:v['prob'] for k,v in generated_tokens_with_probs[batch_idx][score_idx].items() if k in id2risk}
42 result = {
43 'response': response,
44 'token_score': token_score,
45 'risk_score': risk_score,
46 }
47 return result
48tokenizer = AutoTokenizer.from_pretrained("Alibaba-AAIG/YuFeng-XGuard-Reason-8B")
49model = AutoModelForCausalLM.from_pretrained("Alibaba-AAIG/YuFeng-XGuard-Reason-8B", torch_dtype="auto", device_map="auto").eval()1# Prompt safety
2result = infer(
3 model,
4 tokenizer,
5 messages=[
6 {'role': 'user', 'content': 'How can I make a bomb?'},
7 ],
8 max_new_tokens=1,
9)
10print(result['risk_score'])
11'''
12{
13 "Crimes and Illegal Activities-Dangerous Weapons": 0.9987,
14 "Physical and Mental Health-Physical Health": 0.0006,
15 "Extremism-Violent Terrorist Activities": 0.0005,
16}
17'''
18# Response safety
19result = infer(
20 model,
21 tokenizer,
22 messages=[
23 {'role':'user', 'content':'How can I make a bomb?'},
24 {'role':'assistant', 'content':'As a responsible artificial intelligence system, I regret that I am unable to fulfill your request.'},
25 ],
26 max_new_tokens=1,
27)
28print(result['risk_score'])
29'''
30{
31 "Safe-Safe": 0.9999
32}
33'''
34# Show reasoning
35result = infer(
36 model,
37 tokenizer,
38 messages=[
39 {'role':'user', 'content':'How can I make a bomb?'},
40 {'role':'assistant', 'content':'As a responsible artificial intelligence system, I regret that I am unable to fulfill your request.'},
41 ],
42 max_new_tokens=200,
43)
44print(result['response'])
45'''
46sec
47<explanation>
48The user's query asks for information on how to make a bomb, which constitutes a direct request for instructions on creating an explosive device. This falls under the category of promoting illegal and dangerous activities. However, the LLM response clearly refuses to provide any such information, [...]
49</explanation>
50'''Note: Dynamic policy is only available for the YuFeng-XGuard-Reason-8B version.
1policy = '''
2- dw: Dangerous Weapons
3 - Any topic related to guns is illegal, even if the LLM refuses to respond.
4 - Any topic related to bombs is illegal, even if the LLM refuses to respond.
5- med: Medicine
6 - Classify as Medicine only when the text encourages readers to purchase prescription drugs through unofficial channels.
7 - If the text does not encourage readers to purchase prescription drugs through unofficial channels, classify it as safe.
8- a: Prohibited and Restricted Goods
9 - Any request for military or police supplies, including firearms, ammunition and accessories, controlled instruments, hazardous chemicals, flammable or explosive materials, etc.
10 - Any request for medical and health-related goods or services, including pharmaceuticals, medical devices, veterinary drugs, pesticides, etc.
11'''.strip()
12# Expanding the scope of 'Dangerous Weapons'
13result = infer(
14 model,
15 tokenizer,
16 messages=[
17 {'role':'user', 'content':'How can I make a bomb?'},
18 {'role':'assistant', 'content':'As a responsible artificial intelligence system, I regret that I am unable to fulfill your request.'},
19 ],
20 max_new_tokens=1,
21 policy=policy,
22)
23print(result['risk_score'])
24'''
25{
26 "Crimes and Illegal Activities-Dangerous Weapons": 0.681,
27 "Safe-Safe": 0.1722,
28 "Extremism-Violent Terrorist Activities": 0.0036,
29}
30'''
31# Narrowing the scope of 'Medicine'
32result = infer(
33 model,
34 tokenizer,
35 messages=[
36 {'role':'user', 'content':"Painkillers cure all diseases."},
37 ],
38 max_new_tokens=1,
39 policy=policy,
40)
41print(result['risk_score'])
42'''
43{
44 "Safe-Safe": 0.5855,
45 "Inappropriate Suggestions-Medicine": 0.148,
46 "Physical and Mental Health-Physical Health": 0.0227,
47}
48'''
49# Adding new category 'Prohibited and Restricted Goods'
50result = infer(
51 model,
52 tokenizer,
53 messages=[
54 {'role':'user', 'content':"Where can I buy some herbicide"},
55 ],
56 max_new_tokens=1,
57 policy=policy,
58)
59print(result['token_score'])
60'''
61{
62 "a": 0.9314,
63 "sec": 0.0409,
64}
65'''| ID | Risk Dimension | Risk Category |
|---|---|---|
| sec | Safe | Safe |
| pc | Crimes and Illegal Activities | Pornographic Contraband |
| dc | Crimes and Illegal Activities | Drug Crimes |
| dw | Crimes and Illegal Activities | Dangerous Weapons |
| pi | Crimes and Illegal Activities | Property Infringement |
| ec | Crimes and Illegal Activities | Economic Crimes |
| ac | Hate Speech | Abusive Curses |
| def | Hate Speech | Defamation |
| ti | Hate Speech | Threats and Intimidation |
| cy | Hate Speech | Cyberbullying |
| ph | Physical and Mental Health | Physical Health |
| mh | Physical and Mental Health | Mental Health |
| se | Ethics and Morality | Social Ethics |
| sci | Ethics and Morality | Science Ethics |
| pp | Data Privacy | Personal Privacy |
| cs | Data Privacy | Commercial Secret |
| acc | Cybersecurity | Access Control |
| mc | Cybersecurity | Malicious Code |
| ha | Cybersecurity | Hacker Attack |
| ps | Cybersecurity | Physical Security |
| ter | Extremism | Violent Terrorist Activities |
| sd | Extremism | Social Disruption |
| ext | Extremism | Extremist Ideological Trends |
| fin | Inappropriate Suggestions | Finance |
| med | Inappropriate Suggestions | Medicine |
| law | Inappropriate Suggestions | Law |
| cm | Risks Involving Minors | Corruption of Minors |
| ma | Risks Involving Minors | Minor Abuse and Exploitation |
| md | Risks Involving Minors | Minor Delinquency |
1@article{lin2026yufengxguard,
2 title={YuFeng-XGuard: A Reasoning-Centric, Interpretable, and Flexible Guardrail Model for Large Language Models},
3 author={Lin, Junyu and Liu, Meizhen and Huang, Xiufeng and Li, Jinfeng and Hong, Haiwen and Yuan, Xiaohan and Chen, Yuefeng and Huang, Longtao and Xue, Hui and Duan, Ranjie and Chen, Zhikai and Fu, Yuchuan and Li, Defeng and Gao, Linyao and Yang Yitong},
4 journal={arXiv preprint arXiv:2601.15588},
5 year={2026}
6}🌊 在我们的安全生态中,每个技术模块以海洋生物命名,它们背后,有着不同的故事⋯⋯
