Add tqdm for attribute_context

app.py

@@ -1,5 +1,6 @@
 import os
 import re
+import threading
 
 import bm25s
 import gradio as gr
@@ -12,6 +13,7 @@ from lxt.functional import add2, mul2, softmax
 from lxt.models.llama import LlamaForCausalLM, attnlrp
 from rerankers import Reranker
 from style import custom_css
+from tqdm import tqdm
 from transformers import AutoTokenizer
 
 from inseq import load_model, register_step_function
@@ -120,7 +122,7 @@ def generate(
     rag_setting,
     custom_context,
     model_size,
-    progress=gr.Progress(),
+    progress=gr.Progress(track_tqdm=True),
 ):
     global model, model_id
     if rag_setting == "Use Custom Context":
@@ -144,7 +146,8 @@ def generate(
     if model is None or model.model_name != curr_model_id:
         progress(0.2, desc="Loading model...")
         model = get_model(model_size)
-    progress(0.3, desc="Attributing with LXT...")
+    estimated_time = 20
+    tstep = 1
     lm_rag_prompting_example = AttributeContextArgs(
         model_name_or_path=model_id,
         input_context_text="\n\n".join(docs),
@@ -169,7 +172,20 @@ def generate(
         save_path=os.path.join(os.path.dirname(__file__), "outputs/output.json"),
         viz_path=os.path.join(os.path.dirname(__file__), "outputs/output.html"),
     )
-    out = attribute_context_with_model(lm_rag_prompting_example, model)
+
+    ret = [None]
+
+    def run_attribute_context():
+        ret[0] = attribute_context_with_model(lm_rag_prompting_example, model)
+
+    thread = threading.Thread(target=run_attribute_context)
+    pbar = tqdm(total=estimated_time, desc="Attributing with LXT...")
+    thread.start()
+    while thread.is_alive():
+        thread.join(timeout=tstep)
+        pbar.update(tstep)
+    pbar.close()
+    out = ret[0]
     html = visualize_attribute_context(out, show_viz=False, return_html=True)
     return [
         gradio_iframe.iFrame(html, height=500, visible=True),