Update database.py

database.py

@@ -1,6 +1,6 @@
 # database.py
 import chromadb
-from parser import parse_python_code
+from parser import parse_python_code, create_vector
 import os
 from sklearn.metrics.pairwise import cosine_similarity
 import numpy as np
@@ -83,7 +83,7 @@ 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 CodeBERT embeddings
+        # Semantic search using a 6D vector generated from the description
         query_vector = generate_semantic_vector(semantic_query)
         results = collection.query(
             query_embeddings=[query_vector],
@@ -165,29 +165,29 @@ def generate_description_tokens(sequence, vectors):
             tokens.append(f"span:{vec[3]:.2f}")
     return tokens
 
-def generate_semantic_vector(description, use_gpu=USE_GPU):
-    """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)
+def generate_semantic_vector(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
+    }
     
-    # 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()}
+    # Parse description for key terms
+    tokens = description.lower().split()
+    vector = [0] * 6  # Initialize 6D vector
     
-    # 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()
+    # 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
     
-    # 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=os.getenv("HF_KEY")):