start using local granite guardian model

.env.example

@@ -0,0 +1,2 @@
+MODEL_PATH='../dmf_models/granite-guardian-8b-pipecleaner-r241024a'
+USE_CONDA='true'

app.py

@@ -1,9 +1,13 @@
 import gradio as gr
 from dotenv import load_dotenv
+
+from utils import get_prompt_from_test_case
+load_dotenv()
 import json
-from generate import generate_text, get_prompt_from_test_case
-from time import sleep
-# load_dotenv()
+from model import generate_text
+import logging
+
+logging.getLogger('demo')
 
 catalog = {}
 all_test_cases = []
@@ -16,7 +20,7 @@ test_case_name = gr.HTML(f'<h2>{starting_test_case["name"]}</h2>')
 criteria = gr.Textbox(label="Definition", lines=3, interactive=False, value=starting_test_case['criteria'])
 context = gr.Textbox(label="Context", lines=3, interactive=True, value=starting_test_case['context'], visible=False)
 user_message = gr.Textbox(label="User Message", lines=3, interactive=True, value=starting_test_case['user_message'])
-assistant_message = gr.Textbox(label="Assistant Message", lines=3, interactive=True, value=starting_test_case['assistant_message'])
+assistant_message = gr.Textbox(label="Assistant Message", lines=3, interactive=True, visible=False, value=starting_test_case['assistant_message'])
 catalog_buttons: dict[str,dict[str,gr.Button]] = {}
 result_text = gr.Textbox(label="Result", interactive=False)
 result_certainty = gr.Number(label="Certainty", interactive=False, value='')
@@ -31,32 +35,16 @@ for sub_catalog_name, sub_catalog in catalog.items():
         catalog_buttons[sub_catalog_name][test_case['name']] = \
             gr.Button(test_case['name'], elem_classes=elem_classes, variant='secondary', size='sm', elem_id=elem_id)
 
-# watsonx_api_url = os.getenv("WATSONX_URL", None) 
-# watsonx_project_id = os.getenv("WATSONX_API_KEY", None) 
-# watsonx_api_key = os.getenv("WATSONX_PROJECT_ID", None) 
-
-# client = APIClient(credentials={
-#     "url": watsonx_api_url,
-#     "project_id": watsonx_project_id,
-#     "api_key": watsonx_api_key
-# })
-# model = ModelInference(model_id=ModelTypes.LLAMA_3_8B_INSTRUCT, api_client=client)
-
-# client.set.default_project(watsonx_project_id)
-
-# bam_api_url = os.getenv("BAM_URL", None) 
-# bam_api_key = os.getenv("BAM_API_KEY", None) 
-# client = Client(credentials=Credentials(api_endpoint=bam_api_url, api_key=bam_api_key ))
-
-
-def on_test_case_click(link):
-    selected_test_case = [t for sub_catalog in catalog.values() for t in sub_catalog if t['name'] == link][0]
+def on_test_case_click(link, event: gr.EventData):
+    target_sub_catalog_name, target_test_case_name = event.target.elem_id.split('_')
+    selected_test_case = [t for sub_catalog_name, sub_catalog in catalog.items() for t in sub_catalog if t['name'] == link and sub_catalog_name == target_sub_catalog_name][0]
+    print(selected_test_case['assistant_message'])
     return {
         test_case_name: f'<h2>{selected_test_case["name"]}</h2>',
         criteria: selected_test_case['criteria'],
-        context: selected_test_case['context'] if selected_test_case['context'] is not None else gr.update(visible=False),
+        context: selected_test_case['context'] if selected_test_case['context'] is not None else gr.update(visible=False, value=''),
         user_message: selected_test_case['user_message'],
-        assistant_message: selected_test_case['assistant_message'],
+        assistant_message: gr.update(value=selected_test_case['assistant_message'], visible=True) if selected_test_case['assistant_message'] is not None else gr.update(visible=False, value=''),
         result_text: gr.update(value=''),
         result_certainty: gr.update(value='')
     }
