Adapt for deployment on HF ZeroSpace

.env.example

@@ -1,5 +1,4 @@
-MODEL_PATH='../dmf_models/granite-guardian-8b-pipecleaner-r241024a'
-USE_CONDA='true'
-INFERENCE_ENGINE='' # one of [WATSONX, MOCK, VLLM]
-WATSONX_API_KEY=""
-WATSONX_PROJECT_ID=""
+MODEL_PATH='ibm-granite/granite-guardian-3.0-8b'
+INFERENCE_ENGINE='VLLM' # one of [WATSONX, MOCK, VLLM]
+WATSONX_API_KEY=''
+WATSONX_PROJECT_ID=''

requirements.txt

@@ -4,4 +4,5 @@ tqdm
 jinja2
 ibm_watsonx_ai
 transformers
-gradio_modal
+gradio_modal
+spaces

src/model.py

@@ -7,19 +7,20 @@ from ibm_watsonx_ai.client import APIClient
 from ibm_watsonx_ai.foundation_models import ModelInference
 from transformers import AutoTokenizer
 import math
+import spaces
 
 safe_token = "No"
 risky_token = "Yes"
 nlogprobs = 5
 
-inference_engine = os.getenv('INFERENCE_ENGINE')
+inference_engine = os.getenv('INFERENCE_ENGINE', 'VLLM')
 logger.debug(f"Inference engine is: '{inference_engine}'")
 
 if inference_engine == 'VLLM':
     import torch
     from vllm import LLM, SamplingParams
     from transformers import AutoTokenizer
-    model_path = os.getenv('MODEL_PATH') #"granite-guardian-3b-pipecleaner-r241024a"
+    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)
@@ -138,6 +139,7 @@ def parse_output_watsonx(generated_tokens_list):
 
     return label, prob_of_risk
 
+@spaces.GPU
 def generate_text(messages, criteria_name):
     logger.debug(f'Messages used to create the prompt are: \n{messages}')