Update database.py

database.py

@@ -165,7 +165,7 @@ def generate_description_tokens(sequence, vectors):
             tokens.append(f"span:{vec[3]:.2f}")
     return tokens
 
-def generate_semantic_vector(description, total_lines=100):
+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 = {
@@ -190,6 +190,32 @@ def generate_semantic_vector(description, total_lines=100):
     
     return vector
 
+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)
+    
+    # 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 project to 6D (simplified projection: take first 6 dimensions)
+    if len(vector) < 6:
+        vector.extend([0] * (6 - len(vector)))
+    elif len(vector) > 6:
+        vector = vector[:6]  # Truncate to 6D
+    
+    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()