prepare for openshift deployment

.dockerignore

@@ -0,0 +1,8 @@
+
+.*
+*.yml
+*.yaml
+*.sh
+*.md
+__pycache__/
+flagged/

.env.example

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

.gitignore

@@ -2,5 +2,6 @@
 .env
 parse.py
 unparsed_catalog.json
-__pycache__
-logs
+__pycache__/
+logs.txt
+secrets.yaml

Dockerfile

@@ -0,0 +1,7 @@
+FROM python:3.12-slim
+WORKDIR /usr/src/app
+COPY . .
+RUN pip --disable-pip-version-check --no-cache-dir --no-input install -r requirements.txt
+ENV GRADIO_SERVER_NAME="0.0.0.0"
+EXPOSE 7860
+CMD ["python", "src/app.py"]

cicd/build.sh

@@ -0,0 +1,2 @@
+docker build --platform=linux/amd64 . -t granite-guardian
+docker tag granite-guardian us.icr.io/research3/granite-guardian

cicd/deploy.sh

@@ -0,0 +1,3 @@
+ibmcloud cr login
+oc delete -f deployment.yaml
+oc apply -f deployment.yaml

cicd/push_image.sh

@@ -0,0 +1,2 @@
+ibmcloud target -g aipt-experiments
+docker push us.icr.io/research3/granite-guardian

cicd/run.sh

@@ -0,0 +1,3 @@
+./build.sh
+./push_image.sh
+./deploy.sh

deployment.yaml

@@ -0,0 +1,86 @@
+apiVersion: apps/v1
+kind: Deployment
+metadata:
+  name: granite-guardian-pod
+  labels:
+    app: granite-guardian
+spec:
+  selector:
+    matchLabels:
+      run: granite-guardian
+  replicas: 1
+  template:
+    metadata:
+      labels:
+        run: granite-guardian
+    spec:
+      containers:
+        - name: granite-guardian
+          image: us.icr.io/research3/granite-guardian
+          resources:
+            limits:
+              cpu: 1
+              memory: 2Gi
+            requests:
+              cpu: 1
+              memory: 2Gi
+          ports:
+            - containerPort: 7860
+          env:
+            - name: WATSONX_API_KEY
+              valueFrom:
+                secretKeyRef:
+                  name: granite-guardian-secrets
+                  key: WATSONX_API_KEY
+            - name: WATSONX_PROJECT_ID
+              valueFrom:
+                secretKeyRef:
+                  name: granite-guardian-secrets
+                  key: WATSONX_PROJECT_ID
+            - name: INFERENCE_ENGINE
+              valueFrom:
+                secretKeyRef:
+                  name: granite-guardian-secrets
+                  key: INFERENCE_ENGINE
+      imagePullSecrets:
+        - name: all-icr-io
+---
+apiVersion: v1
+kind: Service
+metadata:
+  name: granite-guardian-service
+spec:
+  type: NodePort
+  sessionAffinity: "ClientIP"
+  selector:
+    run: granite-guardian
+  ports:
+    - port: 80
+      targetPort: 7860
+      protocol: TCP
+---
+apiVersion: networking.k8s.io/v1
+kind: Ingress
+metadata:
+  annotations:
+    ingress.kubernetes.io/allow-http: 'false'
+    ingress.kubernetes.io/ssl-redirect: 'true'
+    kubernetes.io/ingress.class: f5
+    virtual-server.f5.com/balance: round-robin
+    virtual-server.f5.com/ip: 9.12.246.36
+    virtual-server.f5.com/partition: RIS3-INT-OCP-DAL12
+    virtual-server.f5.com/clientssl: '[ { "bigIpProfile": "/Common/BlueMix" } ]'
+  name: granite-guardian-ingress
+  namespace: granite-guardian
+spec:
+  rules:
+  - host: granite-guardian.bx.cloud9.ibm.com
+    http:
+      paths:
+      - backend:
+          service:
+            name: granite-guardian-service
+            port:
+              number: 80
+        path: /
+        pathType: ImplementationSpecific

requirements.txt

@@ -1,4 +1,7 @@
-gradio
+gradio>=4,<5
 python-dotenv
 tqdm
 jinja2
+ibm_watsonx_ai
+transformers
+gradio_modal

run_cicd.sh

@@ -0,0 +1,3 @@
+./cicd/build.sh
+./cicd/push_image.sh
+./cicd/deploy.sh

src/app.py