@@ -65,18 +53,19 @@ def change_button_color(event: gr.EventData):
     return [gr.update(elem_classes=['catalog-button', 'selected']) if v.elem_id == event.target.elem_id else gr.update(elem_classes=['catalog-button']) for c in catalog_buttons.values() for v in c.values()]
 
 def on_submit(inputs):
-    # prompt = get_prompt_from_test_case({
-    #     'criteria': inputs[criteria],
-    #     'context': inputs[context],
-    #     'user_message': inputs[user_message],
-    #     'assistant_message': inputs[assistant_message],
-    # })
-    # result = generate_text(prompt)
-    # return result['assessment'], result['certainty']
-    sleep(3)
-    return 'Yes', 0.97
+    prompt = get_prompt_from_test_case({
+        'criteria': inputs[criteria],
+        'context': inputs[context],
+        'user_message': inputs[user_message],
+        'assistant_message': inputs[assistant_message],
+    })
+    result = generate_text(prompt)
+    return result['assessment'], result['certainty']
+    # sleep(3)
+    # return 'Yes', 0.97
 
 with gr.Blocks(
+    title='Granite Guardian',
      theme=gr.themes.Soft(font=[gr.themes.GoogleFont("IBM Plex Sans")]), css='styles.css') as demo:
     with gr.Row():
         gr.HTML('<h1>Granite Guardian</h1>', elem_classes='title')

logger.py

@@ -0,0 +1,7 @@
+import logging
+
+logger = logging.getLogger('demo')
+logger.setLevel(logging.DEBUG)
+handler = logging.StreamHandler()
+handler.setLevel(logging.DEBUG)
+logger.addHandler(handler)

model.py

@@ -0,0 +1,53 @@
+import logging.handlers
+import torch
+from torch.nn.functional import softmax
+from transformers import AutoTokenizer, AutoModelForCausalLM, AutoModel
+import jinja2
+import os
+from time import time
+from logger import logger
+
+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}"')
+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...')
+    logger.debug(f'Prompts content is: \n{prompt["content"]}')
+    start = time()
+    tokenized_chat = tokenizer.apply_chat_template(
+        [prompt], 
+        tokenize=True, 
+        add_generation_prompt=True, 
+        return_tensors="pt")#.to(device)
+    if use_conda:
+        tokenized_chat.to(device)
+    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)}'
+
+    end = time()
+    total = end - start
+    logger.debug(f'it took {round(total/60, 2)} mins')
+
+    return {'assessment': assessment, 'certainty': certainty}

generate.py → utils.py

@@ -25,18 +25,16 @@ def turn_section_content(test_case):
     result = ''
     if test_case['context'] != '':
         result += 'Context: ' + test_case['context'] + '\n'
-    result += 'User message: ' + test_case['user_message'] #+ '\n' + 'Assistant message: ' + test_case['assistant_message']
+
+    result += 'User message: ' + test_case['user_message']
+    
+    if test_case['assistant_message'] != '':
+        result += '\n\nAssistant message: ' + test_case['assistant_message'] + '\n'
     return result
 
 def get_prompt_from_test_case(test_case):
-    return json.dumps(assessment_prompt(assessment_prompt_content().format(
+    print(json.dumps(test_case, indent=4))
+    return assessment_prompt(assessment_prompt_content().format(
         turn_section_content=turn_section_content(test_case),
         criteria=test_case['criteria']
-    )))
-
-def generate_text(prompt):
-    result = requests.post('http://localhost:8081/generate', json={'input': prompt}).json()
-    assessment = 'Yes' if result['certainty'] > 0.5 else 'No'
-    certainty = 1 - result['certainty'] if result['certainty'] < 0.5 else result['certainty']
-    certainty = f'{round(certainty,3)}'
-    return {'assessment': assessment, 'certainty': certainty}
+    ))