feat/fix: several minor fixes and additions

explanation/interpret_captum.py

@@ -0,0 +1,55 @@
+# external imports
+from captum.attr import LLMAttribution, TextTokenInput, KernelShap
+import torch
+
+# internal imports
+from utils import formatting as fmt
+from .markup import markup_text
+
+
+# function to extract sequence attribution
+def cpt_extract_seq_att(attr):
+
+    # getting values from captum
+    values = attr.seq_attr.to(torch.device("cpu")).numpy()
+
+    # format the input tokens nicely and check for mismatch
+    input_tokens = fmt.format_tokens(attr.input_tokens)
+    if len(attr.input_tokens) != len(values):
+        raise RuntimeError("values and input len mismatch")
+
+    # return a list of tuples with token and value
+    return list(zip(input_tokens, values))
+
+
+# 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

@@ -0,0 +1,82 @@
+# 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
+
+
+# function to extract summarized sequence wise attribution
+def extract_seq_att(shap_values):
+
+    # extracting summed up shap values
+    values = fmt.flatten_attribution(shap_values.values[0], 1)
+
+    # returning list of tuples of token and value
+    return list(zip(shap_values.data[0], values))
+
+
+# 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/plotting.py

@@ -0,0 +1,58 @@
+# plotting functions
+
+# external imports
+import numpy as np
+import matplotlib.pyplot as plt
+
+
+def plot_seq(seq_values: list, method_model: tuple = ("", "")):
+
+    # Separate the tokens and their corresponding importance values
+    tokens, importance = zip(*seq_values)
+
+    # Convert importance values to numpy array for conditional coloring
+    importance = np.array(importance)
+    importance = importance.log
+
+    # Determine the colors based on the sign of the importance values
+    colors = ["#ff0051" if val > 0 else "#008bfb" for val in importance]
+
+    # Create a bar plot
+    plt.figure(figsize=(len(tokens) * 0.9, np.max(importance)))
+    x_positions = range(len(tokens))  # Positions for the bars
+
+    # Creating vertical bar plot
+    bar_width = 0.8  # Increase this value to make the bars wider
+    plt.bar(x_positions, importance, color=colors, align="center", width=bar_width)
+    plt.yscale("symlog")
+
+    # Annotating each bar with its value
+    padding = 0.1  # Padding for text annotation
+    for x, (y, color) in enumerate(zip(importance, colors)):
+        sign = "+" if y > 0 else ""
+        plt.annotate(
+            f"{sign}{y:.2f}",  # Format the value with sign
+            xy=(x, y + padding if y > 0 else y - padding),
+            ha="center",
+            color=color,
+            va="bottom" if y > 0 else "top",  # Vertical alignment
+            fontweight="bold",  # Bold text
+            bbox={
+                "facecolor": "white",
+                "edgecolor": "none",
+                "boxstyle": "round,pad=0.1",
+            },  # White background
+        )
+
+    plt.axhline(0, color="black", linewidth=1)
+    plt.title(f"Input Token Attribution with {method_model[0]} on {method_model[1]}")
+    plt.xlabel("Input Tokens", labelpad=0.5)
+    plt.ylabel("Attribution")
+    plt.xticks(x_positions, tokens, rotation=45)
+
+    # Adjust y-axis limits to ensure there's enough space for labels
+    y_min, y_max = plt.ylim()
+    y_range = y_max - y_min
+    plt.ylim(y_min - 0.1 * y_range, y_max + 0.1 * y_range)
+
+    return plt

main.py

