refactor: simplify to use direct Space storage instead of dataset

config/settings.py

@@ -34,7 +34,6 @@ class Settings:
     
     # HuggingFace Settings
     HF_SPACE = "SongLift/LyrGen2"
-    HF_DATASET = "SongLift/LyrGen2_DB"
     
     @classmethod
     def is_huggingface(cls) -> bool:
@@ -44,6 +43,5 @@ class Settings:
     @classmethod
     def get_embeddings_path(cls) -> Path:
         """Get appropriate embeddings path based on deployment mode"""
-        if cls.is_huggingface():
-            return Path('/tmp/chroma_db/chroma')  # Path where HF dataset is downloaded
-        return cls.EMBEDDINGS_DIR
+        # Use same structure in both environments
+        return Path("/data/processed/embeddings")

scripts/upload_embeddings.py

@@ -1,4 +1,4 @@
-"""Upload embeddings to HuggingFace"""
+"""Upload embeddings to HuggingFace Space"""
 import sys
 from pathlib import Path
 from huggingface_hub import HfApi
@@ -10,7 +10,7 @@ from config.settings import Settings
 
 
 def main():
-    """Upload embeddings directory to HuggingFace dataset"""
+    """Upload embeddings directory to HuggingFace Space"""
     print("Starting upload process...")
     
     # Print size info
@@ -21,11 +21,12 @@ def main():
     
     api = HfApi(token=Settings.HF_TOKEN)
     
-    print(f"Uploading to {Settings.HF_DATASET}...")
+    print(f"Uploading to Space: {Settings.HF_SPACE}...")
     api.upload_folder(
         folder_path=str(Settings.EMBEDDINGS_DIR),
-        repo_id=Settings.HF_DATASET,
-        repo_type="dataset"
+        repo_id=Settings.HF_SPACE,
+        repo_type="space",
+        path_in_repo="data/processed/embeddings"
     )
     
     print("Upload complete!")

src/generator/generator.py

@@ -29,88 +29,51 @@ class LyricGenerator:
 
     def _load_embeddings(self) -> None:
         """Load existing embeddings based on environment"""
-        if Settings.is_huggingface():
-            try:
-                print(f"Loading embeddings from HF dataset: {Settings.HF_DATASET}")
-                # Download dataset to Space's storage
-                local_dir = Path(snapshot_download(
-                    repo_id=Settings.HF_DATASET,
-                    repo_type="dataset",
-                    local_dir="/tmp/chroma_db"
-                ))
-                print(f"Dataset downloaded to: {local_dir}")
-                
-                # List contents to debug
-                print("Contents of downloaded directory:")
-                for path in local_dir.rglob("*"):
-                    print(f"  {path}")
-                
-                # Check if chroma directory exists
-                chroma_dir = local_dir / "chroma"
-                if not chroma_dir.exists():
-                    # Try looking in the root
-                    chroma_dir = local_dir
-                    if not (chroma_dir / "chroma.sqlite3").exists():
-                        raise RuntimeError(
-                            f"Chroma files not found in {local_dir}"
-                        )
-                
-                print(f"Loading Chroma DB from: {chroma_dir}")
-                # Initialize vector store from the cached location
-                self.vector_store = Chroma(
-                    persist_directory=str(chroma_dir),
-                    embedding_function=self.embeddings,
-                    collection_name="lyrics"
-                )
-                print("Successfully loaded vector store")
-                
-                # Verify collection has documents
-                try:
-                    collection = self.vector_store._collection
-                    count = collection.count()
-                    print(f"Collection contains {count} documents")
-                    
-                    if count == 0:
-                        # Try to peek at the collection data
-                        print("Checking collection details...")
-                        peek = collection.peek()
-                        print(f"Collection peek: {peek}")
-                        
-                        raise RuntimeError(
-                            "Chroma DB is empty. Please ensure embeddings were "
-                            "properly uploaded to the dataset."
-                        )
-                except Exception as e:
-                    print(f"Error checking collection: {str(e)}")
-                    raise RuntimeError(f"Failed to verify collection: {str(e)}")
-                
-            except Exception as e:
-                print(f"Error loading HF embeddings: {str(e)}")
-                raise RuntimeError(f"Failed to load HF embeddings: {str(e)}")
-        else:
-            if not self.embeddings_dir.exists():
+        try:
+            print(f"Loading vector store from: {self.embeddings_dir}")
+            # Check Chroma directory structure
+            chroma_dir = self.embeddings_dir / "chroma"
+            print(f"Checking Chroma directory: {chroma_dir}")
+            if not chroma_dir.exists():
+                raise RuntimeError(f"Chroma directory not found at {chroma_dir}")
+            
+            sqlite_file = chroma_dir / "chroma.sqlite3"
+            print(f"Checking SQLite file: {sqlite_file}")
+            if not sqlite_file.exists():
+                raise RuntimeError(f"Chroma database not found at {sqlite_file}")
+            print(f"SQLite file size: {sqlite_file.stat().st_size / (1024*1024):.2f} MB")
+            
+            # Load vector store using environment-aware settings
+            print("Initializing Chroma with settings:")
+            print(f"  persist_directory: {str(chroma_dir)}")
+            print(f"  collection_name: lyrics")
+            
+            self.vector_store = Chroma(
+                persist_directory=str(chroma_dir),
+                embedding_function=self.embeddings,
+                collection_name="lyrics"
+            )
+            
+            # Verify collection has documents
+            collection = self.vector_store._collection
+            count = collection.count()
+            print(f"Collection contains {count} documents")
+            
+            if count == 0:
+                print("Collection is empty, checking details...")
+                # Try to peek at the collection data
+                peek = collection.peek()
+                print(f"Collection peek: {peek}")
                 raise RuntimeError(
-                    "Embeddings not found locally. "
-                    "Please run process_lyrics.py first."
+                    "Chroma DB is empty. Please ensure embeddings "
+                    "were properly generated and uploaded."
                 )
+            else:
+                print("Successfully loaded embeddings")
             
-            try:
-                print(f"Loading local vector store from: {self.embeddings_dir}")
-                # Load vector store using environment-aware settings
-                self.vector_store = Chroma(
-                    persist_directory=str(self.embeddings_dir),
-                    embedding_function=self.embeddings,
-                    collection_name="lyrics"
-                )
-                
-                # Verify collection has documents
-                collection = self.vector_store._collection
-                count = collection.count()
-                print(f"Collection contains {count} documents")
-                
-            except Exception as e:
-                print(f"Error loading local embeddings: {str(e)}")
-                raise RuntimeError(f"Failed to load local embeddings: {str(e)}")
+        except Exception as e:
+            print(f"Error loading embeddings: {str(e)}")
+            raise RuntimeError(f"Failed to load embeddings: {str(e)}")
         
         # Setup QA chain
         self._setup_qa_chain()