Separate functions

app.py

@@ -41,6 +41,34 @@ def get_docs_from_parquet(parquet_urls, column, offset, limit):
 
 
 @spaces.GPU
+def calculate_embeddings(sentence_model, docs):
+    embeddings = sentence_model.encode(docs, show_progress_bar=True, batch_size=100)
+    logging.info(f"Embeddings shape: {embeddings.shape}")
+    return embeddings
+
+
+@spaces.GPU
+def fit_model(base_model, sentence_model, representation_model, docs, embeddings):
+    new_model = BERTopic(
+        "english",
+        embedding_model=sentence_model,
+        representation_model=representation_model,
+        min_topic_size=15,  # umap_model=umap_model, hdbscan_model=hdbscan_model
+    )
+    logging.info("Fitting new model")
+    new_model.fit(docs, embeddings)
+    logging.info("End fitting new model")
+    if base_model is None:
+        return new_model, new_model
+
+    updated_model = BERTopic.merge_models([base_model, new_model])
+    nr_new_topics = len(set(updated_model.topics_)) - len(set(base_model.topics_))
+    new_topics = list(updated_model.topic_labels_.values())[-nr_new_topics:]
+    logging.info("The following topics are newly found:")
+    logging.info(f"{new_topics}\n")
+    return updated_model, new_model
+
+
 def generate_topics(dataset, config, split, column, nested_column):
     logging.info(
         f"Generating topics for {dataset} with config {config} {split} {column} {nested_column}"
@@ -67,43 +95,45 @@ def generate_topics(dataset, config, split, column, nested_column):
     while True:
         docs = get_docs_from_parquet(parquet_urls, column, offset, chunk_size)
         logging.info(f"------------> New chunk data {offset=} {chunk_size=}")
-        embeddings = sentence_model.encode(docs, show_progress_bar=True, batch_size=100)
-        logging.info(f"Embeddings shape: {embeddings.shape}")
+        embeddings = calculate_embeddings(sentence_model, docs)
         offset = offset + chunk_size
         if not docs or offset >= limit:
             break
 
-        new_model = BERTopic(
-            "english",
-            embedding_model=sentence_model,
-            representation_model=representation_model,
-            min_topic_size=15,  # umap_model=umap_model, hdbscan_model=hdbscan_model
+        # new_model = BERTopic(
+        #     "english",
+        #     embedding_model=sentence_model,
+        #     representation_model=representation_model,
+        #     min_topic_size=15,  # umap_model=umap_model, hdbscan_model=hdbscan_model
+        # )
+        # logging.info("Fitting new model")
+        # new_model.fit(docs, embeddings)
+        # logging.info("End fitting new model")
+        # if base_model is not None:
+        #     updated_model = BERTopic.merge_models([base_model, new_model])
+        #     nr_new_topics = len(set(updated_model.topics_)) - len(
+        #         set(base_model.topics_)
+        #     )
+        #     new_topics = list(updated_model.topic_labels_.values())[-nr_new_topics:]
+        #     logging.info("The following topics are newly found:")
+        #     logging.info(f"{new_topics}\n")
+        #     base_model = updated_model
+        # else:
+        #     base_model = new_model
+        # logging.info(base_model.get_topic_info())
+        base_model, new_model = fit_model(
+            base_model, sentence_model, representation_model, docs, embeddings
         )
-        logging.info("Fitting new model")
-        new_model.fit(docs, embeddings)
-        logging.info("End fitting new model")
-        if base_model is not None:
-            updated_model = BERTopic.merge_models([base_model, new_model])
-            nr_new_topics = len(set(updated_model.topics_)) - len(
-                set(base_model.topics_)
-            )
-            new_topics = list(updated_model.topic_labels_.values())[-nr_new_topics:]
-            logging.info("The following topics are newly found:")
-            logging.info(f"{new_topics}\n")
-            base_model = updated_model
-        else:
-            base_model = new_model
-        logging.info(base_model.get_topic_info())
-        reduced_embeddings = UMAP(
-            n_neighbors=10, n_components=2, min_dist=0.0, metric="cosine"
-        ).fit_transform(embeddings)
-        logging.info(f"Reduced embeddings shape: {reduced_embeddings.shape}")
+        # reduced_embeddings = UMAP(
+        #     n_neighbors=10, n_components=2, min_dist=0.0, metric="cosine"
+        # ).fit_transform(embeddings)
+        # logging.info(f"Reduced embeddings shape: {reduced_embeddings.shape}")
         yield (
             base_model.get_topic_info(),
             new_model.visualize_documents(
                 docs, embeddings=embeddings
-            ), # TODO: Visualize the merged models  
-        ) 
+            ),  # TODO: Visualize the merged models
+        )
     logging.info("Finished processing all data")
     return base_model.get_topic_info(), base_model.visualize_topics()