Adapt model to guardian cookbook

model.py

@@ -15,19 +15,21 @@ from logger import logger
 
 safe_token = "No"
 risky_token = "Yes"
-nlogprobs = 5
+nlogprobs = 20
 
 inference_engine = os.getenv("INFERENCE_ENGINE", "VLLM")
 logger.debug(f"Inference engine is: '{inference_engine}'")
 
 if inference_engine == "VLLM":
+    device = torch.device("gpu")
 
     model_path = os.getenv("MODEL_PATH", "ibm-granite/granite-guardian-3.0-8b")
     logger.debug(f"model_path is {model_path}")
     tokenizer = AutoTokenizer.from_pretrained(model_path)
     # sampling_params = SamplingParams(temperature=0.0, logprobs=nlogprobs)
     # model = LLM(model=model_path, tensor_parallel_size=1)
-    model = AutoModelForCausalLM.from_pretrained(model_path, torch_dtype=torch.float16, device_map="auto")
+    model = AutoModelForCausalLM.from_pretrained(model_path)
+    model = model.to(device).eval()
 
 elif inference_engine == "WATSONX":
     client = APIClient(
@@ -42,16 +44,17 @@ elif inference_engine == "WATSONX":
     model = ModelInference(model_id=model_id, api_client=client)
 
 
-def parse_output(output):
-    label, prob = None, None
-
+def parse_output(output, input_len):
+    label, prob_of_risk = None, None
     if nlogprobs > 0:
-        logprobs = next(iter(output.outputs)).logprobs
-        if logprobs is not None:
-            prob = get_probablities(logprobs)
+
+        list_index_logprobs_i = [torch.topk(token_i, k=nlogprobs, largest=True, sorted=True)
+                                 for token_i in list(output.scores)[:-1]]
+        if list_index_logprobs_i is not None:
+            prob = get_probablities(list_index_logprobs_i)
             prob_of_risk = prob[1]
 
-    res = next(iter(output.outputs)).text.strip()
+    res = tokenizer.decode(output.sequences[:,input_len:][0],skip_special_tokens=True).strip()
     if risky_token.lower() == res.lower():
         label = risky_token
     elif safe_token.lower() == res.lower():
@@ -61,29 +64,29 @@ def parse_output(output):
 
     return label, prob_of_risk.item()
 
-
-def softmax(values):
-    exp_values = [math.exp(v) for v in values]
-    total = sum(exp_values)
-    return [v / total for v in exp_values]
-
-
 def get_probablities(logprobs):
     safe_token_prob = 1e-50
     unsafe_token_prob = 1e-50
     for gen_token_i in logprobs:
-        for token_prob in gen_token_i.values():
-            decoded_token = token_prob.decoded_token
+        for logprob, index in zip(gen_token_i.values.tolist()[0], gen_token_i.indices.tolist()[0]):
+            decoded_token = tokenizer.convert_ids_to_tokens(index)
             if decoded_token.strip().lower() == safe_token.lower():
-                safe_token_prob += math.exp(token_prob.logprob)
+                safe_token_prob += math.exp(logprob)
             if decoded_token.strip().lower() == risky_token.lower():
-                unsafe_token_prob += math.exp(token_prob.logprob)
+                unsafe_token_prob += math.exp(logprob)
 
-    probabilities = torch.softmax(torch.tensor([math.log(safe_token_prob), math.log(unsafe_token_prob)]), dim=0)
+    probabilities = torch.softmax(
+        torch.tensor([math.log(safe_token_prob), math.log(unsafe_token_prob)]), dim=0
+    )
 
     return probabilities
 
 
+def softmax(values):
+    exp_values = [math.exp(v) for v in values]
+    total = sum(exp_values)
+    return [v / total for v in exp_values]
+
 def get_probablities_watsonx(top_tokens_list):
     safe_token_prob = 1e-50
     risky_token_prob = 1e-50
@@ -99,10 +102,10 @@ def get_probablities_watsonx(top_tokens_list):
     return probabilities
 
 
-def get_prompt(messages, criteria_name):
+def get_prompt(messages, criteria_name, return_tensors=None):
     guardian_config = {"risk_name": criteria_name if criteria_name != "general_harm" else "harm"}
     return tokenizer.apply_chat_template(
-        messages, guardian_config=guardian_config, tokenize=False, add_generation_prompt=True
+        messages, guardian_config=guardian_config, tokenize=False, add_generation_prompt=True, return_tensors=return_tensors
     )
 
 
@@ -155,15 +158,24 @@ def generate_text(messages, criteria_name):
         label, prob_of_risk = "Yes", 0.97
 
     elif inference_engine == "WATSONX":
+        chat = get_prompt(messages, criteria_name)
         generated_tokens = generate_tokens(chat)
         label, prob_of_risk = parse_output_watsonx(generated_tokens)
 
     elif inference_engine == "VLLM":
+        input_ids = tokenizer.apply_chat_template(messages, add_generation_prompt=True, return_tensors="pt").to(model.device)
+        input_len = input_ids.shape[1]
+
         with torch.no_grad():
             # output = model.generate(chat, sampling_params, use_tqdm=False)
-            output = model.generate(chat)
-
-        label, prob_of_risk = parse_output(output[0])
+            output = model.generate(
+                chat,
+                do_sample=False,
+                max_new_tokens=nlogprobs,
+                return_dict_in_generate=True,
+                output_scores=True,)
+
+            label, prob_of_risk = parse_output(output, input_len)
     else:
         raise Exception("Environment variable 'INFERENCE_ENGINE' must be one of [WATSONX, MOCK, VLLM]")