update

app.py

@@ -9,6 +9,7 @@ import numpy as np
 import spaces
 import nltk
 import base64
+import traceback
 from src.utils import split_into_sentences as split_into_sentences_utils
 # --- AttnTrace imports (from app_full.py) ---
 from src.models import create_model
@@ -114,6 +115,7 @@ def initialize_model_and_attr():
     except Exception as e:
         error_msg = f"Error initializing model/traceback: {str(e)}"
         print(error_msg)
+        traceback.print_exc()
         return None, None, error_msg
 
 # Initialize with defaults
@@ -608,6 +610,7 @@ def basic_get_scores_and_sources_full_response(state: State):
             state,
         )
     except Exception as e:
+        traceback.print_exc()
         return (
             gr.update(value=[("", None)], visible=False),
             gr.update(selected=0),
@@ -822,6 +825,7 @@ def basic_get_scores_and_sources(
             state,
         )
     except Exception as e:
+        traceback.print_exc()
         return (
             gr.update(value=[("", None)], visible=False),
             gr.update(selected=0),

src/attribution/attntrace.py

@@ -10,13 +10,14 @@ from .attention_utils import *
 class AttnTraceAttribution(Attribution):
     def __init__(self,  llm,explanation_level = "segment",K=5, avg_k=5, q=0.4, B=30, verbose =1):
         super().__init__(llm,explanation_level,K,verbose)
-        self.model = llm.model # Use float16 for the model
+        self.llm = llm # Use float16 for the model
+        self.model = None
         self.model_type = llm.provider
         self.tokenizer = llm.tokenizer
         self.avg_k = avg_k
         self.q = q
         self.B = B
-        self.layers = range(len(self.model.model.layers))
+        
         self.explanation_level = explanation_level
  
     def loss_to_importance(self,losses, sentences_id_list):
@@ -37,6 +38,11 @@ class AttnTraceAttribution(Attribution):
         return importances
     def attribute(self, question: str, contexts: list, answer: str,explained_answer: str, customized_template: str = None):
         start_time = time.time()
+        if self.llm.model!=None:
+            self.model = self.llm.model
+        else:
+            self.model = self.llm._load_model_if_needed()
+        self.layers = range(len(self.model.model.layers))
         model = self.model
         tokenizer = self.tokenizer
         model.eval()  # Set model to evaluation mode

src/models/Llama.py

@@ -42,7 +42,6 @@ class Llama(Model):
             messages,
             add_generation_prompt=True,
             return_tensors="pt",
-            padding=True,
             truncation=True
         ).to(model.device)
 

src/prompts.py

@@ -1,5 +1,5 @@
 MULTIPLE_PROMPT_FORCE = 'You are a helpful assistant, below is a query from a user and some relevant contexts. \
-Answer the question given the information in those contexts.\
+Answer the query given the information in those contexts.\
 \n\nContexts: [context] \n\nQuery: [question] \n\nAnswer:'
 
 SELF_CITATION_PROMPT = """You are a helpful assistant, below is a query from a user, some relevant contexts, and an answer to the query.