@@ -102,45 +102,46 @@ with gr.Blocks(
                 """)
         # row with columns for the different settings
         with gr.Row(equal_height=True):
-            # column that takes up 3/4 of the row
-            with gr.Column(scale=2):
-                # textbox to enter the system prompt
-                system_prompt = gr.Textbox(
-                    label="System Prompt",
-                    info="Set the models system prompt, dictating how it answers.",
-                    # default system prompt is set to this in the backend
-                    placeholder=(
-                        "You are a helpful, respectful and honest assistant. Always"
-                        " answer as helpfully as possible, while being safe."
-                    ),
-                )
-            # column that takes up 1/4 of the row
-            with gr.Column(scale=1):
-                # checkbox group to select the xai method
-                xai_selection = gr.Radio(
-                    ["None", "SHAP", "Attention"],
-                    label="Interpretability Settings",
-                    info="Select a Interpretability Implementation to use.",
-                    value="None",
-                    interactive=True,
-                    show_label=True,
-                )
-            # column that takes up 1/4 of the row
-            with gr.Column(scale=1):
-                # checkbox group to select the xai method
-                model_selection = gr.Radio(
-                    ["GODEL", "Mistral"],
-                    label="Model Settings",
-                    info="Select a Model to use.",
-                    value="GODEL",
-                    interactive=True,
-                    show_label=True,
-                )
+            with gr.Accordion("Application Settings", open=False):
+                # column that takes up 3/4 of the row
+                with gr.Column(scale=2):
+                    # textbox to enter the system prompt
+                    system_prompt = gr.Textbox(
+                        label="System Prompt",
+                        info="Set the models system prompt, dictating how it answers.",
+                        # default system prompt is set to this in the backend
+                        placeholder=(
+                            "You are a helpful, respectful and honest assistant. Always"
+                            " answer as helpfully as possible, while being safe."
+                        ),
+                    )
+                # column that takes up 1/4 of the row
+                with gr.Column(scale=1):
+                    # checkbox group to select the xai method
+                    xai_selection = gr.Radio(
+                        ["None", "SHAP", "Attention"],
+                        label="Interpretability Settings",
+                        info="Select a Interpretability Implementation to use.",
+                        value="None",
+                        interactive=True,
+                        show_label=True,
+                    )
+                # column that takes up 1/4 of the row
+                with gr.Column(scale=1):
+                    # checkbox group to select the xai method
+                    model_selection = gr.Radio(
+                        ["GODEL", "Mistral"],
+                        label="Model Settings",
+                        info="Select a Model to use.",
+                        value="GODEL",
+                        interactive=True,
+                        show_label=True,
+                    )
 
-            # calling info functions on inputs/submits for different settings
-            system_prompt.submit(system_prompt_info, [system_prompt])
-            xai_selection.input(xai_info, [xai_selection])
-            model_selection.input(model_info, [model_selection])
+                # calling info functions on inputs/submits for different settings
+                system_prompt.submit(system_prompt_info, [system_prompt])
+                xai_selection.input(xai_info, [xai_selection])
+                model_selection.input(model_info, [model_selection])
 
         # row with chatbot ui displaying "conversation" with the model
         with gr.Row(equal_height=True):
@@ -251,6 +252,11 @@ with gr.Blocks(
                 show_label=True,
                 height="400px",
             )
+        with gr.Row():
+            with gr.Accordion("Explanation Plot", open=False):
+                xai_plot = gr.Plot(
+                    label="Input Sequence Attribution Plot", show_label=True
+                )
 
     # functions to trigger the controller
     ## takes information for the chat and the xai selection

model/mistral.py

@@ -25,7 +25,6 @@ else:
     MODEL = AutoModelForCausalLM.from_pretrained("mistralai/Mistral-7B-Instruct-v0.2")
     MODEL.to(device)
 TOKENIZER = AutoTokenizer.from_pretrained("mistralai/Mistral-7B-Instruct-v0.2")
-TOKENIZER.pad_token = TOKENIZER.eos_token
 
 # default model config
 CONFIG = GenerationConfig.from_pretrained("mistralai/Mistral-7B-Instruct-v0.2")
@@ -103,6 +102,7 @@ def format_answer(answer: str):
         # Return an empty string if there are fewer than two occurrences of [/INST]
         formatted_answer = ""
 
+    print(f"Cut {answer} into {formatted_answer}.")
     return formatted_answer
 
 

requirements.txt

@@ -10,7 +10,6 @@ markdown~=3.5.1
 huggingface_hub~=0.19.4
 fastapi~=0.104.1
 uvicorn~=0.24.0
-tinydb~=4.8.0
 black~=23.12.0
 pylint~=3.0.0
 numpy