Update app.py

app.py

@@ -5,8 +5,9 @@ import os
 import json
 import io
 import subprocess  # To call process_hf_dataset.py
-from database import init_chromadb, store_program, query_programs, load_chromadb_from_hf, DB_NAME, create_collection
+from database import init_chromadb, store_program, query_programs, load_chromadb_from_hf, DB_NAME, create_collection, save_chromadb_to_hf
 import logging
+from datasets import Dataset
 
 # Set up logging
 logging.basicConfig(level=logging.INFO, format='%(asctime)s - %(levelname)s - %(message)s')
@@ -127,6 +128,25 @@ def index():
                 if collection is None or not hasattr(collection, 'add'):
                     raise ValueError("ChromaDB collection creation failed")
                 logger.info("Verified ChromaDB collection is valid")
+                # Verify collection is empty
+                count = collection.count()
+                logger.info(f"ChromaDB now contains {count} entries after reset (should be 0)")
+                
+                # Reset Hugging Face dataset (replace with empty dataset)
+                try:
+                    empty_data = {
+                        "code": [],
+                        "sequence": [],
+                        "vectors": [],
+                        "description_tokens": [],
+                        "program_vectors": []
+                    }
+                    empty_dataset = Dataset.from_dict(empty_data)
+                    empty_dataset.push_to_hub(HF_DATASET_NAME, token=os.getenv("HF_KEY"))
+                    logger.info(f"Replaced Hugging Face dataset {HF_DATASET_NAME} with empty dataset")
+                except Exception as e:
+                    logger.error(f"Error replacing Hugging Face dataset: {e}")
+                    raise
                 
                 # Process dataset
                 result = subprocess.run(['python', 'process_hf_dataset.py'], check=True, capture_output=True, text=True, cwd=os.path.dirname(__file__))
@@ -144,7 +164,7 @@ def index():
                     reconstructed_code=None,
                     code_input=None,
                     query_results=None,
-                    message="Hugging Face dataset processed and stored successfully with fresh database."
+                    message="Hugging Face dataset processed and stored successfully with fresh database and empty dataset."
                 )
             except subprocess.CalledProcessError as e:
                 logger.error(f"Error processing Hugging Face dataset: {e.stderr}")
@@ -163,6 +183,9 @@ def index():
                 if collection is None or not hasattr(collection, 'add'):
                     raise ValueError("ChromaDB collection access failed")
                 logger.info("Verified ChromaDB collection is valid")
+                # Verify collection state
+                count = collection.count()
+                logger.info(f"ChromaDB contains {count} entries before loading")
                 
                 # Process dataset
                 result = subprocess.run(['python', 'process_hf_dataset.py'], check=True, capture_output=True, text=True, cwd=os.path.dirname(__file__))
@@ -189,7 +212,7 @@ def index():
                 logger.error(f"Unexpected error loading Hugging Face dataset: {e}")
                 return f"Unexpected error loading Hugging Face dataset: {e}", 500
         elif 'reset_db' in request.form:
-            # Reset ChromaDB collection (no repopulation with samples)
+            # Reset ChromaDB collection and Hugging Face dataset (no repopulation with samples)
             try:
                 client = init_chromadb()
                 try:
@@ -206,6 +229,23 @@ def index():
                 # Verify collection is empty
                 count = collection.count()
                 logger.info(f"ChromaDB now contains {count} entries after reset (should be 0)")
+                
+                # Reset Hugging Face dataset (replace with empty dataset)
+                try:
+                    empty_data = {
+                        "code": [],
+                        "sequence": [],
+                        "vectors": [],
+                        "description_tokens": [],
+                        "program_vectors": []
+                    }
+                    empty_dataset = Dataset.from_dict(empty_data)
+                    empty_dataset.push_to_hub(HF_DATASET_NAME, token=os.getenv("HF_KEY"))
+                    logger.info(f"Replaced Hugging Face dataset {HF_DATASET_NAME} with empty dataset")
+                except Exception as e:
+                    logger.error(f"Error replacing Hugging Face dataset: {e}")
+                    raise
+                
                 return render_template(
                     'results_partial.html',
                     parts=None,
@@ -213,7 +253,7 @@ def index():
                     reconstructed_code=None,
                     code_input=None,
                     query_results=None,
-                    message="Database reset successfully."
+                    message="Database and Hugging Face dataset reset successfully."
                 )
             except Exception as e:
                 logger.error(f"Error resetting database: {e}")