Adding progress bar

app.py

@@ -38,6 +38,9 @@ logging.basicConfig(
     level=logging.INFO, format="%(asctime)s - %(name)s - %(levelname)s - %(message)s"
 )
 
+MAX_ROWS = 1_000
+CHUNK_SIZE = 300
+
 
 session = requests.Session()
 sentence_model = SentenceTransformer("all-MiniLM-L6-v2")
@@ -113,6 +116,22 @@ reduce_umap_model = UMAP(
 global_topic_model = None
 
 
+def get_split_rows(dataset, config, split):
+    config_size = session.get(
+        f"https://datasets-server.huggingface.co/size?dataset={dataset}&config={config}",
+        timeout=20,
+    ).json()
+    if "error" in config_size:
+        raise Exception(f"Error fetching config size: {config_size['error']}")
+    split_size = next(
+        (s for s in config_size["size"]["splits"] if s["split"] == split),
+        None,
+    )
+    if split_size is None:
+        raise Exception(f"Error fetching split{split} in config {config}")
+    return split_size["num_rows"]
+
+
 def get_parquet_urls(dataset, config, split):
     parquet_files = session.get(
         f"https://datasets-server.huggingface.co/parquet?dataset={dataset}&config={config}&split={split}",
@@ -170,9 +189,13 @@ def generate_topics(dataset, config, split, column, nested_column):
     )
 
     parquet_urls = get_parquet_urls(dataset, config, split)
-    limit = 1_000
-    chunk_size = 300
+    split_rows = get_split_rows(dataset, config, split)
+    logging.info(f"Split rows: {split_rows}")
+
+    limit = min(split_rows, MAX_ROWS)
     offset = 0
+    rows_processed = 0
+
     base_model = None
     all_docs = []
     reduced_embeddings_list = []
@@ -180,15 +203,17 @@ def generate_topics(dataset, config, split, column, nested_column):
     yield (
         gr.DataFrame(interactive=False, visible=True),
         gr.Plot(visible=True),
-        gr.Label({f"⚙️ Generating topics {dataset}": 0.0}, visible=True),
+        gr.Label(
+            {f"⚙️ Generating topics {dataset}": rows_processed / limit}, visible=True
+        ),
     )
     while offset < limit:
-        docs = get_docs_from_parquet(parquet_urls, column, offset, chunk_size)
+        docs = get_docs_from_parquet(parquet_urls, column, offset, CHUNK_SIZE)
         if not docs:
             break
 
         logging.info(
-            f"----> Processing chunk: {offset=} {chunk_size=} with {len(docs)} docs"
+            f"----> Processing chunk: {offset=} {CHUNK_SIZE=} with {len(docs)} docs"
         )
 
         embeddings = calculate_embeddings(docs)
@@ -225,15 +250,17 @@ def generate_topics(dataset, config, split, column, nested_column):
         )
 
         logging.info(f"Topics: {repr_model_topics}")
-        progress = min(offset / limit, 1.0)
 
+        rows_processed += len(docs)
+        progress = min(rows_processed / limit, 1.0)
+        logging.info(f"Progress: {progress} % - {rows_processed} of {limit}")
         yield (
             topics_info,
             topic_plot,
             gr.Label({f"⚙️ Generating topics {dataset}": progress}, visible=True),
         )
 
-        offset += chunk_size
+        offset += CHUNK_SIZE
 
     logging.info("Finished processing all data")
     cuda.empty_cache()  # Clear cache at the end of each chunk