Upload folder using huggingface_hub

app.py

@@ -1,4 +1,5 @@
 import gradio as gr
+import spaces
 import os
 import shutil
 import time
@@ -306,10 +307,10 @@ class HackerNewsFineTuner:
             gr.Markdown("This project provides a set of tools to fine-tune [EmbeddingGemma](https://huggingface.co/google/embeddinggemma-300m) to understand your personal taste in Hacker News titles and then use it to score and rank new articles based on their \"vibe\". The core idea is to measure the \"vibe\" of a news title by calculating the semantic similarity between its embedding and the embedding of a fixed anchor phrase, **`MY_FAVORITE_NEWS`**.<br>See [README](https://huggingface.co/spaces/google/embeddinggemma-modkit/blob/main/README.md) for more details.")
             with gr.Tab("🚀 Fine-Tuning & Evaluation"):
                 self._build_training_interface()
-            with gr.Tab("💡 News Vibe Check"):
-                self._build_vibe_check_interface()
             with gr.Tab("📰 Hacker News Mood Reader"):
                 self._build_mood_reader_interface()
+            with gr.Tab("💡 Similarity Check"):
+                self._build_vibe_check_interface()
         return demo
 
     def _build_training_interface(self):

model_trainer.py

@@ -1,3 +1,4 @@
+import spaces
 from huggingface_hub import login
 from sentence_transformers import SentenceTransformer, util
 from datasets import Dataset
@@ -21,8 +22,8 @@ def load_embedding_model(model_name: str) -> SentenceTransformer:
     """Initializes the Sentence Transformer model."""
     print(f"Loading Sentence Transformer model: {model_name}")
     try:
-        model = SentenceTransformer(model_name)
-        print("Model loaded successfully.")
+        model = SentenceTransformer(model_name, model_kwargs={"device_map": "auto"})
+        print(f"Model loaded successfully. {model.device}")
         return model
     except Exception as e:
         print(f"Error loading Sentence Transformer model {model_name}: {e}")
@@ -71,6 +72,7 @@ class EvaluationCallback(TrainerCallback):
         print(f"\n{self.search_fn()}\n")
 
 
+@spaces.GPU
 def train_with_dataset(
     model: SentenceTransformer,
     dataset: List[List[str]],
@@ -129,4 +131,4 @@ def train_with_dataset(
     # Save the final fine-tuned model
     trainer.save_model()
 
-    print(f"Model saved locally to: {output_dir}")
+    print(f"Model saved locally to: {output_dir}")