fix: remove sampling params

model.py

@@ -7,10 +7,12 @@ import torch
 from ibm_watsonx_ai.client import APIClient
 from ibm_watsonx_ai.foundation_models import ModelInference
 from transformers import AutoModelForCausalLM, AutoTokenizer
-from vllm import LLM, SamplingParams
 
 from logger import logger
 
+# from vllm import LLM, SamplingParams
+
+
 safe_token = "No"
 risky_token = "Yes"
 nlogprobs = 5
@@ -23,7 +25,7 @@ if inference_engine == "VLLM":
     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)
+    # 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")
 
@@ -158,7 +160,8 @@ def generate_text(messages, criteria_name):
 
     elif inference_engine == "VLLM":
         with torch.no_grad():
-            output = model.generate(chat, sampling_params, use_tqdm=False)
+            # output = model.generate(chat, sampling_params, use_tqdm=False)
+            output = model.generate(chat, use_tqdm=False)
 
         label, prob_of_risk = parse_output(output[0])
     else:

requirements.txt

@@ -3,4 +3,3 @@ python-dotenv
 transformers
 accelerate
 ibm_watsonx_ai
-vllm