Update process_hf_dataset.py

process_hf_dataset.py

@@ -2,8 +2,13 @@
 from datasets import load_dataset
 import re
 from parser import parse_python_code
-from database import init_chromadb, store_program, DB_NAME, HF_DATASET_NAME, HF_KEY
+from database import init_chromadb, store_program, DB_NAME, HF_DATASET_NAME
 import chromadb
+import os
+from dotenv import load_dotenv
+
+# Load environment variables
+load_dotenv()
 
 def rename_variables(code, variable_prefixes=None):
     """Rename variables in Python code to align with vector categories (input_variable, assigned_variable, returned_variable)."""
@@ -29,17 +34,23 @@ def rename_variables(code, variable_prefixes=None):
     # Sort variables by first appearance (simplified, could improve with AST)
     sorted_vars = sorted(list(variables))
     var_map = {}
-    var_count = {prefix: 1 for prefix in variable_prefixes.values()}
+    var_count = {'input_variable': 1, 'assigned_variable': 1, 'returned_variable': 1}
     
     # Assign variables based on context (simplified heuristic)
     for var in sorted_vars:
-        # Determine variable role based on context (simplified)
-        if var in ['expression', 'input']:  # Assume input parameters or initial variables
+        # Determine variable role based on context
+        is_input = any(var in line and 'def' in line for line in code_lines)  # Check if in function definition (input parameter)
+        is_returned = any('return' in line and var in line for line in code_lines)  # Check if used in return statement
+        is_assigned = any('=' in line and var in line.split('=')[0].strip() for line in code_lines)  # Check if assigned
+        
+        if is_input:
             role = 'input_variable'
-        elif var in code.split() and 'return' in line[x] for x in code_lines if var in line[x]:  # Returned variables
+        elif is_returned:
             role = 'returned_variable'
-        else:  # Default to assigned variables
+        elif is_assigned:
             role = 'assigned_variable'
+        else:
+            role = 'assigned_variable'  # Default to assigned if unclear
         
         new_name = f"{role}{var_count[role]}"
         var_map[var] = new_name
@@ -123,7 +134,7 @@ def process_hf_dataset():
     # Save to Hugging Face Dataset
     save_chromadb_to_hf()
 
-def save_chromadb_to_hf(dataset_name=HF_DATASET_NAME, token=HF_KEY):
+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()
     collection = client.get_collection(DB_NAME)