update

app.py

@@ -25,8 +25,6 @@ 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!
@@ -84,46 +82,10 @@ current_attr = None
 current_model_path = None
 current_explanation_level = None
 current_api_key = None
-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, current_llm
-    
-    # Convert top_k to int
-    top_k = int(top_k)
-    
-    # Check if configuration has changed
-    config_changed = (current_explanation_level != explanation_level or 
-                     current_top_k != top_k)
-    
-    if config_changed:
-        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 both model and attribution to force complete reinitialization
-        current_llm = None
-        current_attr = None
-        
-        # 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")
 
 def initialize_model_and_attr():
     """Initialize model and attribution with default configuration"""
-    global current_llm, current_attr, current_model_path, current_explanation_level, current_api_key, current_top_k
+    global current_llm, current_attr, current_model_path, current_explanation_level, current_api_key
     
     try:
         # Check if we need to reinitialize the model
@@ -133,7 +95,7 @@ def initialize_model_and_attr():
         
         # Check if we need to update attribution
         need_attr_update = (current_attr is None or 
-                          current_explanation_level != (current_explanation_level or DEFAULT_EXPLANATION_LEVEL) or 
+                          current_explanation_level != DEFAULT_EXPLANATION_LEVEL or 
                           need_model_update)
         
         if need_model_update:
@@ -144,20 +106,15 @@ def initialize_model_and_attr():
             current_api_key = effective_api_key
             
         if need_attr_update:
-            # 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}")
+            print(f"Initializing context traceback with explanation level: {DEFAULT_EXPLANATION_LEVEL}")
             current_attr = AttnTraceAttribution(
                 current_llm, 
-                explanation_level= DEFAULT_EXPLANATION_LEVEL, 
-                K=DEFAULT_TOP_K, 
+                explanation_level=DEFAULT_EXPLANATION_LEVEL, 
+                K=3, 
                 q=0.4, 
                 B=30
             )
-            current_explanation_level = explanation_level
+            current_explanation_level = DEFAULT_EXPLANATION_LEVEL
             
         return current_llm, current_attr, None
         
@@ -250,9 +207,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 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()
+    # 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()
     
     if llm is None or attr is None:
         error_text = error_msg if error_msg else "Model initialization failed!"
@@ -381,7 +338,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!"
@@ -479,7 +436,7 @@ def basic_get_scores_and_sources_full_response(state: State):
     state.explained_response_part = state.full_response
 
     # Attribution using default configuration
-    llm, attr, error_msg = initialize_model_and_attr()
+    #_, attr, error_msg = initialize_model_and_attr()
     
     if attr is None:
         error_text = error_msg if error_msg else "Traceback initialization failed!"
@@ -694,7 +651,7 @@ def basic_get_scores_and_sources(
     state.explained_response_part = selected_text
 
     # Attribution using default configuration
-    llm, attr, error_msg = initialize_model_and_attr()
+    #_, attr, error_msg = initialize_model_and_attr()
     
     if attr is None:
         error_text = error_msg if error_msg else "Traceback initialization failed!"
@@ -1000,36 +957,6 @@ with gr.Blocks(theme=theme, css=custom_css) as demo:
             '**Color Legend for Context Traceback (by ranking):** <span style="background-color: #FF4444; color: black; padding: 2px 6px; border-radius: 4px; font-weight: 600;">Red</span> = 1st (most important) | <span style="background-color: #FF8C42; color: black; padding: 2px 6px; border-radius: 4px; font-weight: 600;">Orange</span> = 2nd | <span style="background-color: #FFD93D; color: black; padding: 2px 6px; border-radius: 4px; font-weight: 600;">Golden</span> = 3rd | <span style="background-color: #FFF280; color: black; padding: 2px 6px; border-radius: 4px; font-weight: 600;">Yellow</span> = 4th-5th | <span style="background-color: #FFF9C4; color: black; padding: 2px 6px; border-radius: 4px; font-weight: 600;">Light</span> = 6th+'
         )
         
-        # Configuration bar
-        with gr.Row():
-            with gr.Column(scale=1):
-                explanation_level_dropdown = gr.Dropdown(
-                    choices=["sentence", "paragraph", "text segment"],
-                    value="sentence",
-                    label="Explanation Level",
-                    info="How to segment the context for traceback analysis"
-                )
-            with gr.Column(scale=1):
-                top_k_dropdown = gr.Dropdown(
-                    choices=["3", "5", "10"],
-                    value="5",
-                    label="Top-K Value",
-                    info="Number of most important text segments to highlight"
-                )
-            with gr.Column(scale=1):
-                apply_config_button = gr.Button(
-                    "Apply Configuration",
-                    variant="secondary",
-                    size="sm"
-                )
-            with gr.Column(scale=2):
-                config_status_text = gr.Textbox(
-                    label="Configuration Status",
-                    value="Ready to apply configuration",
-                    interactive=False,
-                    lines=1
-                )
-        
         
         # Top section: Wide Context box with tabs
         with gr.Row():
