use vllm

model.py

@@ -1,22 +1,61 @@
-import torch
-from torch.nn.functional import softmax
-from transformers import AutoTokenizer, AutoModelForCausalLM, AutoModel
 import os
-from time import time
+from time import time, sleep
 from logger import logger
-from time import sleep
+import math
+import torch
+from vllm import LLM, SamplingParams
+from transformers import AutoTokenizer
+
+
+
+safe_token = "No"
+unsafe_token = "Yes"
+nlogprobs = 5
+
+def parse_output(output):
+    label, prob = None, None
+
+    if nlogprobs > 0:
+        logprobs = next(iter(output.outputs)).logprobs
+        if logprobs is not None:
+            prob = get_probablities(logprobs)
+            prob_of_risk = prob[1]
+
+    res = next(iter(output.outputs)).text.strip()
+    if unsafe_token.lower() == res.lower():
+        label = unsafe_token
+    elif safe_token.lower() == res.lower():
+        label = safe_token
+    else:
+        label = "Failed"
+
+    return label, prob_of_risk.item()
+
+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
+            if decoded_token.strip().lower() == safe_token.lower():
+                safe_token_prob += math.exp(token_prob.logprob)
+            if decoded_token.strip().lower() == unsafe_token.lower():
+                unsafe_token_prob += math.exp(token_prob.logprob)
+
+    probabilities = torch.softmax(
+        torch.tensor([math.log(safe_token_prob), math.log(unsafe_token_prob)]), dim=0
+    )
+
+    return probabilities
+
 
 mock_model_call = os.getenv('MOCK_MODEL_CALL') == 'true'
 if not mock_model_call:
     use_conda = os.getenv('USE_CONDA', "false") == "true"
-    device = "cuda"
     model_path = os.getenv('MODEL_PATH')#"granite-guardian-3b-pipecleaner-r241024a"
-    logger.info(f'Model path is "{model_path}"')
+    sampling_params = SamplingParams(temperature=0.0, logprobs=nlogprobs)
+    model = LLM(model=model_path, tensor_parallel_size=1)
     tokenizer = AutoTokenizer.from_pretrained(model_path)
-    model = AutoModelForCausalLM.from_pretrained(
-        model_path,
-        device_map=device if use_conda else None
-    )
 
 def generate_text(prompt):
     logger.debug('Starting evaluation...')
@@ -28,32 +67,22 @@ def generate_text(prompt):
         return {'assessment': 'Yes', 'certainty': 0.97}
     else:
         start = time()
-        tokenized_chat = tokenizer.apply_chat_template(
-            [prompt],
-            tokenize=True,
-            add_generation_prompt=True,
-            return_tensors="pt")
-        if use_conda:
-            tokenized_chat = tokenized_chat.to(device)
+        
+        tokenized_chat = tokenizer.apply_chat_template([prompt], tokenize=False, add_generation_prompt=True)
+
         with torch.no_grad():
-            logits = model(tokenized_chat).logits
-            gen_outputs = model.generate(tokenized_chat, max_new_tokens=128)
-
-        generated_text = tokenizer.decode(gen_outputs[0])
-        logger.debug(f'Model generated text: \n{generated_text}')
-        vocab = tokenizer.get_vocab()
-        selected_logits = logits[0, -1, [vocab['No'], vocab['Yes']]]
-        probabilities = softmax(selected_logits, dim=0)
-
-        prob = probabilities[1].item()
-        logger.debug(f'Certainty is: {prob} from probabilities {probabilities}')
-        certainty = prob
-        assessment = 'Yes' if certainty > 0.5 else 'No'
-        certainty = 1 - certainty if certainty < 0.5 else certainty
-        certainty = f'{round(certainty,3)}'
+            output = model.generate(tokenized_chat, sampling_params, use_tqdm=False)
+
+
+        predicted_label = output[0].outputs[0].text.strip()
+
+        label, prob_of_risk = parse_output(output[0])
 
+        logger.debug(f'Model generated label: \n{label}')
+        logger.debug(f'Model prob_of_risk: \n{prob_of_risk}')
+        
         end = time()
         total = end - start
         logger.debug(f'it took {round(total/60, 2)} mins')
 
-        return {'assessment': assessment, 'certainty': certainty}
+        return {'assessment': label, 'certainty': prob_of_risk}

requirements_frozen.txt

@@ -1,67 +1,154 @@
+accelerate==0.34.2
+aiobotocore==2.15.1
 aiofiles==23.2.1
+aiohappyeyeballs==2.4.3
+aiohttp==3.10.8
+aioitertools==0.12.0
+aiosignal==1.3.1
 annotated-types==0.7.0
 anyio==4.6.0
+async-timeout==4.0.3
+attrs==24.2.0
+botocore==1.35.23
 certifi==2024.8.30
 charset-normalizer==3.3.2
 click==8.1.7
