Update process_hf_dataset.py

process_hf_dataset.py

@@ -184,8 +184,12 @@ def process_hf_dataset(batch_size=100, use_gpu=False):
     client = init_chromadb()
     
     # Do not clear or populate with defaults here—let UI buttons handle this
-    collection = client.get_or_create_collection(DB_NAME)
-    logger.info(f"Using existing or new ChromaDB collection: {DB_NAME}")
+    try:
+        collection = client.get_or_create_collection(DB_NAME)
+        logger.info(f"Using existing or new ChromaDB collection: {DB_NAME}")
+    except Exception as e:
+        logger.error(f"Error accessing ChromaDB collection: {e}")
+        raise
     
     # Process in batches with progress bar
     total_entries = len(dataset_list)
@@ -246,25 +250,25 @@ 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."""
-    client = init_chromadb()
-    collection = client.get_collection(DB_NAME)
-    
-    # Fetch all data from ChromaDB
-    results = collection.get(include=["documents", "metadatas", "embeddings"])
-    data = {
-        "code": results["documents"],
-        "sequence": [meta["sequence"] for meta in results["metadatas"]],
-        "vectors": results["embeddings"],  # Semantic 6D vectors
-        "description_tokens": [meta.get('description_tokens', '') for meta in results["metadatas"]],
-        "program_vectors": [eval(meta.get('program_vectors', '[]')) for meta in results["metadatas"]]  # Store structural vectors
-    }
-    
-    # Create a Hugging Face Dataset
-    dataset = Dataset.from_dict(data)
-    
-    # Push to Hugging Face Hub
+    """Save ChromaDB data to Hugging Face Dataset, with error handling."""
     try:
+        client = init_chromadb()
+        collection = client.get_collection(DB_NAME)
+        
+        # Fetch all data from ChromaDB
+        results = collection.get(include=["documents", "metadatas", "embeddings"])
+        data = {
+            "code": results["documents"],
+            "sequence": [meta["sequence"] for meta in results["metadatas"]],
+            "vectors": results["embeddings"],  # Semantic 6D vectors
+            "description_tokens": [meta.get('description_tokens', '') for meta in results["metadatas"]],
+            "program_vectors": [eval(meta.get('program_vectors', '[]')) for meta in results["metadatas"]]  # Store structural vectors
+        }
+        
+        # Create a Hugging Face Dataset
+        dataset = Dataset.from_dict(data)
+        
+        # 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}")
     except Exception as e: