Updated to load HF dataset into HF Space on run

app.py

@@ -3,6 +3,7 @@ import os
 import streamlit as st
 from src.generator.generator import LyricGenerator
 from config.settings import Settings
+from pathlib import Path
 
 # Set SQLite path for local development
 if not Settings.is_huggingface():
@@ -11,14 +12,43 @@ if not Settings.is_huggingface():
 def initialize_generator():
     """Initialize the generator with proper error handling"""
     try:
+        print("\n=== Initializing Generator ===")
+        # Check for HuggingFace environment requirements
+        if Settings.is_huggingface():
+            print("Running in HuggingFace environment")
+            print("Checking environment requirements...")
+            
+            if not Settings.HF_TOKEN:
+                error_msg = "HuggingFace token not found. Please set HF_TOKEN in Space secrets."
+                print(f"Error: {error_msg}")
+                st.error(error_msg)
+                return None
+            else:
+                print("HF_TOKEN found in environment")
+            
+            # Ensure persistent storage directory exists
+            storage_path = Path("/data/processed/embeddings")
+            print(f"Setting up persistent storage at: {storage_path}")
+            storage_path.mkdir(parents=True, exist_ok=True)
+            print(f"Storage directory created/verified")
+            
+            if storage_path.exists():
+                print(f"Storage directory contents: {list(storage_path.glob('**/*'))}")
+        else:
+            print("Running in local environment")
+        
         # Initialize generator
+        print("\nInitializing LyricGenerator...")
         st.info("Loading embeddings...")
         generator = LyricGenerator()
         st.success("Embeddings loaded successfully!")
         return generator
         
     except Exception as e:
-        st.error(f"Failed to initialize generator: {str(e)}")
+        error_msg = f"Failed to initialize generator: {str(e)}"
+        print(f"\nError during initialization: {error_msg}")
+        print(f"Error type: {type(e).__name__}")
+        st.error(error_msg)
         return None
 
 def main():

config/settings.py

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

src/generator/generator.py

@@ -1,10 +1,11 @@
 from typing import Dict, List, Optional
 from pathlib import Path
+import shutil
 from langchain_openai import OpenAIEmbeddings, ChatOpenAI
 from langchain_chroma import Chroma
 from langchain.chains import ConversationalRetrievalChain
 from langchain.prompts import PromptTemplate
-from huggingface_hub import snapshot_download
+from huggingface_hub import snapshot_download, hf_hub_download
 from config.settings import Settings
 
 
@@ -27,19 +28,105 @@ class LyricGenerator:
         # Load embeddings
         self._load_embeddings()
 
+    def _setup_embeddings_from_hf(self) -> None:
+        """Download and setup embeddings from HuggingFace dataset"""
+        print("\n=== Setting up embeddings from HuggingFace dataset ===")
+        try:
+            # Create necessary directories
+            chroma_dir = self.embeddings_dir / "chroma"
+            print(f"Target Chroma directory: {chroma_dir}")
+            print(f"Creating parent directory: {chroma_dir.parent}")
+            chroma_dir.parent.mkdir(parents=True, exist_ok=True)
+            
+            # Check if embeddings already exist in persistent storage
+            if not chroma_dir.exists() or not (chroma_dir / "chroma.sqlite3").exists():
+                print("\nDownloading embeddings from HuggingFace dataset...")
+                print(f"Dataset repo: {Settings.HF_DATASET}")
+                print(f"Using temp directory: /tmp/embeddings")
+                
+                # Download the entire chroma directory from the dataset
+                try:
+                    temp_dir = snapshot_download(
+                        repo_id=Settings.HF_DATASET,
+                        repo_type="dataset",
+                        token=Settings.HF_TOKEN,
+                        local_dir="/tmp/embeddings"
+                    )
+                    print(f"Download completed to: {temp_dir}")
+                except Exception as e:
+                    print(f"Error during snapshot_download: {str(e)}")
+                    raise
+                
+                temp_chroma = Path(temp_dir) / "chroma"
+                print(f"Looking for Chroma in temp dir at: {temp_chroma}")
+                
+                if not temp_chroma.exists():
+                    print(f"Contents of temp_dir: {list(Path(temp_dir).glob('**/*'))}")
+                    raise RuntimeError(
+                        f"Chroma directory not found in dataset at {temp_chroma}"
+                    )
+                
+                print(f"Found Chroma directory in download. Contents: {list(temp_chroma.glob('**/*'))}")
+                
+                # Copy the downloaded chroma directory to persistent storage
+                print(f"\nCopying embeddings to persistent storage...")
+                if chroma_dir.exists():
+                    print(f"Removing existing directory: {chroma_dir}")
+                    shutil.rmtree(chroma_dir)
+                print(f"Copying from {temp_chroma} to {chroma_dir}")
+                shutil.copytree(temp_chroma, chroma_dir)
+                print(f"Embeddings copied to persistent storage at {chroma_dir}")
+                print(f"Persistent storage contents: {list(chroma_dir.glob('**/*'))}")
+                
+                # Clean up temporary directory
+                print("\nCleaning up temporary directory...")
+                shutil.rmtree("/tmp/embeddings")
+                print("Cleanup complete")
+            else:
+                print("Embeddings already exist in persistent storage")
+                print(f"Existing contents: {list(chroma_dir.glob('**/*'))}")
+                
+        except Exception as e:
+            print(f"\n=== Error in _setup_embeddings_from_hf ===")
+            print(f"Error type: {type(e).__name__}")
+            print(f"Error message: {str(e)}")
+            print(f"Current directory structure:")
+            try:
+                print(f"Parent dir exists: {chroma_dir.parent.exists()}")
+                print(f"Parent dir contents: {list(chroma_dir.parent.glob('**/*'))}")
+            except Exception as dir_error:
+                print(f"Error checking directories: {str(dir_error)}")
+            raise RuntimeError(
+                f"Failed to setup embeddings from HuggingFace: {str(e)}"
+            )
+
     def _load_embeddings(self) -> None:
         """Load existing embeddings based on environment"""
         try:
-            print(f"Loading vector store from: {self.embeddings_dir}")
+            print("\n=== Loading Embeddings ===")
+            
+            # If in HuggingFace environment, ensure embeddings are set up
+            if Settings.is_huggingface():
+                print("HuggingFace environment detected, setting up embeddings...")
+                self._setup_embeddings_from_hf()
+            else:
+                print("Local environment detected")
+            
+            print(f"\nLoading 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():
+                print(f"Parent directory exists: {chroma_dir.parent.exists()}")
+                if chroma_dir.parent.exists():
+                    print(f"Parent directory contents: {list(chroma_dir.parent.glob('**/*'))}")
                 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():
+                print(f"Directory contents: {list(chroma_dir.glob('**/*'))}")
                 raise RuntimeError(f"Chroma database not found at {sqlite_file}")
             print(f"SQLite file size: {sqlite_file.stat().st_size / (1024*1024):.2f} MB")