@@ -112,6 +112,7 @@ def on_show_prompt_click(criteria, context, user_message, assistant_message, sta
 
     messages = get_messages(test_case=test_case, sub_catalog_name=state['selected_sub_catalog'])
     prompt = get_prompt(messages, criteria_name)
+    print(prompt)
     prompt = prompt.replace('<', '&lt;').replace('>', '&gt;').replace('\\n', '<br>')
     return gr.Markdown(prompt)
 
@@ -155,7 +156,7 @@ with gr.Blocks(
         ),
         head=head_style,
         fill_width=False,
-        css=os.path.join(os.path.dirname(os.path.abspath(__file__)), 'styles.css')
+        css=os.path.join(os.path.dirname(os.path.abspath(__file__)), './styles.css')
     ) as demo:
 
     state = gr.State(value={

src/logger.py

@@ -7,6 +7,6 @@ stream_handler = logging.StreamHandler()
 stream_handler.setLevel(logging.DEBUG)
 logger.addHandler(stream_handler)
 
-file_handler = logging.FileHandler('logs')
+file_handler = logging.FileHandler('logs.txt')
 file_handler.setFormatter(logging.Formatter("%(asctime)s - %(filename)s:%(lineno)d - %(message)s", datefmt="%Y-%m-%d %H:%M:%S"))
 logger.addHandler(file_handler)

src/model.py

@@ -2,13 +2,20 @@ import os
 from time import time, sleep
 from logger import logger
 import math
+import os
+from ibm_watsonx_ai.client import APIClient
+from ibm_watsonx_ai.foundation_models import ModelInference
+from transformers import AutoTokenizer
+import math
 
 safe_token = "No"
-unsafe_token = "Yes"
+risky_token = "Yes"
 nlogprobs = 5
 
-mock_model_call = os.getenv('MOCK_MODEL_CALL') == 'true'
-if not mock_model_call:
+inference_engine = os.getenv('INFERENCE_ENGINE')
+logger.debug(f"Inference engine is: '{inference_engine}'")
+
+if inference_engine == 'VLLM':
     import torch
     from vllm import LLM, SamplingParams
     from transformers import AutoTokenizer
@@ -18,6 +25,21 @@ if not mock_model_call:
     sampling_params = SamplingParams(temperature=0.0, logprobs=nlogprobs)
     model = LLM(model=model_path, tensor_parallel_size=1)
 
+elif inference_engine == "WATSONX":
+    client = APIClient(credentials={
+        'api_key': os.getenv('WATSONX_API_KEY'), 
+        'url': 'https://us-south.ml.cloud.ibm.com'})
+    
+    client.set.default_project(os.getenv('WATSONX_PROJECT_ID'))
+    hf_model_path = "ibm-granite/granite-guardian-3.0-8b"
+    tokenizer = AutoTokenizer.from_pretrained(hf_model_path)
+
+    model_id = "ibm/granite-guardian-3-8b" # 8B Model: "ibm/granite-guardian-3-8b"
+    model = ModelInference(
+        model_id=model_id,
+        api_client=client
+    )
+    
 def parse_output(output):
     label, prob = None, None
 
@@ -28,8 +50,8 @@ def parse_output(output):
             prob_of_risk = prob[1]
 
     res = next(iter(output.outputs)).text.strip()
-    if unsafe_token.lower() == res.lower():
-        label = unsafe_token
+    if risky_token.lower() == res.lower():
+        label = risky_token
     elif safe_token.lower() == res.lower():
         label = safe_token
     else:
@@ -37,6 +59,11 @@ 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
@@ -45,7 +72,7 @@ def get_probablities(logprobs):
             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():
+            if decoded_token.strip().lower() == risky_token.lower():
                 unsafe_token_prob += math.exp(token_prob.logprob)
 
     probabilities = torch.softmax(
@@ -54,6 +81,20 @@ def get_probablities(logprobs):
 
     return probabilities
 
+def get_probablities_watsonx(top_tokens_list):
+    safe_token_prob = 1e-50
+    risky_token_prob = 1e-50
+    for top_tokens in top_tokens_list:
+        for token in top_tokens:
+            if token['text'].strip().lower() == safe_token.lower():
+                safe_token_prob += math.exp(token['logprob'])
+            if token['text'].strip().lower() == risky_token.lower():
+                risky_token_prob += math.exp(token['logprob'])
+
+    probabilities = softmax([math.log(safe_token_prob), math.log(risky_token_prob)])
+
+    return probabilities
+
 def get_prompt(messages, criteria_name):
     guardian_config = {"risk_name": criteria_name if criteria_name != 'general_harm' else 'harm'}
     return tokenizer.apply_chat_template(
@@ -62,26 +103,65 @@ def get_prompt(messages, criteria_name):
         tokenize=False,
         add_generation_prompt=True)
 
+def generate_tokens(prompt):
+    result = model.generate(
+        prompt=[prompt],
+        params={
+            'decoding_method':'greedy',
+            'max_new_tokens': 20, 
+            "temperature": 0,
+            "return_options": {
+                "token_logprobs": True,
+                "generated_tokens": True,
+                "input_text": True,
+                "top_n_tokens": 5
+            }
+        })
+    return result[0]['results'][0]['generated_tokens']
+
+def parse_output_watsonx(generated_tokens_list):
+    label, prob_of_risk = None, None
+
+    if nlogprobs > 0:
+        top_tokens_list = [generated_tokens['top_tokens'] for generated_tokens in generated_tokens_list]
+        prob = get_probablities_watsonx(top_tokens_list)
+        prob_of_risk = prob[1]
+
+    res = next(iter(generated_tokens_list))['text'].strip()
+
+    if risky_token.lower() == res.lower():
+        label = risky_token
+    elif safe_token.lower() == res.lower():
+        label = safe_token
+    else:
+        label = "Failed"
+
+    return label, prob_of_risk
 
 def generate_text(messages, criteria_name):
-    logger.debug(f'Messages are: \n{messages}')
-    
-    mock_model_call = os.getenv('MOCK_MODEL_CALL') == 'true'
-    if mock_model_call:
-        logger.debug('Returning mocked model result.')
-        sleep(1)
-        return {'assessment': 'Yes', 'certainty': 0.97}
+    logger.debug(f'Messages used to create the prompt are: \n{messages}')
     
     start = time()
+
     chat = get_prompt(messages, criteria_name)
     logger.debug(f'Prompt is \n{chat}')
+
+    if inference_engine=="MOCK":
+        logger.debug('Returning mocked model result.')
+        sleep(1)
+        label, prob_of_risk = 'Yes', 0.97
     
-    with torch.no_grad():
-        output = model.generate(chat, sampling_params, use_tqdm=False)
+    elif inference_engine=="WATSONX":
+        generated_tokens = generate_tokens(chat)
+        label, prob_of_risk = parse_output_watsonx(generated_tokens)
 
-    # predicted_label = output[0].outputs[0].text.strip()
+    elif inference_engine=="VLLM":
+        with torch.no_grad():
+            output = model.generate(chat, sampling_params, use_tqdm=False)
 
-    label, prob_of_risk = parse_output(output[0])
+        label, prob_of_risk = parse_output(output[0])
+    else:
+        raise Exception("Environment variable 'INFERENCE_ENGINE' must be one of [WATSONX, MOCK, VLLM]")
 
     logger.debug(f'Model generated label: \n{label}')
     logger.debug(f'Model prob_of_risk: \n{prob_of_risk}')

src/utils.py

@@ -1,27 +1,6 @@
-import json
-from jinja2 import Template
 import argparse
 import os
 
-# with open('prompt_templates.json', mode='r', encoding="utf-8") as f:
-#     prompt_templates = json.load(f)
-
-# def assessment_prompt(content):
-#     return {"role": "user", "content": content}
-
-# def get_prompt_template(test_case, sub_catalog_name):
-#     test_case_name = test_case['name']
-#     if sub_catalog_name == 'harmful_content_in_user_prompt':
-#         template_type = 'prompt'
-#     elif sub_catalog_name == 'harmful_content_in_assistant_response':
-#         template_type = 'prompt_response'
-#     elif sub_catalog_name == 'rag_hallucination_risks':
-#         template_type = test_case_name
-#     return prompt_templates[f'{test_case_name}>{template_type}']
-
-# def get_prompt_from_test_case(test_case, sub_catalog_name):
-#     return assessment_prompt(Template(get_prompt_template(test_case, sub_catalog_name)).render(**test_case))
-
 def get_messages(test_case, sub_catalog_name) -> list[dict[str,str]]:
     messages = []
 
@@ -76,14 +55,16 @@ def get_evaluated_component(sub_catalog_name, criteria_name):
     return component
 
 def to_title_case(input_string):
-    if input_string == 'rag_hallucination_risks': 
+    if input_string == 'rag_hallucination_risks':
         return 'RAG Hallucination Risks'
     return ' '.join(word.capitalize() for word in input_string.split('_'))
 
+def capitalize_first_word(input_string):
+    return ' '.join(word.capitalize() if i == 0 else word for i, word in enumerate(input_string.split('_')))
+
 def to_snake_case(text):
     return text.lower().replace(" ", "_")
 
-
 def load_command_line_args():
     parser = argparse.ArgumentParser()
     parser.add_argument("--model_path", type=str, default=None, help="Path to the model or HF repo")