Adding logs

app.py

@@ -15,6 +15,7 @@ from bertopic.representation import KeyBERTInspired
 from huggingface_hub import HfApi, InferenceClient
 from sklearn.feature_extraction.text import CountVectorizer
 from sentence_transformers import SentenceTransformer
+from torch import cuda
 
 from src.hub import create_space_with_content
 from src.templates import LLAMA_3_8B_PROMPT, SPACE_REPO_CARD_CONTENT
@@ -167,14 +168,11 @@ def generate_topics(dataset, config, split, column, plot_type):
 
     try:
         while offset < limit:
+            logging.info(f"----> Getting records from {offset=} with {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"
-            )
-
+            logging.info(f"Got {len(docs)} docs  ✓")
             embeddings = calculate_embeddings(docs)
             new_model = fit_model(docs, embeddings, n_neighbors, n_components)
 
@@ -192,14 +190,18 @@ def generate_topics(dataset, config, split, column, plot_type):
                 logging.info(f"The following topics are newly found: {new_topics}")
                 base_model = updated_model
 
+            logging.info("Reducing embeddings to 2D")
             reduced_embeddings = reduce_umap_model.fit_transform(embeddings)
             reduced_embeddings_list.append(reduced_embeddings)
 
             all_docs.extend(docs)
             reduced_embeddings_array = np.vstack(reduced_embeddings_list)
+            logging.info("Reducing embeddings to 2D ✓")
 
             topics_info = base_model.get_topic_info()
             all_topics = base_model.topics_
+            logging.info(f"Preparing topics {plot_type} plot")
+
             topic_plot = (
                 base_model.visualize_document_datamap(
                     docs=all_docs,
@@ -224,11 +226,13 @@ def generate_topics(dataset, config, split, column, plot_type):
                 if plot_type == "DataMapPlot"
                 else base_model.visualize_documents(
                     docs=all_docs,
+                    topics=all_topics,
                     reduced_embeddings=reduced_embeddings_array,
                     custom_labels=True,
                     title="",
                 )
             )
+            logging.info("Plot done ✓")
             rows_processed += len(docs)
             progress = min(rows_processed / limit, 1.0)
             logging.info(f"Progress: {progress} % - {rows_processed} of {limit}")
@@ -403,7 +407,7 @@ def generate_topics(dataset, config, split, column, plot_type):
         del (
             base_model,
             all_topics,
-            topic_info,
+            topics_info,
             topic_names_array,
             interactive_plot,
         )