Update process_hf_dataset.py

process_hf_dataset.py

@@ -3,11 +3,11 @@ from datasets import load_dataset
 import re
 from parser import parse_python_code, create_vector
 from database import init_chromadb, store_program, DB_NAME, HF_DATASET_NAME
-from transformers import AutoTokenizer, AutoModel
-import torch
 import chromadb
 import os
 from dotenv import load_dotenv
+from transformers import AutoTokenizer, AutoModel
+import torch
 
 # Load environment variables
 load_dotenv()
@@ -97,41 +97,13 @@ def generate_description_tokens(sequence, vectors, var_map=None):
     
     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
-
-
-"""Generate a 6D semantic vector for a textual description using CodeBERT, projecting to 6D."""
-# Load CodeBERT model and tokenizer
-use_gpu=False
-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, 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)
@@ -150,6 +122,7 @@ def generate_semantic_vector(description, total_lines=100, use_gpu=False):
         vector = vector[:6]  # Truncate to 6D
     
     return vector
+
 def process_hf_dataset():
     """Process the Hugging Face dataset and store programs in ChromaDB, aligning with vector categories and including instruction in vectors."""
     # Load the dataset
@@ -168,29 +141,30 @@ def process_hf_dataset():
         
         # Parse the code to get parts and sequence, generating our 6D vectors
         parts, sequence = parse_python_code(processed_code)
-        vectors = [part['vector'] for part in parts]  # Use parser's 6D vectors
+        program_vectors = [part['vector'] for part in parts]  # Use parser's 6D vectors for program structure
         
         # Generate description tokens including variable roles
         description_tokens = f"task:{instruction.replace(' ', '_')}"
-        description_tokens_list = generate_description_tokens(sequence, vectors, var_map)
+        description_tokens_list = generate_description_tokens(sequence, program_vectors, var_map)
         description_tokens += " " + " ".join(description_tokens_list)
         
-        # Generate a 6D semantic vector for the instruction, incorporating it into the program vector
+        # Generate a 6D semantic vector for the instruction
         semantic_vector = generate_semantic_vector(instruction)
         
-        # Combine program vectors with instruction vector (average or concatenate, but ensure 6D)
-        combined_vector = semantic_vector  # Use semantic vector as primary for semantic search
+        # Combine program vectors with semantic vector (use semantic vector for semantic search, store program vectors separately)
+        # Store both semantic and program vectors, but ensure ChromaDB uses 6D
+        combined_vector = semantic_vector  # Use semantic vector for ChromaDB embedding (6D)
         
         # Store in ChromaDB with description and combined vector
         store_program(client, processed_code, sequence, [combined_vector], DB_NAME)
         
-        # Update metadata with instruction and variable roles as description
+        # Update metadata with instruction and variable roles as description, and store program vectors
         collection = client.get_collection(DB_NAME)
         program_id = str(hash(processed_code))
         collection.update(
             ids=[program_id],
-            metadatas=[{"sequence": ",".join(sequence), "description_tokens": description_tokens}],
-            embeddings=[combined_vector]  # Ensure 6D embedding
+            metadatas=[{"sequence": ",".join(sequence), "description_tokens": description_tokens, "program_vectors": str(program_vectors)}],
+            embeddings=[combined_vector]  # Ensure 6D embedding for semantic search
         )
     
     # Save to Hugging Face Dataset
@@ -206,8 +180,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