Update database.py

database.py

@@ -7,14 +7,14 @@ import numpy as np
 from datasets import Dataset, load_dataset
 from transformers import AutoTokenizer, AutoModel
 import torch
-import dotenv
-dotenv.load()
+from dotenv import load_dotenv
 
+# Load environment variables
+load_dotenv()
 
-# User-configurable variables
+# User-configurable variables (no HF_KEY hardcoded here)
 DB_NAME = "python_programs"  # ChromaDB collection name
 HF_DATASET_NAME = "python_program_vectors"  # Hugging Face Dataset name
-HF_KEY = os.getenv("HF_KEY")  # Replace with your Hugging Face API token
 PERSIST_DIR = "./chroma_data"  # Directory for persistent storage (optional)
 USE_GPU = False  # Default to CPU, set to True for GPU if available
 
@@ -190,7 +190,7 @@ def generate_semantic_vector(description, use_gpu=USE_GPU):
         vector = vector[:6]
     return vector
 
-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 = create_collection(client)
@@ -211,7 +211,7 @@ def save_chromadb_to_hf(dataset_name=HF_DATASET_NAME, token=HF_KEY):
     dataset.push_to_hub(dataset_name, token=token)
     print(f"Dataset pushed to Hugging Face Hub as {dataset_name}")
 
-def load_chromadb_from_hf(dataset_name=HF_DATASET_NAME, token=HF_KEY):
+def load_chromadb_from_hf(dataset_name=HF_DATASET_NAME, token=os.getenv("HF_KEY")):
     """Load ChromaDB data from Hugging Face Dataset, handle empty dataset."""
     try:
         dataset = load_dataset(dataset_name, split="train", token=token)