Update process_hf_dataset.py

process_hf_dataset.py

@@ -10,6 +10,11 @@ from transformers import AutoTokenizer, AutoModel
 import torch
 from tqdm import tqdm  # For progress bar
 import time
+import logging
+
+# Set up logging
+logging.basicConfig(level=logging.INFO, format='%(asctime)s - %(levelname)s - %(message)s')
+logger = logging.getLogger(__name__)
 
 # Load environment variables
 load_dotenv()
@@ -28,9 +33,9 @@ def load_codebert_model(use_gpu=False):
             device = torch.device("cuda" if use_gpu and torch.cuda.is_available() else "cpu")
             tokenizer = AutoTokenizer.from_pretrained(model_name)
             model = AutoModel.from_pretrained(model_name).to(device)
-            print(f"CodeBERT model loaded on {device}")
+            logger.info(f"CodeBERT model loaded on {device}")
         except Exception as e:
-            print(f"Error loading CodeBERT model: {e}")
+            logger.error(f"Error loading CodeBERT model: {e}")
             raise
     return tokenizer, model, device
 
@@ -143,6 +148,7 @@ def generate_semantic_vector(description, total_lines=100, use_gpu=False):
     
     # Ensure vector isn’t all zeros or defaults
     if all(v == 0 for v in vector):
+        logger.warning(f"Default vector detected for description: {description}")
         # Fallback: Use heuristic if CodeBERT fails to generate meaningful embeddings
         category_map = {
             'import': 1, 'function': 2, 'assign': 17, 'input': 18, 'return': 19, 'if': 5, 'try': 8, 'except': 14
@@ -160,15 +166,19 @@ def generate_semantic_vector(description, total_lines=100, use_gpu=False):
                     vector[5] = cat_id / len(category_map)  # parent_weight
                     break
     
+    logger.debug(f"Generated semantic vector for '{description}': {vector}")
     return vector
 
 def process_hf_dataset(batch_size=100, use_gpu=False):
     """Process the Hugging Face dataset in batches and store programs in ChromaDB, aligning with vector categories."""
     # Load the dataset
-    dataset = load_dataset("iamtarun/python_code_instructions_18k_alpaca", split="train")
-    
-    # Convert dataset to list of dictionaries for iteration
-    dataset_list = list(dataset)
+    try:
+        dataset = load_dataset("iamtarun/python_code_instructions_18k_alpaca", split="train")
+        dataset_list = list(dataset)
+        logger.info(f"Loaded dataset with {len(dataset_list)} entries")
+    except Exception as e:
+        logger.error(f"Error loading dataset: {e}")
+        raise
     
     # Initialize ChromaDB client
     client = init_chromadb()
@@ -187,33 +197,36 @@ def process_hf_dataset(batch_size=100, use_gpu=False):
         batch_embeddings = []
         
         for entry in batch:
-            instruction = entry['instruction']
-            output = entry['output']
-            
-            # Rename variables to align with vector categories
-            processed_code, var_map = rename_variables(output)
-            
-            # Parse the code to get parts and sequence, generating our 6D vectors
-            parts, sequence = parse_python_code(processed_code)
-            program_vectors = [part['vector'] for part in parts]  # Use parser's 6D vectors for program structure
-            
-            # Generate description tokens including variable roles
-            description_tokens = f"task:{instruction.replace(' ', '_')}"
-            description_tokens_list = generate_description_tokens(sequence, program_vectors, var_map)
-            description_tokens += " " + " ".join(description_tokens_list)
-            
-            # Generate a 6D semantic vector for the instruction
-            semantic_vector = generate_semantic_vector(instruction, use_gpu=use_gpu)
-            
-            # Store program data
-            program_id = str(hash(processed_code))
-            batch_ids.append(program_id)
-            batch_documents.append(processed_code)
-            batch_metadatas.append({"sequence": ",".join(sequence), "description_tokens": description_tokens, "program_vectors": str(program_vectors)})
-            batch_embeddings.append(semantic_vector)
-            
-            # Add small delay to prevent freezing (optional, adjust as needed)
-            time.sleep(0.01)
+            try:
+                instruction = entry['instruction']
+                output = entry['output']
+                
+                # Rename variables to align with vector categories
+                processed_code, var_map = rename_variables(output)
+                
+                # Parse the code to get parts and sequence, generating our 6D vectors
+                parts, sequence = parse_python_code(processed_code)
+                program_vectors = [part['vector'] for part in parts]  # Use parser's 6D vectors for program structure
+                
+                # Generate description tokens including variable roles
+                description_tokens = f"task:{instruction.replace(' ', '_')}"
+                description_tokens_list = generate_description_tokens(sequence, program_vectors, var_map)
+                description_tokens += " " + " ".join(description_tokens_list)
+                
+                # Generate a 6D semantic vector for the instruction
+                semantic_vector = generate_semantic_vector(instruction, use_gpu=use_gpu)
+                
+                # Store program data
+                program_id = str(hash(processed_code))
+                batch_ids.append(program_id)
+                batch_documents.append(processed_code)
+                batch_metadatas.append({"sequence": ",".join(sequence), "description_tokens": description_tokens, "program_vectors": str(program_vectors)})
+                batch_embeddings.append(semantic_vector)
+                
+                logger.debug(f"Processed entry: {program_id}, Vector: {semantic_vector}")
+            except Exception as e:
+                logger.error(f"Error processing entry {i}: {e}")
+                continue  # Skip failed entries but continue processing
         
         # Batch add to ChromaDB
         try:
@@ -223,8 +236,9 @@ def process_hf_dataset(batch_size=100, use_gpu=False):
                 ids=batch_ids,
                 embeddings=batch_embeddings
             )
+            logger.info(f"Added batch {i//batch_size + 1} to ChromaDB with {len(batch_ids)} entries")
         except Exception as e:
-            print(f"Error adding batch to ChromaDB: {e}")
+            logger.error(f"Error adding batch to ChromaDB: {e}")
             raise
     
     # Save to Hugging Face Dataset
@@ -249,8 +263,12 @@ def save_chromadb_to_hf(dataset_name=HF_DATASET_NAME, token=os.getenv("HF_KEY"))
     dataset = Dataset.from_dict(data)
     
     # Push to Hugging Face Hub
-    dataset.push_to_hub(dataset_name, token=token)
-    print(f"Dataset pushed to Hugging Face Hub as {dataset_name}")
+    try:
+        dataset.push_to_hub(dataset_name, token=token)
+        logger.info(f"Dataset pushed to Hugging Face Hub as {dataset_name}")
+    except Exception as e:
+        logger.error(f"Error pushing dataset to Hugging Face Hub: {e}")
+        raise
 
 if __name__ == "__main__":
     process_hf_dataset(batch_size=100, use_gpu=False)