Update database.py

database.py

@@ -8,118 +8,142 @@ from datasets import Dataset, load_dataset
 from transformers import AutoTokenizer, AutoModel
 import torch
 from dotenv import load_dotenv
+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()
 
-# User-configurable variables (no HF_KEY hardcoded here)
+# User-configurable variables
 DB_NAME = "python_programs"  # ChromaDB collection name
 HF_DATASET_NAME = "python_program_vectors"  # Hugging Face Dataset name
 PERSIST_DIR = "./chroma_data"  # Directory for persistent storage (optional)
 USE_GPU = False  # Default to CPU, set to True for GPU if available
 
 def init_chromadb(persist_dir=PERSIST_DIR):
-    """Initialize ChromaDB client, optionally with persistent storage."""
+    """Initialize ChromaDB client, optionally with persistent storage, with error handling."""
     try:
         # Use persistent storage if directory exists, otherwise in-memory
         if os.path.exists(persist_dir):
+            logger.info(f"Initializing ChromaDB with persistent storage at {persist_dir}")
             client = chromadb.PersistentClient(path=persist_dir)
         else:
+            logger.info("Initializing ChromaDB with in-memory storage")
             client = chromadb.Client()
         return client
     except Exception as e:
-        print(f"Error initializing ChromaDB: {e}")
-        return chromadb.Client()  # Fallback to in-memory
+        logger.error(f"Error initializing ChromaDB: {e}")
+        raise
 
 def create_collection(client, collection_name=DB_NAME):
-    """Create or get a ChromaDB collection for Python programs."""
+    """Create or get a ChromaDB collection for Python programs, with error handling."""
     try:
-        collection = client.get_collection(name=collection_name)
-    except:
-        collection = client.create_collection(name=collection_name)
-    return collection
+        collection = client.get_or_create_collection(name=collection_name)
+        logger.info(f"Using ChromaDB collection: {collection_name}")
+        return collection
+    except Exception as e:
+        logger.error(f"Error creating or getting collection {collection_name}: {e}")
+        raise
 
 def store_program(client, code, sequence, vectors, collection_name=DB_NAME):
-    """Store a program in ChromaDB with its code, sequence, and vectors."""
-    collection = create_collection(client, collection_name)
-    
-    # Flatten vectors to ensure they are a list of numbers (ChromaDB expects flat embeddings)
-    # Use the first vector (semantic vector) for ChromaDB embedding
-    flattened_vectors = vectors[0] if vectors else [0] * 6  # Ensure 6D
-    
-    # Store program data (ID, code, sequence, vectors)
-    program_id = str(hash(code))  # Use hash of code as ID for uniqueness
-    collection.add(
-        documents=[code],
-        metadatas=[{"sequence": ",".join(sequence), "description_tokens": " ".join(generate_description_tokens(sequence, vectors)), "program_vectors": str(vectors)}],
-        ids=[program_id],
-        embeddings=[flattened_vectors]  # Pass as 6D semantic vector
-    )
-    return program_id
+    """Store a program in ChromaDB with its code, sequence, and vectors, with error handling."""
+    try:
+        collection = create_collection(client, collection_name)
+        
+        # Flatten vectors to ensure they are a list of numbers (ChromaDB expects flat embeddings)
+        # Use the first vector (semantic or program vector) for ChromaDB embedding, ensuring 6D
+        flattened_vectors = vectors[0] if vectors and len(vectors) > 0 and len(vectors[0]) == 6 else [0] * 6
+        
+        # Store program data (ID, code, sequence, vectors)
+        program_id = str(hash(code))  # Use hash of code as ID for uniqueness
+        collection.add(
+            documents=[code],
+            metadatas=[{"sequence": ",".join(sequence), "description_tokens": " ".join(generate_description_tokens(sequence, vectors)), "program_vectors": str(vectors)}],
+            ids=[program_id],
+            embeddings=[flattened_vectors]  # Pass as 6D vector
+        )
+        logger.info(f"Stored program in ChromaDB: {program_id}")
+        return program_id
+    except Exception as e:
+        logger.error(f"Error storing program in ChromaDB: {e}")
+        raise
 
 def populate_sample_db(client):
