Update process_hf_dataset.py

process_hf_dataset.py

@@ -2,7 +2,7 @@
 from datasets import load_dataset
 import re
 from parser import parse_python_code, create_vector
-from database import init_chromadb, store_program, DB_NAME, HF_DATASET_NAME
+from database import init_chromadb, store_program, DB_NAME, HF_DATASET_NAME, create_collection
 import chromadb
 import os
 from dotenv import load_dotenv
@@ -187,6 +187,10 @@ def process_hf_dataset(batch_size=100, use_gpu=False):
     try:
         collection = client.get_or_create_collection(DB_NAME)
         logger.info(f"Using existing or new ChromaDB collection: {DB_NAME}")
+        # Verify collection is valid
+        if collection is None or not hasattr(collection, 'add'):
+            raise ValueError("ChromaDB collection access failed")
+        logger.info(f"ChromaDB collection verified, contains {collection.count()} entries")
     except Exception as e:
         logger.error(f"Error accessing ChromaDB collection: {e}")
         raise
@@ -242,6 +246,9 @@ def process_hf_dataset(batch_size=100, use_gpu=False):
                 embeddings=batch_embeddings
             )
             logger.info(f"Added batch {i//batch_size + 1} to ChromaDB with {len(batch_ids)} entries")
+            # Verify addition
+            count = collection.count()
+            logger.info(f"ChromaDB now contains {count} entries after adding batch")
         except Exception as e:
             logger.error(f"Error adding batch to ChromaDB: {e}")
             raise
@@ -250,7 +257,7 @@ def process_hf_dataset(batch_size=100, use_gpu=False):
     save_chromadb_to_hf()
 
 def save_chromadb_to_hf(dataset_name=HF_DATASET_NAME, token=os.getenv("HF_KEY")):
-    """Save ChromaDB data to Hugging Face Dataset, with error handling."""
+    """Save ChromaDB data to Hugging Face Dataset, with error handling and logging."""
     try:
         client = init_chromadb()
         collection = client.get_collection(DB_NAME)
@@ -267,10 +274,13 @@ def save_chromadb_to_hf(dataset_name=HF_DATASET_NAME, token=os.getenv("HF_KEY"))
         
         # Create a Hugging Face Dataset
         dataset = Dataset.from_dict(data)
+        logger.info(f"Created Hugging Face Dataset with {len(data['code'])} entries")
         
         # Push to Hugging Face Hub
         dataset.push_to_hub(dataset_name, token=token)
         logger.info(f"Dataset pushed to Hugging Face Hub as {dataset_name}")
+        # Verify push (optional, could check dataset on Hub)
+        logger.info(f"Verified Hugging Face dataset push with {len(dataset)} entries")
     except Exception as e:
         logger.error(f"Error pushing dataset to Hugging Face Hub: {e}")
         raise