Update database.py

database.py

@@ -5,6 +5,8 @@ import os
 from sklearn.metrics.pairwise import cosine_similarity
 import numpy as np
 from datasets import Dataset, load_dataset
+from transformers import AutoTokenizer, AutoModel
+import torch
 
 # User-configurable variables
 DB_NAME = "python_programs"  # ChromaDB collection name
@@ -37,7 +39,7 @@ 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
+    # Flatten vectors to ensure they are a list of numbers (ChromaDB expects flat embeddings)
     flattened_vectors = [item for sublist in vectors for item in sublist]
     
     # Store program data (ID, code, sequence, vectors)
@@ -77,10 +79,10 @@ def query_programs(client, operations, collection_name=DB_NAME, top_k=5, semanti
     collection = create_collection(client, collection_name)
     
     if semantic_query:
-        # Semantic search using description tokens
+        # Semantic search using CodeBERT embeddings
         query_vector = generate_semantic_vector(semantic_query)
         results = collection.query(
-            query_texts=[semantic_query],
+            query_embeddings=[query_vector],
             n_results=top_k,
             include=["documents", "metadatas"]
         )
@@ -99,8 +101,12 @@ def query_programs(client, operations, collection_name=DB_NAME, top_k=5, semanti
         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(doc['vectors']) if isinstance(doc['vectors'], str) else doc['vectors']
-                program_vector = np.mean([v for v in doc_vectors if isinstance(v, (list, np.ndarray))], axis=0).tolist()
+                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
             similarity = cosine_similarity([query_vector], [program_vector])[0][0] if program_vector and query_vector else 0
@@ -155,23 +161,29 @@ def generate_description_tokens(sequence, vectors):
             tokens.append(f"span:{vec[3]:.2f}")
     return tokens
 
-def generate_semantic_vector(description):
-    """Generate a semantic vector for a textual description (simplified for now)."""
-    # This is a placeholder—use an embedding model (e.g., CodeBERT, BERT) for real semantic search
-    tokens = description.lower().split()
-    category_weights = {
-        'import': 1, 'function': 2, 'assign': 17, 'input': 18, 'return': 19, 'if': 5, 'try': 8, 'except': 14
-    }
-    vector = [0] * 6
-    for token in tokens:
-        for cat, weight in category_weights.items():
-            if cat in token:
-                vector[0] = weight  # Use category_id as primary feature
-                vector[1] = 1  # Assume level 1 for simplicity
-                vector[2] = 0.5  # Center position (midpoint)
-                vector[3] = 0.1  # Span (small for simplicity)
-                vector[4] = 1  # Parent depth (shallow)
-                vector[5] = weight / len(category_weights)  # Parent weight
+def generate_semantic_vector(description, use_gpu=False):
+    """Generate a semantic vector for a textual description using CodeBERT, with CPU/GPU option."""
+    # 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)
+    
+    # Tokenize and encode the description
+    inputs = tokenizer(description, return_tensors="pt", padding=True, truncation=True, max_length=512)
+    inputs = {k: v.to(device) for k, v in inputs.items()}
+    
+    # Generate embeddings
+    with torch.no_grad():
+        outputs = model(**inputs)
+        # Use mean pooling of the last hidden states
+        vector = outputs.last_hidden_state.mean(dim=1).squeeze().cpu().numpy().tolist()
+    
+    # Truncate or pad to 6D to match our vectors
+    if len(vector) < 6:
+        vector.extend([0] * (6 - len(vector)))
+    elif len(vector) > 6:
+        vector = vector[:6]
     return vector
 
 def save_chromadb_to_hf(dataset_name=HF_DATASET_NAME, token=HF_TOKEN):
@@ -184,7 +196,7 @@ def save_chromadb_to_hf(dataset_name=HF_DATASET_NAME, token=HF_TOKEN):
     data = {
         "code": results["documents"],
         "sequence": [meta["sequence"] for meta in results["metadatas"]],
-        "vectors": [[item for sublist in vec for item in sublist] for vec in results["embeddings"]],  # Flatten vectors
+        "vectors": results["embeddings"],  # ChromaDB already flattens embeddings
         "description_tokens": [meta.get('description_tokens', '') for meta in results["metadatas"]]
     }