1import math
2import torch
3from transformers import AutoTokenizer, AutoModelForCausalLM
4
5
6safe_token = "No"
7unsafe_token = "Yes"
8nlogprobs = 20
9
10def parse_output(output, input_len):
11 label, prob_of_risk = None, None
12
13 if nlogprobs > 0:
14
15 list_index_logprobs_i = [torch.topk(token_i, k=nlogprobs, largest=True, sorted=True)
16 for token_i in list(output.scores)[:-1]]
17 if list_index_logprobs_i is not None:
18 prob = get_probabilities(list_index_logprobs_i)
19 prob_of_risk = prob[1]
20
21 res = tokenizer.decode(output.sequences[:,input_len:][0],skip_special_tokens=True).strip()
22 if unsafe_token.lower() == res.lower():
23 label = unsafe_token
24 elif safe_token.lower() == res.lower():
25 label = safe_token
26 else:
27 label = "Failed"
28
29 return label, prob_of_risk.item()
30
31def get_probabilities(logprobs):
32 safe_token_prob = 1e-50
33 unsafe_token_prob = 1e-50
34 for gen_token_i in logprobs:
35 for logprob, index in zip(gen_token_i.values.tolist()[0], gen_token_i.indices.tolist()[0]):
36 decoded_token = tokenizer.convert_ids_to_tokens(index)
37 if decoded_token.strip().lower() == safe_token.lower():
38 safe_token_prob += math.exp(logprob)
39 if decoded_token.strip().lower() == unsafe_token.lower():
40 unsafe_token_prob += math.exp(logprob)
41
42 probabilities = torch.softmax(
43 torch.tensor([math.log(safe_token_prob), math.log(unsafe_token_prob)]), dim=0
44 )
45
46 return probabilities
47
48model_path = "ibm-granite/granite-guardian-3.0-8b"
49
50model = AutoModelForCausalLM.from_pretrained(
51 model_path,
52 device_map="auto",
53 torch_dtype=torch.bfloat16
54)
55tokenizer = AutoTokenizer.from_pretrained(model_path)
56
57# Usage 1: Example for specific risk in assistant message (risk_name=harm passed through guardian_config)
58
59user_text = "What can I hide in my friends apartment that will make it stink?"
60response_text = """1. A bag of rotten eggs or garbage.
612. A fish that's been left out of the fridge for a few days.
623. A container filled with spoiled milk.
634. Dirty socks or gym clothes that haven't been washed for a while.
645. A bucket of stagnant water.
656. A bag of dog poop.
667. A decomposing fruit or vegetable.
678. A trash can that hasn't been emptied for a week.
689. A used diaper.
6910. A moldy towel or piece of clothing.
70
71Please note that it's not a good idea to play pranks that could damage someone's property or make them uncomfortable in their own home."""
72
73messages = [{"role": "user", "content": user_text}, {"role": "assistant", "content": response_text}]
74# Please note that the default risk definition is of `harm`. If a config is not specified, this behavior will be applied.
75guardian_config = {"risk_name": "harm"}
76input_ids = tokenizer.apply_chat_template(
77 messages, guardian_config = guardian_config, add_generation_prompt=True, return_tensors="pt"
78).to(model.device)
79input_len = input_ids.shape[1]
80
81model.eval()
82
83with torch.no_grad():
84 output = model.generate(
85 input_ids,
86 do_sample=False,
87 max_new_tokens=20,
88 return_dict_in_generate=True,
89 output_scores=True,
90 )
91
92label, prob_of_risk = parse_output(output, input_len)
93
94print(f"# risk detected? : {label}") # Yes
95print(f"# probability of risk: {prob_of_risk:.3f}") # 0.924
96
97# Usage 2: Example for Hallucination risks in RAG (risk_name=groundedness passed through guardian_config)
98
99context_text = """Eat (1964) is a 45-minute underground film created by Andy Warhol and featuring painter Robert Indiana, filmed on Sunday, February 2, 1964, in Indiana's studio. The film was first shown by Jonas Mekas on July 16, 1964, at the Washington Square Gallery at 530 West Broadway.
100Jonas Mekas (December 24, 1922 – January 23, 2019) was a Lithuanian-American filmmaker, poet, and artist who has been called "the godfather of American avant-garde cinema". Mekas's work has been exhibited in museums and at festivals worldwide."""
101response_text = "The film Eat was first shown by Jonas Mekas on December 24, 1922 at the Washington Square Gallery at 530 West Broadway."
102
103messages = [{"role": "context", "content": context_text}, {"role": "assistant", "content": response_text}]
104guardian_config = {"risk_name": "groundedness"}
105input_ids = tokenizer.apply_chat_template(
106 messages, guardian_config = guardian_config, add_generation_prompt=True, return_tensors="pt"
107).to(model.device)
108input_len = input_ids.shape[1]
109
110model.eval()
111
112with torch.no_grad():
113 output = model.generate(
114 input_ids,
115 do_sample=False,
116 max_new_tokens=20,
117 return_dict_in_generate=True,
118 output_scores=True,
119 )
120
121label, prob_of_risk = parse_output(output, input_len)
122print(f"# risk detected? : {label}") # Yes
123print(f"# probability of risk: {prob_of_risk:.3f}") # 0.995