feat: add comprehensive embeddings and vector store debugging

src/generator/generator.py

@@ -11,8 +11,16 @@ from config.settings import Settings
 class LyricGenerator:
     def __init__(self):
         """Initialize the generator with embeddings"""
+        print("Initializing LyricGenerator...")
         self.embeddings_dir = Settings.get_embeddings_path()
-        self.embeddings = OpenAIEmbeddings()
+        print(f"Embeddings directory: {self.embeddings_dir}")
+        
+        # Initialize OpenAI embeddings
+        print("Setting up OpenAI embeddings...")
+        self.embeddings = OpenAIEmbeddings(
+            openai_api_key=Settings.OPENAI_API_KEY
+        )
+        
         self.vector_store = None
         self.qa_chain = None
         
@@ -24,12 +32,11 @@ class LyricGenerator:
         if Settings.is_huggingface():
             try:
                 print(f"Loading embeddings from HF dataset: {Settings.HF_DATASET}")
-                # Download dataset to Space's storage (cached after first download)
+                # Download dataset to Space's storage
                 local_dir = Path(snapshot_download(
                     repo_id=Settings.HF_DATASET,
                     repo_type="dataset",
-                    token=Settings.HF_TOKEN,
-                    local_dir="/tmp/chroma_db"  # Specify a fixed location
+                    local_dir="/tmp/chroma_db"
                 ))
                 print(f"Dataset downloaded to: {local_dir}")
                 
@@ -37,18 +44,25 @@ class LyricGenerator:
                 chroma_dir = local_dir / "chroma"
                 if not chroma_dir.exists():
                     raise RuntimeError(
-                        f"Chroma directory not found at {chroma_dir}. "
-                        "Please ensure the dataset contains a 'chroma' directory."
+                        f"Chroma directory not found at {chroma_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 from Space storage")
+                print("Successfully loaded vector store")
+                
+                # 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 HF embeddings: {str(e)}")
                 raise RuntimeError(f"Failed to load HF embeddings: {str(e)}")
         else:
             if not self.embeddings_dir.exists():
@@ -58,13 +72,21 @@ class LyricGenerator:
                 )
             
             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)}")
         
         # Setup QA chain
@@ -206,11 +228,27 @@ class LyricGenerator:
             
             # Get source documents with scores first
             print("Searching for similar documents...")
-            docs_and_scores = self.vector_store.similarity_search_with_score(
-                prompt,
-                k=20
-            )
-            print(f"Found {len(docs_and_scores)} similar documents")
+            try:
+                # Test embeddings function first
+                print("Testing embeddings function...")
+                test_embedding = self.embeddings.embed_query("test")
+                print(f"Embeddings function working (vector size: {len(test_embedding)})")
+                
+                # Now try similarity search
+                docs_and_scores = self.vector_store.similarity_search_with_score(
+                    prompt,
+                    k=20
+                )
+                print(f"Found {len(docs_and_scores)} similar documents")
+                
+                if not docs_and_scores:
+                    print("Warning: No similar documents found. This may affect generation quality.")
+                
+            except Exception as e:
+                print(f"Error during similarity search: {str(e)}")
+                raise RuntimeError(
+                    f"Failed to search vector store: {str(e)}"
+                )
 
             # Sort by similarity (convert distance to similarity)
             docs_and_scores.sort(key=lambda x: x[1], reverse=False)