Update database.py

database.py

@@ -44,15 +44,16 @@ def store_program(client, code, sequence, vectors, collection_name=DB_NAME):
     collection = create_collection(client, collection_name)
     
     # 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]
+    # 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))}],
+        metadatas=[{"sequence": ",".join(sequence), "description_tokens": " ".join(generate_description_tokens(sequence, vectors)), "program_vectors": str(vectors)}],
         ids=[program_id],
-        embeddings=[flattened_vectors]  # Pass as flat list
+        embeddings=[flattened_vectors]  # Pass as 6D semantic vector
     )
     return program_id
 
@@ -106,15 +107,17 @@ def query_programs(client, operations, collection_name=DB_NAME, top_k=5, semanti
         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']
+                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
-            similarity = cosine_similarity([query_vector], [program_vector])[0][0] if program_vector and query_vector else 0
-            matching_programs.append({'id': meta['id'], 'code': doc, 'similarity': similarity, 'description': meta.get('description_tokens', '')})
+            # 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)
 
@@ -165,32 +168,7 @@ def generate_description_tokens(sequence, vectors):
             tokens.append(f"span:{vec[3]:.2f}")
     return tokens
 
-def generate_semantic_vector_og(description, total_lines=100):
-    """Generate a 6D semantic vector for a textual description, matching our vector format."""
-    # Use a simplified heuristic to map description to our 6D vector format
-    category_map = {
-        'import': 1, 'function': 2, 'assign': 17, 'input': 18, 'return': 19, 'if': 5, 'try': 8, 'except': 14
-    }
-    
-    # Parse description for key terms
-    tokens = description.lower().split()
-    vector = [0] * 6  # Initialize 6D vector
-    
-    # Map description tokens to categories and assign basic vector values
-    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 (assume top-level for simplicity)
-                vector[2] = 0.5  # center_pos (midpoint of code)
-                vector[3] = 0.1  # span (small for simplicity)
-                vector[4] = 1  # parent_depth (shallow)
-                vector[5] = cat_id / len(category_map)  # parent_weight (normalized)
-                break
-    
-    return vector
-
-def generate_semantic_vector(description, total_lines=100, use_gpu=False):
+def generate_semantic_vector(description, total_lines=100, use_gpu=USE_GPU):
     """Generate a 6D semantic vector for a textual description using CodeBERT, projecting to 6D."""
     # Load CodeBERT model and tokenizer
     model_name = "microsoft/codebert-base"
@@ -226,8 +204,9 @@ def save_chromadb_to_hf(dataset_name=HF_DATASET_NAME, token=os.getenv("HF_KEY"))
     data = {
         "code": results["documents"],
         "sequence": [meta["sequence"] for meta in results["metadatas"]],
-        "vectors": results["embeddings"],  # ChromaDB already flattens embeddings
-        "description_tokens": [meta.get('description_tokens', '') 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
@@ -254,9 +233,9 @@ def load_chromadb_from_hf(dataset_name=HF_DATASET_NAME, token=os.getenv("HF_KEY"
     for item in dataset:
         collection.add(
             documents=[item["code"]],
-            metadatas=[{"sequence": item["sequence"], "description_tokens": item["description_tokens"]}],
+            metadatas=[{"sequence": item["sequence"], "description_tokens": item["description_tokens"], "program_vectors": str(item["program_vectors"])}],
             ids=[str(hash(item["code"]))],
-            embeddings=[item["vectors"]]
+            embeddings=[item["vectors"]]  # Use semantic 6D vectors
         )
     return client