fix: another set of attention fixes

backend/controller.py

@@ -43,7 +43,6 @@ def explained_chat(
     #    message, history, system_prompt, knowledge
     # )
     prompt = model.format_prompt(message, history, system_prompt, knowledge)
-    print(f"Formatted prompt: {prompt}")
 
     # generating an answer using the methods chat function
     answer, xai_graphic, xai_markup, xai_plot = xai.chat_explained(model, prompt)
@@ -66,10 +65,10 @@ def interference(
 ):
     # if no proper system prompt is given, use a default one
     if system_prompt in ("", " "):
-        system_prompt = """
-            You are a helpful, respectful and honest assistant.
-            Always answer as helpfully as possible, while being safe.
-        """
+        system_prompt = (
+            "You are a helpful, respectful and honest assistant."
+            "Always answer as helpfully as possible, while being safe."
+        )
 
     # if a model is selected, grab the model instance
     if model_selection.lower() == "mistral":

explanation/attention.py

@@ -37,6 +37,9 @@ def chat_explained(model, prompt):
         attention_output = mdl.format_mistral_attention(attention_output)
         averaged_attention = fmt.avg_attention(attention_output, model="mistral")
 
+        response_text = fmt.format_output_text(output_text)
+        response_text = mistral.format_answer(response_text)
+
     # otherwise use attention visualization for godel
     else:
         # get attention values for the input and output vectors
@@ -49,9 +52,8 @@ def chat_explained(model, prompt):
 
         # averaging attention across layers
         averaged_attention = fmt.avg_attention(attention_output, model="godel")
+        response_text = fmt.format_output_text(output_text)
 
-    # format response text for clean output
-    response_text = fmt.format_output_text(output_text)
     # setting placeholder for iFrame graphic
     graphic = (
         "<div style='text-align: center; font-family:arial;'><h4>Attention"

explanation/markup.py

@@ -10,6 +10,8 @@ from utils import formatting as fmt
 
 # main function that assigns each text snipped a marked bucket
 def markup_text(input_text: list, text_values: ndarray, variant: str):
+    print(f"Marking up text {input_text} and {text_values} for {variant}.")
+
     # naming of the 11 buckets
     bucket_tags = ["-5", "-4", "-3", "-2", "-1", "0", "+1", "+2", "+3", "+4", "+5"]
 
@@ -21,6 +23,12 @@ def markup_text(input_text: list, text_values: ndarray, variant: str):
     elif variant == "visualizer":
         text_values = fmt.flatten_attention(text_values)
 
+    if text_values.size != len(input_text):
+        raise ValueError(
+            "Length of input text and attribution values do not match. "
+            f"Text: {len(input_text)}, Attributions: {len(text_values)}"
+        )
+
     # determine the minimum and maximum values
     min_val, max_val = np.min(text_values), np.max(text_values)
 
@@ -47,17 +55,17 @@ def markup_text(input_text: list, text_values: ndarray, variant: str):
     for text, value in zip(input_text, text_values):
 
         # validating text and skipping empty text/special tokens
-        if text not in ("", fmt.SPECIAL_TOKENS):
-            # setting initial bucket at lowest
-            bucket = "-5"
-
-            # looping over all bucket and their threshold
-            for i, threshold in zip(bucket_tags, thresholds):
-                # updating assigned bucket if value is above threshold
-                if value >= threshold:
-                    bucket = i
-            # finally adding text and bucket assignment to list of tuples
-            marked_text.append((text, str(bucket)))
+        # if text not in fmt.SPECIAL_TOKENS:
+        # setting initial bucket at lowest
+        bucket = "-5"
+
+        # looping over all bucket and their threshold
+        for i, threshold in zip(bucket_tags, thresholds):
+            # updating assigned bucket if value is above threshold
+            if value >= threshold:
+                bucket = i
+        # finally adding text and bucket assignment to list of tuples
+        marked_text.append((text, str(bucket)))
 
     # returning list of marked text snippets as list of tuples
     return marked_text

main.py

@@ -216,36 +216,46 @@ with gr.Blocks(
                 gr.Examples(
                     label="Example Questions",
                     examples=[
-                        ["Does money buy happiness?", "", "", "Mistral", "None"],
-                        ["Does money buy happiness?", "", "", "Mistral", "SHAP"],
-                        ["Does money buy happiness?", "", "", "Mistral", "Attention"],
+                        ["Does money buy happiness?", "None", "", "Mistral", ""],
+                        ["Does money buy happiness?", "SHAP", "", "Mistral", ""],
+                        ["Does money buy happiness?", "Attention", "", "Mistral", ""],
                         [
                             "Does money buy happiness?",
-                            "",
+                            "None",
                             (
                                 "Respond from the perspective of billionaire heir"
                                 " living his best life with his father's money."
                             ),
                             "Mistral",
-                            "None",
+                            "",
                         ],
                         [
                             "Does money buy happiness?",
+                            "SHAP",
+                            (
+                                "Respond from the perspective of billionaire heir"
+                                " living his best life with his father's money."
+                            ),
+                            "Mistral",
                             "",
+                        ],
+                        [
+                            "Does money buy happiness?",
+                            "Attention",
                             (
                                 "Respond from the perspective of billionaire heir"
                                 " living his best life with his father's money."
                             ),
                             "Mistral",
-                            "SHAP",
+                            "",
                         ],
                     ],
                     inputs=[
                         user_prompt,
-                        knowledge_input,
                         system_prompt,
-                        model_selection,
                         xai_selection,
+                        model_selection,
+                        knowledge_input,
                     ],
                 )
             with gr.Accordion("GODEL Model Examples", open=False):

model/godel.py

@@ -14,7 +14,7 @@ MODEL = AutoModelForSeq2SeqLM.from_pretrained("microsoft/GODEL-v1_1-large-seq2se
 # model config definition
 CONFIG = GenerationConfig.from_pretrained("microsoft/GODEL-v1_1-large-seq2seq")
 base_config_dict = {
-    "max_new_tokens": 50,
+    "max_new_tokens": 64,
     "min_length": 8,
     "top_p": 0.9,
     "do_sample": True,

utils/modelling.py

@@ -82,23 +82,20 @@ def get_device():
         device = torch.device("cuda")
     else:
         device = torch.device("cpu")
-
     return device
 
 
 # function to set device config
-# CREDIT: Adapted from captum llama 2 example
+# CREDIT: Copied from captum llama 2 example
 # see https://captum.ai/tutorials/Llama2_LLM_Attribution
 def gpu_loading_config(max_memory: str = "15000MB"):
     n_gpus = torch.cuda.device_count()
-
     bnb_config = BitsAndBytesConfig(
         load_in_4bit=True,
         bnb_4bit_use_double_quant=True,
         bnb_4bit_quant_type="nf4",
         bnb_4bit_compute_dtype=torch.bfloat16,
     )
-
     return n_gpus, max_memory, bnb_config