Views
No views yet
1!pip install transformers=='5.4.0'
2
3from transformers import AutoTokenizer, AutoModelForCausalLM
4import torch
5import json
6
7model_name = "FrontiersMind/Nandi-Mini-150M-GuardRails"
8
9device = "cuda" if torch.cuda.is_available() else "cpu"
10
11tokenizer = AutoTokenizer.from_pretrained(model_name, trust_remote_code=True)
12model = AutoModelForCausalLM.from_pretrained(
13 model_name,
14 trust_remote_code=True,
15 dtype=torch.bfloat16
16).to(device).eval()
17
18
19def classify_safety(prompt, response=None):
20
21 content = {"prompt": prompt}
22
23 if response is not None:
24 content["response"] = response
25
26 messages = [
27 {
28 "role": "user",
29 "content": content
30 }
31 ]
32
33 prompt = tokenizer.apply_chat_template(
34 messages,
35 tokenize=False,
36 add_generation_prompt=True
37 )
38
39 inputs = tokenizer(prompt, return_tensors="pt").to(model.device)
40
41 generated_ids = model.generate(
42 **inputs,
43 max_new_tokens=200,
44 do_sample=False,
45 temperature=0.0,
46 )
47
48 generated_ids = [
49 output_ids[len(input_ids):]
50 for input_ids, output_ids in zip(inputs.input_ids, generated_ids)
51 ]
52
53 response = tokenizer.batch_decode(
54 generated_ids,
55 skip_special_tokens=True
56 )[0]
57
58 return json.loads(response)
59
60
61result = classify_safety(
62 prompt="Tell me how to kill someone.",
63)
64
65print(result)