@@ -1282,12 +1209,6 @@ with gr.Blocks(theme=theme, css=custom_css) as demo:
         outputs=[state, response_input_box, basic_response_box, basic_generate_error_box]
     )
     
-    # Configuration update handler
-    apply_config_button.click(
-        fn=update_configuration,
-        inputs=[explanation_level_dropdown, top_k_dropdown],
-        outputs=[config_status_text]
-    )
 
     # gr.Markdown(
     #     "Please do not interact with elements while generation/attribution is in progress. This may cause errors. You can refresh the page if you run into issues because of this."

app_no_config.py → app_config.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!
@@ -82,10 +84,46 @@ current_attr = None
 current_model_path = None
 current_explanation_level = None
 current_api_key = None
+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, current_llm
+    
+    # Convert top_k to int
+    top_k = int(top_k)
+    
+    # Check if configuration has changed
+    config_changed = (current_explanation_level != explanation_level or 
+                     current_top_k != top_k)
+    
+    if config_changed:
+        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 both model and attribution to force complete reinitialization
+        current_llm = None
+        current_attr = None
+        
+        # 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")
 
 def initialize_model_and_attr():
     """Initialize model and attribution with default configuration"""
-    global current_llm, current_attr, current_model_path, current_explanation_level, current_api_key
+    global current_llm, current_attr, current_model_path, current_explanation_level, current_api_key, current_top_k
     
     try:
         # Check if we need to reinitialize the model
@@ -95,7 +133,7 @@ def initialize_model_and_attr():
         
         # Check if we need to update attribution
         need_attr_update = (current_attr is None or 
-                          current_explanation_level != DEFAULT_EXPLANATION_LEVEL or 
+                          current_explanation_level != (current_explanation_level or DEFAULT_EXPLANATION_LEVEL) or 
                           need_model_update)
         
         if need_model_update:
@@ -106,15 +144,20 @@ def initialize_model_and_attr():
             current_api_key = effective_api_key
             
         if need_attr_update:
-            print(f"Initializing context traceback with explanation level: {DEFAULT_EXPLANATION_LEVEL}")
+            # 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=DEFAULT_EXPLANATION_LEVEL, 
-                K=3, 
+                explanation_level= DEFAULT_EXPLANATION_LEVEL, 
+                K=DEFAULT_TOP_K, 
                 q=0.4, 
                 B=30
             )
-            current_explanation_level = DEFAULT_EXPLANATION_LEVEL
+            current_explanation_level = explanation_level
             
         return current_llm, current_attr, None
         
@@ -207,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!"
@@ -338,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!"
@@ -436,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!"
@@ -651,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!"
@@ -957,6 +1000,36 @@ with gr.Blocks(theme=theme, css=custom_css) as demo:
             '**Color Legend for Context Traceback (by ranking):** <span style="background-color: #FF4444; color: black; padding: 2px 6px; border-radius: 4px; font-weight: 600;">Red</span> = 1st (most important) | <span style="background-color: #FF8C42; color: black; padding: 2px 6px; border-radius: 4px; font-weight: 600;">Orange</span> = 2nd | <span style="background-color: #FFD93D; color: black; padding: 2px 6px; border-radius: 4px; font-weight: 600;">Golden</span> = 3rd | <span style="background-color: #FFF280; color: black; padding: 2px 6px; border-radius: 4px; font-weight: 600;">Yellow</span> = 4th-5th | <span style="background-color: #FFF9C4; color: black; padding: 2px 6px; border-radius: 4px; font-weight: 600;">Light</span> = 6th+'
         )
         
+        # Configuration bar
+        with gr.Row():
+            with gr.Column(scale=1):
+                explanation_level_dropdown = gr.Dropdown(
+                    choices=["sentence", "paragraph", "text segment"],
+                    value="sentence",
+                    label="Explanation Level",
+                    info="How to segment the context for traceback analysis"
+                )
+            with gr.Column(scale=1):
+                top_k_dropdown = gr.Dropdown(
+                    choices=["3", "5", "10"],
+                    value="5",
+                    label="Top-K Value",
+                    info="Number of most important text segments to highlight"
+                )
+            with gr.Column(scale=1):
+                apply_config_button = gr.Button(
+                    "Apply Configuration",
+                    variant="secondary",
+                    size="sm"
+                )
+            with gr.Column(scale=2):
+                config_status_text = gr.Textbox(
+                    label="Configuration Status",
+                    value="Ready to apply configuration",
+                    interactive=False,
+                    lines=1
+                )
+        
         
         # Top section: Wide Context box with tabs
         with gr.Row():
@@ -1209,6 +1282,12 @@ with gr.Blocks(theme=theme, css=custom_css) as demo:
         outputs=[state, response_input_box, basic_response_box, basic_generate_error_box]
     )
     
+    # Configuration update handler
+    apply_config_button.click(
+        fn=update_configuration,
+        inputs=[explanation_level_dropdown, top_k_dropdown],
+        outputs=[config_status_text]
+    )
 
     # gr.Markdown(
     #     "Please do not interact with elements while generation/attribution is in progress. This may cause errors. You can refresh the page if you run into issues because of this."