+cloudpickle==3.0.0
 contourpy==1.3.0
 cycler==0.12.1
+datasets==3.0.1
+deprecation==2.1.0
+dill==0.3.8
+diskcache==5.6.3
+distro==1.9.0
+dmf-lib @ git+ssh://git@github.ibm.com/arc/dmf-library.git@6acf931132183153684c1c9a8edd6dbfec6f0372
+duckdb==1.1.1
+einops==0.8.0
 exceptiongroup==1.2.2
 fastapi==0.115.0
 ffmpy==0.4.0
 filelock==3.16.1
 fonttools==4.54.1
-fsspec==2024.9.0
+frozenlist==1.4.1
+fsspec==2024.6.1
+gguf==0.10.0
 gradio==4.44.1
 gradio_client==1.3.0
 h11==0.14.0
 httpcore==1.0.6
+httptools==0.6.1
 httpx==0.27.2
 huggingface-hub==0.25.1
+ibm-cos-sdk==2.13.6
+ibm-cos-sdk-core==2.13.6
+ibm-cos-sdk-s3transfer==2.13.6
 idna==3.10
+importlib_metadata==8.5.0
 importlib_resources==6.4.5
+interegular==0.3.3
 Jinja2==3.1.4
+jiter==0.6.1
+jmespath==1.0.1
+jsonschema==4.23.0
+jsonschema-specifications==2024.10.1
 kiwisolver==1.4.7
+lark==1.2.2
+llvmlite==0.43.0
+lm-format-enforcer==0.10.6
 markdown-it-py==3.0.0
 MarkupSafe==2.1.5
 matplotlib==3.9.2
 mdurl==0.1.2
+mistral_common==1.4.4
+mmh3==4.1.0
 mpmath==1.3.0
+msgpack==1.1.0
+msgspec==0.18.6
+multidict==6.1.0
+multiprocess==0.70.16
+nest-asyncio==1.6.0
 networkx==3.3
+numba==0.60.0
 numpy==1.26.4
+nvidia-cublas-cu12==12.1.3.1
+nvidia-cuda-cupti-cu12==12.1.105
+nvidia-cuda-nvrtc-cu12==12.1.105
+nvidia-cuda-runtime-cu12==12.1.105
+nvidia-cudnn-cu12==9.1.0.70
+nvidia-cufft-cu12==11.0.2.54
+nvidia-curand-cu12==10.3.2.106
+nvidia-cusolver-cu12==11.4.5.107
+nvidia-cusparse-cu12==12.1.0.106
+nvidia-ml-py==12.560.30
+nvidia-nccl-cu12==2.20.5
+nvidia-nvjitlink-cu12==12.6.77
+nvidia-nvtx-cu12==12.1.105
+openai==1.51.2
 orjson==3.10.7
+outlines==0.0.46
 packaging==24.1
 pandas==2.2.3
+partial-json-parser==0.2.1.1.post4
 pillow==10.4.0
+progressbar==2.5
+prometheus-fastapi-instrumentator==7.0.0
+prometheus_client==0.21.0
+protobuf==5.28.2
+psutil==6.0.0
+py-cpuinfo==9.0.0
+pyairports==2.1.1
+pyarrow==17.0.0
+pycountry==24.6.1
 pydantic==2.9.2
 pydantic_core==2.23.4
 pydub==0.25.1
 Pygments==2.18.0
+pyiceberg==0.7.1
 pyparsing==3.1.4
 python-dateutil==2.9.0.post0
 python-dotenv==1.0.1
 python-multipart==0.0.12
 pytz==2024.2
 PyYAML==6.0.2
+pyzmq==26.2.0
+ray==2.37.0
+referencing==0.35.1
 regex==2024.9.11
 requests==2.32.3
 rich==13.9.2
+rpds-py==0.20.0
 ruff==0.6.9
+s3fs==2023.12.2
 safetensors==0.4.5
 semantic-version==2.10.0
+sentencepiece==0.2.0
 shellingham==1.5.4
 six==1.16.0
 sniffio==1.3.1
+sortedcontainers==2.4.0
 starlette==0.38.6
+strictyaml==1.7.3
 sympy==1.13.3
+tenacity==8.5.0
+tiktoken==0.7.0
 tokenizers==0.20.0
 tomlkit==0.12.0
-torch==2.2.2
+torch==2.4.0
+torchvision==0.19.0
 tqdm==4.66.5
 transformers==4.45.1
+triton==3.0.0
 typer==0.12.5
 typing_extensions==4.12.2
 tzdata==2024.2
 urllib3==2.2.3
 uvicorn==0.31.0
+uvloop==0.20.0
+vllm==0.6.2
+watchfiles==0.24.0
 websockets==12.0
+wrapt==1.16.0
+xformers==0.0.27.post2
+xxhash==3.5.0
+yarl==1.13.1
+zipp==3.20.2