-    """Populate ChromaDB with sample Python programs."""
-    samples = [
-        """
-        import os
-        def add_one(x):
-            y = x + 1
-            return y
-        """,
-        """
-        def multiply(a, b):
-            c = a * b
-            if c > 0:
-                return c
-        """
-    ]
-    
-    for code in samples:
-        parts, sequence = parse_python_code(code)
-        vectors = [part['vector'] for part in parts]
-        store_program(client, code, sequence, vectors)
+    """Populate ChromaDB with sample Python programs, with logging."""
+    try:
+        samples = [
+            """
+            import os
+            def add_one(x):
+                y = x + 1
+                return y
+            """,
+            """
+            def multiply(a, b):
+                c = a * b
+                if c > 0:
+                    return c
+            """
+        ]
+        
+        for code in samples:
+            parts, sequence = parse_python_code(code)
+            vectors = [part['vector'] for part in parts]
+            store_program(client, code, sequence, vectors)
+        logger.info("Populated ChromaDB with sample programs")
+    except Exception as e:
+        logger.error(f"Error populating sample database: {e}")
+        raise
 
 def query_programs(client, operations, collection_name=DB_NAME, top_k=5, semantic_query=None):
-    """Query ChromaDB for programs matching the operations sequence or semantic description."""
-    collection = create_collection(client, collection_name)
-    
-    if semantic_query:
-        # Semantic search using a 6D vector generated from the description
-        query_vector = generate_semantic_vector(semantic_query)
-        results = collection.query(
-            query_embeddings=[query_vector],
-            n_results=top_k,
-            include=["documents", "metadatas"]
-        )
-    else:
-        # Vector-based search for operations sequence
-        query_vector = sum([create_vector(op, 0, (1, 1), 100, []) for op in operations], []) / len(operations) if operations else [0] * 6
-        results = collection.query(
-            query_embeddings=[query_vector],
-            n_results=top_k,
-            include=["documents", "metadatas"]
-        )
-    
-    # Process results
-    matching_programs = []
-    for doc, meta in zip(results['documents'][0], results['metadatas'][0]):
-        sequence = meta['sequence'].split(',')
-        if not semantic_query or is_subsequence(operations, sequence):  # Ensure sequence match for operations
-            try:
-                # Reconstruct program vectors (flatten if needed)
-                doc_vectors = eval(meta['program_vectors']) if isinstance(meta['program_vectors'], str) else meta['program_vectors']
-                if isinstance(doc_vectors, (list, np.ndarray)) and len(doc_vectors) == 6:
-                    program_vector = doc_vectors  # Single flat vector
-                else:
-                    program_vector = np.mean([v for v in doc_vectors if isinstance(v, (list, np.ndarray))], axis=0).tolist()
-            except:
-                program_vector = [0] * 6  # Fallback for malformed vectors
-            # Use the semantic embedding for similarity
-            semantic_vector = eval(doc['vectors']) if isinstance(doc['vectors'], str) else doc['vectors']
-            similarity = cosine_similarity([query_vector], [semantic_vector])[0][0] if semantic_vector and query_vector else 0
-            matching_programs.append({'id': meta['id'], 'code': doc, 'similarity': similarity, 'description': meta.get('description_tokens', ''), 'program_vectors': meta.get('program_vectors', '[]')})
-    
-    return sorted(matching_programs, key=lambda x: x['similarity'], reverse=True)
+    """Query ChromaDB for programs matching the operations sequence or semantic description, with error handling."""
+    try:
+        collection = create_collection(client, collection_name)
+        
+        if semantic_query:
+            # Semantic search using a 6D vector generated from the description
+            query_vector = generate_semantic_vector(semantic_query)
+            results = collection.query(
+                query_embeddings=[query_vector],
+                n_results=top_k,
+                include=["documents", "metadatas"]
+            )
+        else:
+            # Vector-based search for operations sequence
+            query_vector = sum([create_vector(op, 0, (1, 1), 100, []) for op in operations], []) / len(operations) if operations else [0] * 6
+            results = collection.query(
+                query_embeddings=[query_vector],
+                n_results=top_k,
+                include=["documents", "metadatas"]
+            )
+        
+        # Process results
+        matching_programs = []
+        for doc, meta in zip(results['documents'][0], results['metadatas'][0]):
+            sequence = meta['sequence'].split(',')
+            if not semantic_query or is_subsequence(operations, sequence):  # Ensure sequence match for operations
+                try:
+                    # Reconstruct program vectors (flatten if needed)
+                    doc_vectors = eval(meta['program_vectors']) if isinstance(meta['program_vectors'], str) else meta['program_vectors']
+                    if isinstance(doc_vectors, (list, np.ndarray)) and len(doc_vectors) == 6:
+                        program_vector = doc_vectors  # Single flat vector
+                    else:
+                        program_vector = np.mean([v for v in doc_vectors if isinstance(v, (list, np.ndarray))], axis=0).tolist()
+                except:
+                    program_vector = [0] * 6  # Fallback for malformed vectors
+                # Use the semantic embedding for similarity
+                semantic_vector = eval(doc['vectors']) if isinstance(doc['vectors'], str) else doc['vectors']
+                similarity = cosine_similarity([query_vector], [semantic_vector])[0][0] if semantic_vector and query_vector else 0
+                matching_programs.append({'id': meta['id'], 'code': doc, 'similarity': similarity, 'description': meta.get('description_tokens', ''), 'program_vectors': meta.get('program_vectors', '[]')})
+        
+        logger.info(f"Queried {len(matching_programs)} programs from ChromaDB")
+        return sorted(matching_programs, key=lambda x: x['similarity'], reverse=True)
+    except Exception as e:
+        logger.error(f"Error querying programs from ChromaDB: {e}")
+        raise
 
 def create_vector(category, level, location, total_lines, parent_path):
     """Helper to create a vector for query (matches parser's create_vector)."""
