Removing TextGeneration layer temporally

app.py

@@ -7,17 +7,8 @@ from gradio_huggingfacehub_search import HuggingfaceHubSearch
 from bertopic import BERTopic
 from bertopic.representation import (
     KeyBERTInspired,
-    TextGeneration,
 )
 from umap import UMAP
-from torch import cuda, bfloat16
-from transformers import (
-    BitsAndBytesConfig,
-    AutoTokenizer,
-    AutoModelForCausalLM,
-    pipeline,
-)
-from prompts import REPRESENTATION_PROMPT
 from hdbscan import HDBSCAN
 from sklearn.feature_extraction.text import CountVectorizer
 
@@ -26,7 +17,7 @@ from sentence_transformers import SentenceTransformer
 from dotenv import load_dotenv
 import os
 
-import spaces
+# import spaces
 import gradio as gr
 
 
@@ -38,8 +29,8 @@ logging.basicConfig(
     level=logging.INFO, format="%(asctime)s - %(name)s - %(levelname)s - %(message)s"
 )
 
-MAX_ROWS = 1_000
-CHUNK_SIZE = 300
+MAX_ROWS = 5_000
+CHUNK_SIZE = 1_000
 
 
 session = requests.Session()
@@ -47,71 +38,7 @@ sentence_model = SentenceTransformer("all-MiniLM-L6-v2")
 keybert = KeyBERTInspired()
 vectorizer_model = CountVectorizer(stop_words="english")
 
-model_id = "meta-llama/Llama-2-7b-chat-hf"
-device = f"cuda:{cuda.current_device()}" if cuda.is_available() else "cpu"
-logging.info(device)
-
-bnb_config = BitsAndBytesConfig(
-    load_in_4bit=True,  # 4-bit quantization
-    bnb_4bit_quant_type="nf4",  # Normalized float 4
-    bnb_4bit_use_double_quant=True,  # Second quantization after the first
-    bnb_4bit_compute_dtype=bfloat16,  # Computation type
-)
-
-tokenizer = AutoTokenizer.from_pretrained(model_id)
-
-model = AutoModelForCausalLM.from_pretrained(
-    model_id,
-    trust_remote_code=True,
-    quantization_config=bnb_config,
-    device_map="auto",
-    offload_folder="offload",  # Offloading part of the model to CPU to save GPU memory
-)
-
-# Enable gradient checkpointing for memory efficiency during backprop?
-model.gradient_checkpointing_enable()
-
-generator = pipeline(
-    model=model,
-    tokenizer=tokenizer,
-    task="text-generation",
-    temperature=0.1,
-    max_new_tokens=200,  # Reduced max_new_tokens to limit memory consumption
-    repetition_penalty=1.1,
-)
-
-llama2 = TextGeneration(generator, prompt=REPRESENTATION_PROMPT)
-representation_model = {
-    "KeyBERT": keybert,
-    "Llama2": llama2,
-}
-
-# TODO: It should be proporcional to the number of rows
-# For small datasets (1-200 rows) it worked fine with 2 neighbors
-N_NEIGHBORS = 15
-
-umap_model = UMAP(
-    n_neighbors=N_NEIGHBORS,
-    n_components=5,
-    min_dist=0.0,
-    metric="cosine",
-    random_state=42,
-)
-
-hdbscan_model = HDBSCAN(
-    min_cluster_size=N_NEIGHBORS,
-    metric="euclidean",
-    cluster_selection_method="eom",
-    prediction_data=True,
-)
-
-reduce_umap_model = UMAP(
-    n_neighbors=N_NEIGHBORS,
-    n_components=2,
-    min_dist=0.0,
-    metric="cosine",
-    random_state=42,
-)
+representation_model = KeyBERTInspired()
 
 global_topic_model = None
 
@@ -151,16 +78,30 @@ def get_docs_from_parquet(parquet_urls, column, offset, limit):
     return df[column].tolist()
 
 
-@spaces.GPU
-# TODO: Modify batch size to reduce memory consumption during embedding calculation, which value is better?
+# @spaces.GPU
 def calculate_embeddings(docs):
     return sentence_model.encode(docs, show_progress_bar=True, batch_size=32)
 
 
-@spaces.GPU
-def fit_model(docs, embeddings):
+# @spaces.GPU
+def fit_model(docs, embeddings, n_neighbors):
     global global_topic_model
 
+    umap_model = UMAP(
+        n_neighbors=n_neighbors,
+        n_components=5,
+        min_dist=0.0,
+        metric="cosine",
+        random_state=42,
+    )
+
+    hdbscan_model = HDBSCAN(
+        min_cluster_size=n_neighbors,
+        metric="euclidean",
+        cluster_selection_method="eom",
+        prediction_data=True,
+    )
+
     new_model = BERTopic(
         "english",
         # Sub-models
@@ -172,7 +113,7 @@ def fit_model(docs, embeddings):
         # Hyperparameters
         top_n_words=10,
         verbose=True,
-        min_topic_size=15,  # TODO: Should this value be coherent with N_NEIGHBORS?
+        min_topic_size=n_neighbors,  # TODO: Should this value be coherent with N_NEIGHBORS?
     )
     logging.info("Fitting new model")
     new_model.fit(docs, embeddings)
@@ -183,6 +124,10 @@ def fit_model(docs, embeddings):
     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}"
@@ -193,6 +138,16 @@ 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)
+
+    reduce_umap_model = UMAP(
+        n_neighbors=n_neighbors,
+        n_components=2,
+        min_dist=0.0,
+        metric="cosine",
+        random_state=42,
+    )
+
     offset = 0
     rows_processed = 0
 
@@ -201,8 +156,8 @@ def generate_topics(dataset, config, split, column, nested_column):
     reduced_embeddings_list = []
     topics_info, topic_plot = None, None
     yield (
-        gr.DataFrame(interactive=False, visible=True),
-        gr.Plot(visible=True),
+        gr.DataFrame(value=[], interactive=False, visible=True),
+        gr.Plot(value=None, visible=True),
         gr.Label(
             {f"⚙️ Generating topics {dataset}": rows_processed / limit}, visible=True
         ),
@@ -217,7 +172,7 @@ def generate_topics(dataset, config, split, column, nested_column):
         )
 
         embeddings = calculate_embeddings(docs)
-        fit_model(docs, embeddings)
+        fit_model(docs, embeddings, n_neighbors)
 
         if base_model is None:
             base_model = global_topic_model
@@ -230,13 +185,6 @@ def generate_topics(dataset, config, split, column, nested_column):
             logging.info(f"The following topics are newly found: {new_topics}")
             base_model = updated_model
 
-        repr_model_topics = {
-            key: label[0][0].split("\n")[0]
-            for key, label in base_model.get_topics(full=True)["Llama2"].items()
-        }
-
-        base_model.set_topic_labels(repr_model_topics)
-
         reduced_embeddings = reduce_umap_model.fit_transform(embeddings)
         reduced_embeddings_list.append(reduced_embeddings)
 
@@ -249,8 +197,6 @@ def generate_topics(dataset, config, split, column, nested_column):
             custom_labels=True,
         )
 
-        logging.info(f"Topics: {repr_model_topics}")
-
         rows_processed += len(docs)
         progress = min(rows_processed / limit, 1.0)
         logging.info(f"Progress: {progress} % - {rows_processed} of {limit}")