fix: fixing more bugs

explanation/interpret_captum.py

@@ -16,7 +16,9 @@ def chat_explained(model, prompt):
 
     # generation attribution
     attribution_input = TextTokenInput(prompt, model.TOKENIZER)
-    attribution_result = llm_attribution.attribute(attribution_input, gen_args=model.CONFIG.to_dict())
+    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()

model/mistral.py

@@ -25,35 +25,36 @@ 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
+TOKENIZER.pad_token = TOKENIZER.eos_token
 
 # default model config
 CONFIG = GenerationConfig.from_pretrained("mistralai/Mistral-7B-Instruct-v0.2")
 CONFIG.update(**{
-        "temperature": 0.7,
-        "max_new_tokens": 50,
-        "top_p": 0.9,
-        "repetition_penalty": 1.2,
-        "do_sample": True,
-        "seed": 42
+    "temperature": 0.7,
+    "max_new_tokens": 50,
+    "max_length": 50,
+    "top_p": 0.9,
+    "repetition_penalty": 1.2,
+    "do_sample": True,
+    "seed": 42,
 })
 
 
 # function to (re) set config
 def set_config(config: dict):
-    global CONFIG
 
     # if config dict is given, update it
     if config != {}:
         CONFIG.update(**dict)
     else:
         CONFIG.update(**{
-                "temperature": 0.7,
-                "max_new_tokens": 50,
-                "top_p": 0.9,
-                "repetition_penalty": 1.2,
-                "do_sample": True,
-                "seed": 42
+            "temperature": 0.7,
+            "max_new_tokens": 50,
+            "max_length": 50,
+            "top_p": 0.9,
+            "repetition_penalty": 1.2,
+            "do_sample": True,
+            "seed": 42,
         })
 
 
@@ -93,9 +94,6 @@ def format_answer(answer: str):
     # empty answer string
     formatted_answer = ""
 
-    if type(answer) == list:
-        answer = fmt.format_output_text(answer)
-
     # extracting text after INST tokens
     parts = answer.split("[/INST]")
     if len(parts) >= 3:
@@ -116,5 +114,6 @@ def respond(prompt: str):
     # generating text with tokenized input, returning output
     output_ids = MODEL.generate(input_ids, generation_config=CONFIG)
     output_text = TOKENIZER.batch_decode(output_ids)
+    output_text = fmt.format_output_text(output_text)
 
-    return fmt.format_output_text(output_text)
+    return format_answer(output_text)

pyproject.toml

@@ -21,6 +21,7 @@ exclude = '''
 
 [tool.pylint.messages_control]
 disable = [
+    "not-a-mapping",
     "arguments-differ",
     "attribute-defined-outside-init",
     "blacklisted-name",