@@ -168,13 +192,11 @@ def generate_description_tokens(sequence, vectors):
             tokens.append(f"span:{vec[3]:.2f}")
     return tokens
 
-def generate_semantic_vector(description, total_lines=100, use_gpu=USE_GPU):
+def generate_semantic_vector(description, total_lines=100, use_gpu=False):
     """Generate a 6D semantic vector for a textual description using CodeBERT, projecting to 6D."""
-    # Load CodeBERT model and tokenizer
-    model_name = "microsoft/codebert-base"
-    tokenizer = AutoTokenizer.from_pretrained(model_name)
-    device = torch.device("cuda" if use_gpu and torch.cuda.is_available() else "cpu")
-    model = AutoModel.from_pretrained(model_name).to(device)
+    global tokenizer, model, device
+    if tokenizer is None or model is None:
+        tokenizer, model, device = load_codebert_model(use_gpu)
     
     # Tokenize and encode the description
     inputs = tokenizer(description, return_tensors="pt", padding=True, truncation=True, max_length=512)
@@ -192,54 +214,73 @@ def generate_semantic_vector(description, total_lines=100, use_gpu=USE_GPU):
     elif len(vector) > 6:
         vector = vector[:6]  # Truncate to 6D
     
+    # 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
+        }
+        tokens = description.lower().split()
+        vector = [0] * 6
+        for token in tokens:
+            for cat, cat_id in category_map.items():
+                if cat in token:
+                    vector[0] = cat_id  # category_id
+                    vector[1] = 1  # level
+                    vector[2] = 0.5  # center_pos
+                    vector[3] = 0.1  # span
+                    vector[4] = 1  # parent_depth
+                    vector[5] = cat_id / len(category_map)  # parent_weight
+                    break
+    
+    logger.debug(f"Generated semantic vector for '{description}': {vector}")
     return vector
 
 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 = create_collection(client)
-    
-    # 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)
-    print(f"Dataset pushed to Hugging Face Hub as {dataset_name}")
+    """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:
+        logger.error(f"Error pushing dataset to Hugging Face Hub: {e}")
+        raise
 
 def load_chromadb_from_hf(dataset_name=HF_DATASET_NAME, token=os.getenv("HF_KEY")):
-    """Load ChromaDB data from Hugging Face Dataset, handle empty dataset."""
+    """Load ChromaDB data from Hugging Face Dataset, handle empty dataset, with error handling."""
     try:
         dataset = load_dataset(dataset_name, split="train", token=token)
+        client = init_chromadb()
+        collection = create_collection(client)
+        
+        for item in dataset:
+            store_program(client, item["code"], item["sequence"].split(','), item["program_vectors"])
+        logger.info(f"Loaded {len(dataset)} entries from Hugging Face Hub into ChromaDB")
+        return client
     except Exception as e:
-        print(f"Error loading dataset from Hugging Face: {e}. Populating with samples...")
+        logger.error(f"Error loading dataset from Hugging Face: {e}")
+        # Fallback: Create empty collection
         client = init_chromadb()
-        populate_sample_db(client)
-        save_chromadb_to_hf()  # Create and push a new dataset
-        return init_chromadb()
-    
-    client = init_chromadb()
-    collection = create_collection(client)
-    
-    for item in dataset:
-        collection.add(
-            documents=[item["code"]],
-            metadatas=[{"sequence": item["sequence"], "description_tokens": item["description_tokens"], "program_vectors": str(item["program_vectors"])}],
-            ids=[str(hash(item["code"]))],
-            embeddings=[item["vectors"]]  # Use semantic 6D vectors
-        )
-    return client
+        create_collection(client)
+        return client
 
 if __name__ == '__main__':
     client = load_chromadb_from_hf()
-    # Uncomment to save to Hugging Face
-    # save_chromadb_to_hf()
+    logger.info("Database initialized or loaded from Hugging Face Hub")