update

app.py

@@ -81,22 +81,30 @@ current_model_path = None
 current_explanation_level = None
 current_api_key = None
 current_top_k = 3  # Add top-k tracking
+current_B = 30  # Add B parameter tracking
+current_q = 0.4  # Add q parameter tracking
 
-def update_configuration(explanation_level, top_k):
+def update_configuration(explanation_level, top_k, B, q):
     """Update the global configuration and reinitialize attribution if needed"""
-    global current_explanation_level, current_top_k, current_attr, current_llm
+    global current_explanation_level, current_top_k, current_B, current_q, current_attr, current_llm
     
-    # Convert top_k to int
+    # Convert parameters to appropriate types
     top_k = int(top_k)
+    B = int(B)
+    q = float(q)
     
     # Check if configuration has changed
     config_changed = (current_explanation_level != explanation_level or 
-                     current_top_k != top_k)
+                     current_top_k != top_k or
+                     current_B != B or
+                     current_q != q)
     
     if config_changed:
-        print(f"🔄 Updating configuration: explanation_level={explanation_level}, top_k={top_k}")
+        print(f"🔄 Updating configuration: explanation_level={explanation_level}, top_k={top_k}, B={B}, q={q}")
         current_explanation_level = explanation_level
         current_top_k = top_k
+        current_B = B
+        current_q = q
         
         # Reset both model and attribution to force complete reinitialization
         current_llm = None
@@ -106,7 +114,7 @@ def update_configuration(explanation_level, top_k):
         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}")
+                return gr.update(value=f"✅ Configuration updated: {explanation_level} level, top-{top_k}, B={B}, q={q}")
             else:
                 return gr.update(value=f"❌ Error reinitializing: {error_msg}")
         except Exception as e:
@@ -116,7 +124,7 @@ def update_configuration(explanation_level, top_k):
 
 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, current_top_k, current_B, current_q
     
     try:
         # Check if we need to reinitialize the model
@@ -140,15 +148,17 @@ 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
+            B = current_B or 30
+            q = current_q or 0.4
             if "segment" in explanation_level:
                 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: {explanation_level}, top_k: {top_k}, B: {B}, q: {q}")
             current_attr = AttnTraceAttribution(
                 current_llm, 
                 explanation_level= explanation_level, 
                 K=top_k, 
-                q=0.4, 
-                B=30
+                q=q, 
+                B=B
             )
             current_explanation_level = explanation_level
             
@@ -1108,6 +1118,26 @@ with gr.Blocks(theme=theme, css=custom_css) as demo:
                         info="Number of most important text segments to highlight"
                     )
             
+            with gr.Row():
+                with gr.Column(scale=1):
+                    B_slider = gr.Slider(
+                        minimum=10,
+                        maximum=100,
+                        value=30,
+                        step=5,
+                        label="B Parameter",
+                        info="Number of subsamples (higher = more accurate but slower)"
+                    )
+                with gr.Column(scale=1):
+                    q_slider = gr.Slider(
+                        minimum=0.1,
+                        maximum=0.9,
+                        value=0.4,
+                        step=0.1,
+                        label="q Parameter",
+                        info="Sub-sampling ratio (0.1-0.9)"
+                    )
+            
             with gr.Row():
                 with gr.Column(scale=1):
                     apply_config_button = gr.Button(
@@ -1134,6 +1164,14 @@ with gr.Blocks(theme=theme, css=custom_css) as demo:
               - Higher values show more context but may be less focused
               - Lower values provide more focused results but may miss some context
             
+            - **B Parameter**: Number of subsamples
+              - Higher values (50-100): More thorough analysis but slower
+              - Lower values (10-30): Faster analysis but may miss some important segments
+              - Default: 30 (good balance of speed and accuracy)
+            
+            - **q Parameter**: Sub-sampling ratio (0.1-0.9)
+
+            
             **Note**: Configuration changes will take effect immediately for new traceback operations.
             """)
             
@@ -1318,7 +1356,7 @@ with gr.Blocks(theme=theme, css=custom_css) as demo:
     # Configuration update handler
     apply_config_button.click(
         fn=update_configuration,
-        inputs=[explanation_level_dropdown, top_k_dropdown],
+        inputs=[explanation_level_dropdown, top_k_dropdown, B_slider, q_slider],
         outputs=[config_status_text]
     )