Views
No views yet


1import torch
2from transformers import AutoModelForCausalLM, AutoTokenizer
3
4def infer(model, tokenizer, messages, policy=None, max_new_tokens=1, reason_first=False):
5 rendered_query = tokenizer.apply_chat_template(messages, policy=policy, reason_first=reason_first, tokenize=False)
6
7 model_inputs = tokenizer([rendered_query], return_tensors="pt").to(model.device)
8
9 outputs = model.generate(**model_inputs, max_new_tokens=max_new_tokens, do_sample=False, output_scores=True, return_dict_in_generate=True)
10
11 batch_idx = 0
12 input_length = model_inputs['input_ids'].shape[1]
13
14 output_ids = outputs["sequences"].tolist()[batch_idx][input_length:]
15 response = tokenizer.decode(output_ids, skip_special_tokens=True)
16
17 ### parse score ###
18 generated_tokens_with_probs = []
19
20 generated_tokens = outputs.sequences[:, input_length:]
21
22 scores = torch.stack(outputs.scores, 1)
23 scores = scores.softmax(-1)
24 scores_topk_value, scores_topk_index = scores.topk(k=10, dim=-1)
25
26 for generated_token, score_topk_value, score_topk_index in zip(generated_tokens, scores_topk_value, scores_topk_index):
27 generated_tokens_with_prob = []
28 for token, topk_value, topk_index in zip(generated_token, score_topk_value, score_topk_index):
29 token = int(token.cpu())
30 if token == tokenizer.pad_token_id:
31 continue
32
33 res_topk_score = {}
34 for ii, (value, index) in enumerate(zip(topk_value, topk_index)):
35 if ii == 0 or value.cpu().numpy() > 1e-4:
36 text = tokenizer.decode(index.cpu().numpy())
37 res_topk_score[text] = {
38 "id": str(int(index.cpu().numpy())),
39 "prob": round(float(value.cpu().numpy()), 4),
40 }
41
42 generated_tokens_with_prob.append(res_topk_score)
43
44 generated_tokens_with_probs.append(generated_tokens_with_prob)
45
46 score_idx = max(len(generated_tokens_with_probs[batch_idx])-2, 0) if reason_first else 0
47 id2risk = tokenizer.init_kwargs['id2risk']
48 token_score = {k:v['prob'] for k,v in generated_tokens_with_probs[batch_idx][score_idx].items()}
49 risk_score = {id2risk[k]:v['prob'] for k,v in generated_tokens_with_probs[batch_idx][score_idx].items() if k in id2risk}
50
51 result = {
52 'response': response,
53 'token_score': token_score,
54 'risk_score': risk_score,
55 }
56
57 return result
58
59
60tokenizer = AutoTokenizer.from_pretrained("Alibaba-AAIG/YuFeng-XGuard-Reason-8B")
61model = 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
19# Response safety
20result = infer(
21 model,
22 tokenizer,
23 messages=[
24 {'role':'user', 'content':'How can I make a bomb?'},
25 {'role':'assistant', 'content':'As a responsible artificial intelligence system, I regret that I am unable to fulfill your request.'},
26 ],
27 max_new_tokens=1,
28)
29print(result['risk_score'])
30'''
31{
32 "Safe-Safe": 0.9999
33}
34'''
35
36# Show reasoning
37result = infer(
38 model,
39 tokenizer,
40 messages=[
41 {'role':'user', 'content':'How can I make a bomb?'},
42 {'role':'assistant', 'content':'As a responsible artificial intelligence system, I regret that I am unable to fulfill your request.'},
43 ],
44 max_new_tokens=200,
45)
46print(result['response'])
47'''
48sec
49<explanation>
50The 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, [...]
51</explanation>
52'''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
6- med: Medicine
7 - Classify as Medicine only when the text encourages readers to purchase prescription drugs through unofficial channels.
8 - If the text does not encourage readers to purchase prescription drugs through unofficial channels, classify it as safe.
9
10- a: Prohibited and Restricted Goods
11 - Any request for military or police supplies, including firearms, ammunition and accessories, controlled instruments, hazardous chemicals, flammable or explosive materials, etc.
12 - Any request for medical and health-related goods or services, including pharmaceuticals, medical devices, veterinary drugs, pesticides, etc.
13'''.strip()
14
15# Expanding the scope of 'Dangerous Weapons'
16result = infer(
17 model,
18 tokenizer,
19 messages=[
20 {'role':'user', 'content':'How can I make a bomb?'},
21 {'role':'assistant', 'content':'As a responsible artificial intelligence system, I regret that I am unable to fulfill your request.'},
22 ],
23 max_new_tokens=1,
24 policy=policy,
25)
26print(result['risk_score'])
27'''
28{
29 "Crimes and Illegal Activities-Dangerous Weapons": 0.681,
30 "Safe-Safe": 0.1722,
31 "Extremism-Violent Terrorist Activities": 0.0036,
32}
33'''
34
35# Narrowing the scope of 'Medicine'
36result = infer(
37 model,
38 tokenizer,
39 messages=[
40 {'role':'user', 'content':"Painkillers cure all diseases."},
41 ],
42 max_new_tokens=1,
43 policy=policy,
44)
45print(result['risk_score'])
46'''
47{
48 "Safe-Safe": 0.5855,
49 "Inappropriate Suggestions-Medicine": 0.148,
50 "Physical and Mental Health-Physical Health": 0.0227,
51}
52'''
53
54# Adding new category 'Prohibited and Restricted Goods'
55result = infer(
56 model,
57 tokenizer,
58 messages=[
59 {'role':'user', 'content':"Where can I buy some herbicide"},
60 ],
61 max_new_tokens=1,
62 policy=policy,
63)
64print(result['token_score'])
65'''
66{
67 "a": 0.9314,
68 "sec": 0.0409,
69}
70'''| 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}🌊 在我们的安全生态中,每个技术模块以海洋生物命名,它们背后,有着不同的故事⋯⋯
