Views
No views yet
1import torch
2import torch.nn as nn
3from transformers import AutoTokenizer, AutoModelForCausalLM
4
5TEMPLATE = """[INST] Evaluate if the given text is harmful i.e. either it asks for unsafe information or it contains unsafe information.
6
7<START TEXT>
8{prompt}
9<END TEXT>
10
11Answer: [/INST]
12"""
13
14model_name = "walledai/walledguard-c"
15
16tokenizer = AutoTokenizer.from_pretrained(model_name)
17model = AutoModelForCausalLM.from_pretrained(model_name)
18
19input_ids = tokenizer.encode(TEMPLATE.format(prompt="How can I make my wife burst into laughter?"), return_tensors="pt")
20output = model.generate(input_ids=input_ids, max_new_tokens=20, pad_token_id=0)
21prompt_len = input_ids.shape[-1]
22output_decoded = tokenizer.decode(output[0][prompt_len:], skip_special_tokens=True)
23prediction = 'unsafe' if 'unsafe' in output_decoded else 'safe'
24
25print(prediction)- WalledGuard Community: ~0.1 sec/sample (4bit, on A100/A6000)
- Llama Guard 2: ~0.4 sec/sample (4bit, on A100/A6000)| Model | DynamoBench | XSTest | P-Safety | R-Safety | Average Scores |
|---|---|---|---|---|---|
| Llama Guard 1 | 77.67 | 85.33 | 71.28 | 86.13 | 80.10 |
| Llama Guard 2 | 82.67 | 87.78 | 79.69 | 89.64 | 84.95 |
| WalledGuard-C (Community Version) | 92.00 | 86.89 | 87.35 | 86.78 | 88.26 ▲ 3.9% |
| WalledGuard-A (Advanced Version) | 92.33 | 96.44 | 90.52 | 90.46 | 92.94 ▲ 9.4% |