fix: fixing mistral answering and prompt formatting

backend/controller.py

@@ -59,13 +59,15 @@ def interference(
                 raise RuntimeError("There was an error in the selected XAI approach.")
 
         # call the explained chat function with the model instance
-        prompt_output, history_output, xai_graphic, xai_markup = explained_chat(
-            model=model,
-            xai=xai,
-            message=prompt,
-            history=history,
-            system_prompt=system_prompt,
-            knowledge=knowledge,
+        prompt_output, history_output, xai_interactive, xai_markup, xai_plot = (
+            explained_chat(
+                model=model,
+                xai=xai,
+                message=prompt,
+                history=history,
+                system_prompt=system_prompt,
+                knowledge=knowledge,
+            )
         )
     # if no XAI approach is selected call the vanilla chat function
     else:
@@ -78,16 +80,17 @@ def interference(
             knowledge=knowledge,
         )
         # set XAI outputs to disclaimer html/none
-        xai_graphic, xai_markup = (
+        xai_interactive, xai_markup, xai_plot = (
             """
             <div style="text-align: center"><h4>Without Selected XAI Approach,
             no graphic will be displayed</h4></div>
             """,
             [("", "")],
+            None,
         )
 
     # return the outputs
-    return prompt_output, history_output, xai_graphic, xai_markup
+    return prompt_output, history_output, xai_interactive, xai_markup, xai_plot
 
 
 # simple chat function that calls the model
@@ -121,10 +124,10 @@ def explained_chat(
     prompt = model.format_prompt(message, history, system_prompt, knowledge)
 
     # generating an answer using the methods chat function
-    answer, xai_graphic, xai_markup = xai.chat_explained(model, prompt)
+    answer, xai_graphic, xai_markup, xai_plot = xai.chat_explained(model, prompt)
 
     # updating the chat history with the new answer
     history.append((message, answer))
 
     # returning the updated history, xai graphic and xai plot elements
-    return "", history, xai_graphic, xai_markup
+    return "", history, xai_graphic, xai_markup, xai_plot

explanation/interpret_captum.py

@@ -52,4 +52,4 @@ def chat_explained(model, prompt):
     marked_text = markup_text(input_tokens, values, variant="captum")
 
     # return response, graphic and marked_text array
-    return response_text, graphic, marked_text
+    return response_text, graphic, marked_text, None

explanation/interpret_shap.py

@@ -23,29 +23,6 @@ def extract_seq_att(shap_values):
     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
@@ -80,3 +57,26 @@ def create_graphic(shap_values):
 
     # return the html graphic as string to display in iFrame
     return str(graphic_html)
+
+
+# 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, None

model/mistral.py

@@ -58,8 +58,8 @@ def set_config(config_dict: dict):
 
 
 # advanced formatting function that takes into a account a conversation history
-# CREDIT: adapted from Venkata Bhanu Teja Pallakonda in Huggingface discussions
-## see https://huggingface.co/mistralai/Mistral-7B-Instruct-v0.1/discussions/
+# CREDIT: adapated from the Mistral AI Instruct chat template
+# see https://github.com/chujiezheng/chat_templates/blob/main/chat_templates/mistral-instruct.jinja 
 def format_prompt(message: str, history: list, system_prompt: str, knowledge: str = ""):
     prompt = ""
 
@@ -83,8 +83,13 @@ def format_prompt(message: str, history: list, system_prompt: str, knowledge: st
         # adds conversation history to the prompt
         for conversation in history[1:]:
             # takes all the following conversations and adds them as context
-            prompt += "".join(f"[INST] {conversation[0]} [/INST] {conversation[1]}</s>")
+            prompt += "".join(
+                f"\n[INST] {conversation[0]} [/INST] {conversation[1]}</s>"
+            )
 
+        prompt += """\n[INST] {message} [/INST]"""
+
+    # returns full prompt
     return prompt
 
 
@@ -93,16 +98,22 @@ def format_answer(answer: str):
     # empty answer string
     formatted_answer = ""
 
-    # extracting text after INST tokens
-    parts = answer.split("[/INST]")
-    if len(parts) >= 3:
-        # Return the text after the second occurrence of [/INST]
-        formatted_answer = parts[2].strip()
-    else:
-        # Return an empty string if there are fewer than two occurrences of [/INST]
-        formatted_answer = ""
+    # splitting answer by instruction tokens
+    segments = answer.split("[/INST]")
 
-    print(f"Cut {answer} into {formatted_answer}.")
+    # checking if proper history got returned
+    if len(segments) > 1:
+        # return text after the last ['/INST'] - reponse to last message
+        formatted_answer = segments[-1].strip()
+    else:
+        # return warning and full answer if not enough [/INST] tokens found
+        gr.Warning("""
+                   There was an issue with answer formatting...\n
+                   returning the full answer.
+                   """)
+        formatted_answer = answer
+
+    print(f"CUT:\n {answer}\nINTO:\n{formatted_answer}")
     return formatted_answer