Move hub code to another file

app.py

@@ -15,11 +15,9 @@ from bertopic import BERTopic
 from bertopic.representation import KeyBERTInspired
 from bertopic.representation import TextGeneration
 
-
-from huggingface_hub import HfApi, SpaceCard
+from huggingface_hub import HfApi
 from sklearn.feature_extraction.text import CountVectorizer
 from sentence_transformers import SentenceTransformer
-from templates import REPRESENTATION_PROMPT, SPACE_REPO_CARD_CONTENT
 from torch import cuda, bfloat16
 from transformers import (
     BitsAndBytesConfig,
@@ -28,13 +26,8 @@ from transformers import (
     pipeline,
 )
 
-"""
-TODOs:
-- Improve representation layer (Try with llamacpp or TextGeneration)
-- Make it run on Zero GPU
-- Try with more rows (Current: 50_000/10_000 -> Minimal Targett: 1_000_000/20_000)
-- Export interactive plots and serve their HTML content (It doesn't work with gr.HTML)
-"""
+from src.hub import create_space_with_content
+from src.templates import REPRESENTATION_PROMPT, SPACE_REPO_CARD_CONTENT
 
 load_dotenv()
 HF_TOKEN = os.getenv("HF_TOKEN")
@@ -189,54 +182,7 @@ def fit_model(docs, embeddings, n_neighbors, n_components):
     return new_model
 
 
-def _push_to_hub(
-    dataset_id,
-    file_path,
-):
-    logging.info(f"Pushing file to hub: {dataset_id} on file {file_path}")
-
-    file_name = file_path.split("/")[-1]
-    try:
-        logging.info(f"About to push {file_path} - {dataset_id}")
-        api.upload_file(
-            path_or_fileobj=file_path,
-            path_in_repo=file_name,
-            repo_id=EXPORTS_REPOSITORY,
-            repo_type="dataset",
-        )
-    except Exception as e:
-        logging.info("Failed to push file", e)
-        raise
-
-
-def create_space_with_content(dataset_id, html_file_path):
-    repo_id = f"{DATASETS_TOPICS_ORGANIZATION}/{dataset_id.replace('/', '-')}"
-    logging.info(f"Creating space with content: {repo_id} on file {html_file_path}")
-    api.create_repo(
-        repo_id=repo_id,
-        repo_type="space",
-        private=False,
-        exist_ok=True,
-        token=HF_TOKEN,
-        space_sdk="static",
-    )
-
-    SpaceCard(
-        content=SPACE_REPO_CARD_CONTENT.format(dataset_id=dataset_id)
-    ).push_to_hub(repo_id=repo_id, repo_type="space", token=HF_TOKEN)
-
-    api.upload_file(
-        path_or_fileobj=html_file_path,
-        path_in_repo="index.html",
-        repo_type="space",
-        repo_id=repo_id,
-        token=HF_TOKEN,
-    )
-    logging.info(f"Space creation done")
-    return repo_id
-
-
-@spaces.GPU(duration=60*5)
+@spaces.GPU(duration=60 * 5)
 def generate_topics(dataset, config, split, column, nested_column, plot_type):
     logging.info(
         f"Generating topics for {dataset=} {config=} {split=} {column=} {nested_column=} {plot_type=}"
@@ -374,8 +320,6 @@ def generate_topics(dataset, config, split, column, nested_column, plot_type):
     else:
         topic_plot.write_image(plot_png)
 
-    _push_to_hub(dataset, plot_png)
-
     all_topics, _ = base_model.transform(all_docs)
     topic_info = base_model.get_topic_info()
 
@@ -407,10 +351,20 @@ def generate_topics(dataset, config, split, column, nested_column, plot_type):
     with open(html_file_path, "w", encoding="utf-8") as html_file:
         html_file.write(html_content)
 
-    space_id = create_space_with_content(dataset, html_file_path)
+    repo_id = f"{DATASETS_TOPICS_ORGANIZATION}/{dataset.replace('/', '-')}"
+
+    space_id = create_space_with_content(
+        api=api,
+        repo_id=repo_id,
+        dataset_id=dataset,
+        html_file_path=html_file_path,
+        plot_file_path=plot_png,
+        space_card=SPACE_REPO_CARD_CONTENT,
+        token=HF_TOKEN,
+    )
 
     plot_png_link = (
-        f"https://huggingface.co/datasets/{EXPORTS_REPOSITORY}/blob/main/{plot_png}"
+        f"https://huggingface.co/spaces/{space_id}/blob/main/static_plot.png"
     )
 
     space_link = f"https://huggingface.co/spaces/{space_id}"

src/hub.py

@@ -0,0 +1,51 @@
+import logging
+from huggingface_hub import HfApi, SpaceCard
+
+
+def create_space_with_content(
+    api: HfApi,
+    repo_id: str,
+    dataset_id: str,
+    html_file_path: str,
+    plot_file_path: str,
+    space_card: str,
+    token: str,
+):
+    logging.info(f"Creating space with content: {repo_id} on file {html_file_path}")
+    api.create_repo(
+        repo_id=repo_id,
+        repo_type="space",
+        private=False,
+        exist_ok=True,
+        token=token,
+        space_sdk="static",
+    )
+
+    SpaceCard(content=space_card.format(dataset_id=dataset_id)).push_to_hub(
+        repo_id=repo_id, repo_type="space", token=token
+    )
+
+    api.upload_file(
+        path_or_fileobj=html_file_path,
+        path_in_repo="index.html",
+        repo_type="space",
+        repo_id=repo_id,
+        token=token,
+    )
+
+    logging.info(f"Pushing file to hub: {dataset_id} on file {plot_file_path}")
+
+    try:
+        logging.info(f"About to push {plot_file_path} - {dataset_id}")
+        api.upload_file(
+            path_or_fileobj=plot_file_path,
+            path_in_repo="static_plot.png",
+            repo_id=repo_id,
+            repo_type="space",
+        )
+    except Exception as e:
+        logging.info("Failed to push file", e)
+        raise
+
+    logging.info(f"Space creation done")
+    return repo_id

templates.py → src/templates.py

@@ -39,4 +39,5 @@ datasets:
 - {dataset_id}
 ---
 
+![Static plot](static_plot.png)
 """