feat/fix: fixing code issues, adding plotting functions

.gitignore

@@ -2,3 +2,4 @@
 __pycache__/
 /start-venv.sh
 /components/iframe/dist/
+.venv

backend/controller.py

@@ -10,7 +10,7 @@ from model import mistral
 from explanation import (
     interpret_shap as shap_int,
     interpret_captum as cpt_int,
-    visualize as viz,
+    visualize_att as viz,
 )
 
 
@@ -33,10 +33,10 @@ def interference(
 
     if model_selection.lower() == "mistral":
         model = mistral
-        print("Indetified model as Mistral")
+        print("Indentified model as Mistral")
     else:
         model = godel
-        print("Indetified model as GODEL")
+        print("Indentified model as GODEL")
 
     # if a XAI approach is selected, grab the XAI module instance
     if xai_selection in ("SHAP", "Attention"):

explanation/interpret_captum.py

@@ -1,40 +0,0 @@
-# external imports
-from captum.attr import LLMAttribution, TextTokenInput, KernelShap
-import torch
-
-# internal imports
-from utils import formatting as fmt
-from .markup import markup_text
-
-
-# main explain function that returns a chat with explanations
-def chat_explained(model, prompt):
-    model.set_config({})
-
-    # creating llm attribution class with KernelSHAP and Mistal Model, Tokenizer
-    llm_attribution = LLMAttribution(KernelShap(model.MODEL), model.TOKENIZER)
-
-    # generation attribution
-    attribution_input = TextTokenInput(prompt, model.TOKENIZER)
-    attribution_result = llm_attribution.attribute(
-        attribution_input, gen_args=model.CONFIG.to_dict()
-    )
-
-    # extracting values and input tokens
-    values = attribution_result.seq_attr.to(torch.device("cpu")).numpy()
-    input_tokens = fmt.format_tokens(attribution_result.input_tokens)
-
-    # raising error if mismatch occurs
-    if len(attribution_result.input_tokens) != len(values):
-        raise RuntimeError("values and input len mismatch")
-
-    # getting response text, graphic placeholder and marked text object
-    response_text = fmt.format_output_text(attribution_result.output_tokens)
-    graphic = (
-        "<div style='text-align: center; font-family:arial;'><h4>Attention"
-        "Intepretation with Captum doesn't support an interactive graphic.</h4></div>"
-    )
-    marked_text = markup_text(input_tokens, values, variant="captum")
-
-    # return response, graphic and marked_text array
-    return response_text, graphic, marked_text

explanation/interpret_shap.py

@@ -1,72 +0,0 @@
-# interpret module that implements the interpretability method
-
-# external imports
-from shap import models, maskers, plots, PartitionExplainer
-import torch
-
-# internal imports
-from utils import formatting as fmt
-from .markup import markup_text
-
-# global variables
-TEACHER_FORCING = None
-TEXT_MASKER = None
-
-
-# main explain function that returns a chat with explanations
-def chat_explained(model, prompt):
-    model.set_config({})
-
-    # create the shap explainer
-    shap_explainer = PartitionExplainer(model.MODEL, model.TOKENIZER)
-
-    # get the shap values for the prompt
-    shap_values = shap_explainer([prompt])
-
-    # create the explanation graphic and marked text array
-    graphic = create_graphic(shap_values)
-    marked_text = markup_text(
-        shap_values.data[0], shap_values.values[0], variant="shap"
-    )
-
-    # create the response text
-    response_text = fmt.format_output_text(shap_values.output_names)
-
-    # return response, graphic and marked_text array
-    return response_text, graphic, marked_text
-
-
-# function used to wrap the model with a shap model
-def wrap_shap(model):
-    # calling global variants
-    global TEXT_MASKER, TEACHER_FORCING
-
-    # set the device to cuda if gpu is available
-    device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
-
-    # updating the model settings
-    model.set_config()
-
-    # (re)initialize the shap models and masker
-    # creating a shap text_generation model
-    text_generation = models.TextGeneration(model.MODEL, model.TOKENIZER)
-    # wrapping the text generation model in a teacher forcing model
-    TEACHER_FORCING = models.TeacherForcing(
-        text_generation,
-        model.TOKENIZER,
-        device=str(device),
-        similarity_model=model.MODEL,
-        similarity_tokenizer=model.TOKENIZER,
-    )
-    # setting the text masker as an empty string
-    TEXT_MASKER = maskers.Text(model.TOKENIZER, " ", collapse_mask_token=True)
-
-
-# graphic plotting function that creates a html graphic (as string) for the explanation
-def create_graphic(shap_values):
-
-    # create the html graphic using shap text plot function
-    graphic_html = plots.text(shap_values, display=False)
-
-    # return the html graphic as string to display in iFrame
-    return str(graphic_html)

explanation/markup.py

@@ -66,16 +66,16 @@ def color_codes():
     return {
         # -5 to -1: Strong Light Sky Blue to Lighter Sky Blue
         # 0: white (assuming default light mode)
-        # +1 to +5 light pink to string magenta
-        "-5": "#3251a8",  # Strong Light Sky Blue
-        "-4": "#5A7FB2",  # Slightly Lighter Sky Blue
-        "-3": "#8198BC",  # Intermediate Sky Blue
-        "-2": "#A8B1C6",  # Light Sky Blue
-        "-1": "#E6F0FF",  # Very Light Sky Blue
-        "0": "#FFFFFF",  # White
-        "+1": "#FFE6F0",  # Lighter Pink
-        "+2": "#DF8CA3",  # Slightly Stronger Pink
-        "+3": "#D7708E",  # Intermediate Pink
-        "+4": "#CF5480",  # Deep Pink
-        "+5": "#A83273",  # Strong Magenta
+        # +1 to +5 light pink to strng magenta
+        "-5": "#008bfb",
+        "-4": "#68a1fd",
+        "-3": "#96b7fe",
+        "-2": "#bcceff",
+        "-1:": "#dee6ff",
+        "0": "#ffffff",
+        "1": "#ffd9d9",
+        "2": "#ffb3b5",
+        "3": "#ff8b92",
+        "4": "#ff5c71",
+        "5": "#ff0051",
     }

explanation/visualize.py

@@ -1,52 +0,0 @@
-# visualization module that creates an attention visualization
-
-
-# internal imports
-from utils import formatting as fmt
-from .markup import markup_text
-
-
-# chat function that returns an answer
-# and marked text based on attention
-def chat_explained(model, prompt):
-
-    # get encoded input
-    encoder_input_ids = model.TOKENIZER(
-        prompt, return_tensors="pt", add_special_tokens=True
-    ).input_ids
-    # generate output together with attentions of the model
-    decoder_input_ids = model.MODEL.generate(
-        encoder_input_ids, output_attentions=True, **model.CONFIG
-    )
-
-    # get input and output text as list of strings
-    encoder_text = fmt.format_tokens(
-        model.TOKENIZER.convert_ids_to_tokens(encoder_input_ids[0])
-    )
-    decoder_text = fmt.format_tokens(
-        model.TOKENIZER.convert_ids_to_tokens(decoder_input_ids[0])
-    )
-
-    # get attention values for the input and output vectors
-    # using already generated input and output
-    attention_output = model.MODEL(
-        input_ids=encoder_input_ids,
-        decoder_input_ids=decoder_input_ids,
-        output_attentions=True,
-    )
-
-    # averaging attention across layers
-    averaged_attention = fmt.avg_attention(attention_output)
-
-    # format response text for clean output
-    response_text = fmt.format_output_text(decoder_text)
-    # setting placeholder for iFrame graphic
-    graphic = (
-        "<div style='text-align: center; font-family:arial;'><h4>Attention"
-        " Visualization doesn't support an interactive graphic.</h4></div>"
-    )
-    # creating marked text using markup_text function and attention
-    marked_text = markup_text(encoder_text, averaged_attention, variant="visualizer")
-
-    # returning response, graphic and marked text array
-    return response_text, graphic, marked_text

model/mistral.py

@@ -41,13 +41,11 @@ CONFIG.update(**{
 
 
 # function to (re) set config
-def set_config(config: dict):
+def set_config(config_dict: dict):
 
-    # if config dict is given, update it
-    if config != {}:
-        CONFIG.update(**dict)
-    else:
-        CONFIG.update(**{
+    # if config dict is not given, set to default
+    if config_dict == {}:
+        config_dict = {
             "temperature": 0.7,
             "max_new_tokens": 50,
             "max_length": 50,
@@ -55,7 +53,9 @@ def set_config(config: dict):
             "repetition_penalty": 1.2,
             "do_sample": True,
             "seed": 42,
-        })
+        }
+
+    CONFIG.update(**dict)
 
 
 # advanced formatting function that takes into a account a conversation history
@@ -77,9 +77,9 @@ def format_prompt(message: str, history: list, system_prompt: str, knowledge: st
             """
     else:
         # takes the very first exchange and the system prompt as base
-        prompt = (
-            f"<s>[INST] {system_prompt} {history[0][0]} [/INST] {history[0][1]}</s>"
-        )
+        prompt = f"""
+            <s>[INST] {system_prompt} {history[0][0]} [/INST] {history[0][1]}</s>
+            """
 
         # adds conversation history to the prompt
         for conversation in history[1:]:

requirements.txt

@@ -2,7 +2,7 @@ gradio~=4.7.1
 transformers~=4.35.2
 torch~=2.1.1
 shap
-captum
+captum @ git+https://github.com/LennardZuendorf/thesis-captum.git
 bertviz~=1.4.0
 accelerate~=0.24.1
 bitsandbytes
@@ -13,9 +13,7 @@ uvicorn~=0.24.0
 tinydb~=4.8.0
 black~=23.12.0
 pylint~=3.0.0
-seaborn~=0.13.0
 numpy
 matplotlib
 pre-commit
-ipython
 gradio-iframe~=0.0.10