Change parameters by dataset size

app.py

@@ -2,12 +2,10 @@ import requests
 import logging
 import duckdb
 import numpy as np
-
+from torch import cuda
 from gradio_huggingfacehub_search import HuggingfaceHubSearch
 from bertopic import BERTopic
-from bertopic.representation import (
-    KeyBERTInspired,
-)
+from bertopic.representation import KeyBERTInspired
 from umap import UMAP
 from hdbscan import HDBSCAN
 from sklearn.feature_extraction.text import CountVectorizer
@@ -21,6 +19,11 @@ import os
 import gradio as gr
 
 
+"""
+TODOs:
+- Try for small dataset <1000 rows
+"""
+
 load_dotenv()
 HF_TOKEN = os.getenv("HF_TOKEN")
 assert HF_TOKEN is not None, "You need to set HF_TOKEN in your environment variables"
@@ -55,7 +58,7 @@ def get_split_rows(dataset, config, split):
         None,
     )
     if split_size is None:
-        raise Exception(f"Error fetching split{split} in config {config}")
+        raise Exception(f"Error fetching split {split} in config {config}")
     return split_size["num_rows"]
 
 
@@ -83,27 +86,37 @@ def calculate_embeddings(docs):
     return sentence_model.encode(docs, show_progress_bar=True, batch_size=32)
 
 
+# Adjust n_neighbors and n_components based on dataset size
+def calculate_n_neighbors_and_components(n_rows):
+    # Ensure n_neighbors is proportional to the dataset size, with reasonable limits
+    n_neighbors = min(max(n_rows // 20, 15), 100)
+    n_components = 10 if n_rows > 1000 else 5  # Higher components for larger datasets
+    return n_neighbors, n_components
+
+
 # @spaces.GPU
-def fit_model(docs, embeddings, n_neighbors):
+def fit_model(docs, embeddings, n_neighbors, n_components):
     global global_topic_model
 
     umap_model = UMAP(
         n_neighbors=n_neighbors,
-        n_components=5,
+        n_components=n_components,
         min_dist=0.0,
         metric="cosine",
         random_state=42,
     )
 
     hdbscan_model = HDBSCAN(
-        min_cluster_size=n_neighbors,
+        min_cluster_size=max(
+            5, n_neighbors // 2
+        ),  # Reducing min_cluster_size for fewer outliers
         metric="euclidean",
         cluster_selection_method="eom",
         prediction_data=True,
     )
 
     new_model = BERTopic(
-        "english",
+        language="english",
         # Sub-models
         embedding_model=sentence_model,
         umap_model=umap_model,
@@ -113,7 +126,7 @@ def fit_model(docs, embeddings, n_neighbors):
         # Hyperparameters
         top_n_words=10,
         verbose=True,
-        min_topic_size=n_neighbors,  # TODO: Should this value be coherent with N_NEIGHBORS?
+        min_topic_size=n_neighbors,  # Coherent with n_neighbors?
     )
     logging.info("Fitting new model")
     new_model.fit(docs, embeddings)
@@ -124,10 +137,6 @@ def fit_model(docs, embeddings, n_neighbors):
     logging.info("Global model updated")
 
 
-def calculate_n_neighbors(n_rows):
-    return max(n_rows // 20, 2)
-
-
 def generate_topics(dataset, config, split, column, nested_column):
     logging.info(
         f"Generating topics for {dataset} with config {config} {split} {column} {nested_column}"
@@ -138,11 +147,11 @@ def generate_topics(dataset, config, split, column, nested_column):
     logging.info(f"Split rows: {split_rows}")
 
     limit = min(split_rows, MAX_ROWS)
-    n_neighbors = calculate_n_neighbors(limit)
+    n_neighbors, n_components = calculate_n_neighbors_and_components(limit)
 
     reduce_umap_model = UMAP(
         n_neighbors=n_neighbors,
-        n_components=2,
+        n_components=2,  # For visualization, keeping it at 2 (2D)
         min_dist=0.0,
         metric="cosine",
         random_state=42,
@@ -172,7 +181,7 @@ def generate_topics(dataset, config, split, column, nested_column):
         )
 
         embeddings = calculate_embeddings(docs)
-        fit_model(docs, embeddings, n_neighbors)
+        fit_model(docs, embeddings, n_neighbors, n_components)
 
         if base_model is None:
             base_model = global_topic_model