Update process_hf_dataset.py

process_hf_dataset.py

@@ -141,6 +141,25 @@ def generate_semantic_vector(description, total_lines=100, use_gpu=False):
     elif len(vector) > 6:
         vector = vector[:6]  # Truncate to 6D
     
+    # Ensure vector isn’t all zeros or defaults
+    if all(v == 0 for v in vector):
+        # Fallback: Use heuristic if CodeBERT fails to generate meaningful embeddings
+        category_map = {
+            'import': 1, 'function': 2, 'assign': 17, 'input': 18, 'return': 19, 'if': 5, 'try': 8, 'except': 14
+        }
+        tokens = description.lower().split()
+        vector = [0] * 6
+        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
+                    vector[2] = 0.5  # center_pos
+                    vector[3] = 0.1  # span
+                    vector[4] = 1  # parent_depth
+                    vector[5] = cat_id / len(category_map)  # parent_weight
+                    break
+    
     return vector
 
 def process_hf_dataset(batch_size=100, use_gpu=False):
@@ -154,12 +173,8 @@ def process_hf_dataset(batch_size=100, use_gpu=False):
     # Initialize ChromaDB client
     client = init_chromadb()
     
-    # Clear existing collection (fresh install) if needed
-    try:
-        client.delete_collection(DB_NAME)
-    except:
-        pass  # Collection may not exist
-    collection = client.create_collection(DB_NAME)
+    # Do not clear or populate with defaults here—let UI buttons handle this
+    collection = client.get_or_create_collection(DB_NAME)
     
     # Process in batches with progress bar
     total_entries = len(dataset_list)