feat: improve HuggingFace dataset loading with better error handling

src/generator/generator.py

@@ -4,13 +4,15 @@ from langchain_openai import OpenAIEmbeddings, ChatOpenAI
 from langchain_chroma import Chroma
 from langchain.chains import ConversationalRetrievalChain
 from langchain.prompts import PromptTemplate
+from datasets import load_dataset
+import tempfile
+import shutil
 from config.settings import Settings
 
 
 class LyricGenerator:
     def __init__(self):
         """Initialize the generator with embeddings"""
-        # Use Settings to determine the correct path based on environment
         self.embeddings_dir = Settings.get_embeddings_path()
         self.embeddings = OpenAIEmbeddings()
         self.vector_store = None
@@ -21,50 +23,53 @@ class LyricGenerator:
 
     def _load_embeddings(self) -> None:
         """Load existing embeddings based on environment"""
-        print(f"Loading embeddings from: {self.embeddings_dir}")
-        print(f"Deployment mode: {Settings.DEPLOYMENT_MODE}")
-        print(f"Directory exists: {self.embeddings_dir.exists()}")
-        
-        if not self.embeddings_dir.exists():
-            if Settings.is_huggingface():
-                # List contents of root directory in HF
-                try:
-                    print("Contents of root directory:")
-                    import os
-                    print(os.listdir("/"))
-                except Exception as e:
-                    print(f"Error listing root: {str(e)}")
+        if Settings.is_huggingface():
+            try:
+                print(f"Loading dataset from HuggingFace: {Settings.HF_DATASET}")
+                # Create a temporary directory to store the dataset
+                with tempfile.TemporaryDirectory() as temp_dir:
+                    # Load the dataset
+                    dataset = load_dataset(
+                        path=Settings.HF_DATASET,
+                        split="train",
+                        token=Settings.HF_TOKEN  # Add token for private dataset access
+                    )
+                    print("Dataset loaded, downloading files...")
+                    # Download the chroma files
+                    dataset.download(temp_dir)
+                    chroma_path = Path(temp_dir) / "chroma"
+                    print(f"Looking for Chroma files in: {chroma_path}")
+                    if not chroma_path.exists():
+                        raise RuntimeError(f"Chroma directory not found in dataset at {chroma_path}")
                     
-                raise RuntimeError(
-                    "Embeddings not found in HF storage. "
-                    "Please ensure they are uploaded to the dataset."
-                )
-            else:
+                    # Use the downloaded files
+                    self.vector_store = Chroma(
+                        persist_directory=str(chroma_path),
+                        embedding_function=self.embeddings,
+                        collection_name="lyrics"
+                    )
+                    print("Successfully loaded vector store from dataset")
+            except Exception as e:
+                raise RuntimeError(f"Failed to load HuggingFace dataset: {str(e)}")
+        else:
+            if not self.embeddings_dir.exists():
                 raise RuntimeError(
                     "Embeddings not found locally. "
                     "Please run process_lyrics.py first."
                 )
             
-        try:
-            # List contents of embeddings directory
             try:
-                print("Contents of embeddings directory:")
-                print(os.listdir(str(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"
+                )
             except Exception as e:
-                print(f"Error listing embeddings dir: {str(e)}")
-                
-            # Load vector store using environment-aware settings
-            self.vector_store = Chroma(
-                persist_directory=str(self.embeddings_dir),
-                embedding_function=self.embeddings,
-                collection_name="lyrics"
-            )
-            
-            # Setup QA chain
-            self._setup_qa_chain()
-            
-        except Exception as e:
-            raise RuntimeError(f"Failed to load embeddings: {str(e)}")
+                raise RuntimeError(f"Failed to load local embeddings: {str(e)}")
+        
+        # Setup QA chain
+        self._setup_qa_chain()
 
     def _setup_qa_chain(self) -> None:
         """Initialize the QA chain for generating lyrics"""