update4

app.py

@@ -25,6 +25,8 @@ nltk.download("punkt", download_dir="/home/user/nltk_data")
 nltk.data.path.append("/home/user/nltk_data")
 from nltk.tokenize import sent_tokenize
 
+DEFAULT_TOP_K = 3
+
 # Load original app constants
 APP_TITLE = '<div class="app-title"><span class="brand">AttnTrace: </span><span class="subtitle">Attention-based Context Traceback for Long-Context LLMs</span></div>'
 APP_DESCRIPTION = """AttnTrace traces a model's generated statements back to specific parts of the context using attention-based traceback. Try it out with Meta-Llama-3.1-8B-Instruct here! See the [[paper](https://arxiv.org/abs/2506.04202)] and [[code](https://github.com/Wang-Yanting/TracLLM-Kit)] for more!
@@ -86,7 +88,7 @@ current_top_k = 3  # Add top-k tracking
 
 def update_configuration(explanation_level, top_k):
     """Update the global configuration and reinitialize attribution if needed"""
-    global current_explanation_level, current_top_k, current_attr
+    global current_explanation_level, current_top_k, current_attr, current_llm
     
     # Convert top_k to int
     top_k = int(top_k)
@@ -99,11 +101,23 @@ def update_configuration(explanation_level, top_k):
         print(f"🔄 Updating configuration: explanation_level={explanation_level}, top_k={top_k}")
         current_explanation_level = explanation_level
         current_top_k = top_k
+
+        DEFAULT_EXPLANATION_LEVEL = explanation_level
+        DEFAULT_TOP_K = top_k
         
-        # Reset attribution to force reinitialization with new config
+        # Reset both model and attribution to force complete reinitialization
+        current_llm = None
         current_attr = None
         
-        return gr.update(value=f"✅ Configuration updated: {explanation_level} level, top-{top_k}")
+        # Reinitialize with new configuration
+        try:
+            llm, attr, error_msg = initialize_model_and_attr()
+            if llm is not None and attr is not None:
+                return gr.update(value=f"✅ Configuration updated: {explanation_level} level, top-{top_k}")
+            else:
+                return gr.update(value=f"❌ Error reinitializing: {error_msg}")
+        except Exception as e:
+            return gr.update(value=f"❌ Error updating configuration: {str(e)}")
     else:
         return gr.update(value="ℹ️ Configuration unchanged")
 
@@ -133,12 +147,13 @@ def initialize_model_and_attr():
             # Use current configuration or defaults
             explanation_level = current_explanation_level or DEFAULT_EXPLANATION_LEVEL
             top_k = current_top_k or 3
-            
+            if "segment" in DEFAULT_EXPLANATION_LEVEL:
+                DEFAULT_EXPLANATION_LEVEL = "segment"
             print(f"Initializing context traceback with explanation level: {explanation_level}, top_k: {top_k}")
             current_attr = AttnTraceAttribution(
                 current_llm, 
-                explanation_level=explanation_level, 
-                K=top_k, 
+                explanation_level= DEFAULT_EXPLANATION_LEVEL, 
+                K=DEFAULT_TOP_K, 
                 q=0.4, 
                 B=30
             )
@@ -235,9 +250,9 @@ def generate_model_response(state: State):
         print("❌ Validation failed: No query")
         return state, gr.update(value=[("❌ Please enter a query before generating response! If you just changed configuration, try reloading an example.", None)], visible=True)
     
-    # Initialize model and attribution with default configuration
-    print(f"🔧 Generating response with explanation_level: {DEFAULT_EXPLANATION_LEVEL}")
-    #llm, attr, error_msg = initialize_model_and_attr()
+    # Initialize model and attribution with current configuration
+    print(f"🔧 Generating response with explanation_level: {current_explanation_level or DEFAULT_EXPLANATION_LEVEL}, top_k: {current_top_k or 3}")
+    llm, attr, error_msg = initialize_model_and_attr()
     
     if llm is None or attr is None:
         error_text = error_msg if error_msg else "Model initialization failed!"
@@ -366,7 +381,7 @@ def unified_response_handler(response_text: str, state: State):
             )
         
         # Initialize model and generate response
-        #llm, attr, error_msg = initialize_model_and_attr()
+        llm, attr, error_msg = initialize_model_and_attr()
         
         if llm is None:
             error_text = error_msg if error_msg else "Model initialization failed!"
@@ -464,7 +479,7 @@ def basic_get_scores_and_sources_full_response(state: State):
     state.explained_response_part = state.full_response
 
     # Attribution using default configuration
-    #_, attr, error_msg = initialize_model_and_attr()
+    llm, attr, error_msg = initialize_model_and_attr()
     
     if attr is None:
         error_text = error_msg if error_msg else "Traceback initialization failed!"
@@ -679,7 +694,7 @@ def basic_get_scores_and_sources(
     state.explained_response_part = selected_text
 
     # Attribution using default configuration
-    #_, attr, error_msg = initialize_model_and_attr()
+    llm, attr, error_msg = initialize_model_and_attr()
     
     if attr is None:
         error_text = error_msg if error_msg else "Traceback initialization failed!"