fix: fixing various bugs

backend/controller.py

@@ -8,9 +8,9 @@ import gradio as gr
 from model import godel
 from model import mistral
 from explanation import (
+    attention as attention_viz,
     interpret_shap as shap_int,
     interpret_captum as cpt_int,
-    visualize_att as viz,
 )
 
 
@@ -48,7 +48,7 @@ def interference(
                 else:
                     xai = shap_int
             case "attention":
-                xai = viz
+                xai = attention_viz
             case _:
                 # use Gradio warning to display error message
                 gr.Warning(f"""

explanation/attention.py

@@ -0,0 +1,51 @@
+# 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, None

explanation/interpret_captum.py

@@ -46,8 +46,9 @@ def chat_explained(model, prompt):
     # 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>"
+        """<div style='text-align: center; font-family:arial;'><h4>
+        Intepretation with Captum doesn't support an interactive graphic.</h4></div>
+        """
     )
     marked_text = markup_text(input_tokens, values, variant="captum")
 

explanation/markup.py

@@ -71,11 +71,11 @@ def color_codes():
         "-4": "#68a1fd",
         "-3": "#96b7fe",
         "-2": "#bcceff",
-        "-1:": "#dee6ff",
+        "-1": "#dee6ff",
         "0": "#ffffff",
-        "1": "#ffd9d9",
-        "2": "#ffb3b5",
-        "3": "#ff8b92",
-        "4": "#ff5c71",
-        "5": "#ff0051",
+        "+1": "#ffd9d9",
+        "+2": "#ffb3b5",
+        "+3": "#ff8b92",
+        "+4": "#ff5c71",
+        "+5": "#ff0051",
     }