Fixing preferences management to avoid data loss

app.py

@@ -5,6 +5,7 @@ import random
 from datetime import datetime
 from apscheduler.schedulers.background import BackgroundScheduler
 from PIL import Image
+from filelock import FileLock # Added for file locking
 
 import config
 import utils
@@ -13,22 +14,87 @@ import utils
 # Attempt to log in to Hugging Face Hub at startup
 utils.login_hugging_face()
 
-# Load preferences: Try from Hub, then local, then empty
+print(f"Attempting to load preferences from Hugging Face Hub, ensuring local {config.RESULTS_CSV_FILE} is synchronized.")
+# We assume utils.load_preferences_from_hf_hub:
+# 1. Downloads from Hub, overwrites local config.RESULTS_CSV_FILE.
+# 2. If Hub file doesn't exist, local config.RESULTS_CSV_FILE becomes empty (or reflects this).
+# 3. Returns the DataFrame loaded from the (now synchronized) local file.
+# 4. Returns None on major failure (e.g. network, file not found on Hub).
 preferences_df = utils.load_preferences_from_hf_hub(config.HF_DATASET_REPO_ID, config.RESULTS_CSV_FILE)
+
 if preferences_df is None:
+    print(f"Failed to load from Hub or Hub is empty/file not found. Initializing/loading from {config.RESULTS_CSV_FILE} as a fallback.")
     if os.path.exists(config.RESULTS_CSV_FILE):
-        print(f"Loading preferences from local file: {config.RESULTS_CSV_FILE}")
         try:
             preferences_df = pd.read_csv(config.RESULTS_CSV_FILE)
+            if not preferences_df.empty and list(preferences_df.columns) != config.CSV_HEADERS:
+                print(f"Warning: Local CSV {config.RESULTS_CSV_FILE} columns ({list(preferences_df.columns)}) do not match expected headers ({config.CSV_HEADERS}). Re-initializing file and DataFrame.")
+                preferences_df = pd.DataFrame(columns=config.CSV_HEADERS)
+                preferences_df.to_csv(config.RESULTS_CSV_FILE, index=False)
+            elif preferences_df.empty: # Loaded an empty DataFrame
+                # Check if the file itself had incorrect headers or was truly empty
+                current_headers = []
+                if os.path.getsize(config.RESULTS_CSV_FILE) > 0:
+                    try:
+                        current_headers = list(pd.read_csv(config.RESULTS_CSV_FILE, nrows=0).columns)
+                    except Exception: # Handle cases where reading headers might fail
+                        pass # Will be caught by re-initialization if headers are bad
+                if current_headers != config.CSV_HEADERS:
+                    print(f"Local CSV {config.RESULTS_CSV_FILE} is empty or has incorrect headers. Re-initializing file and DataFrame with correct headers.")
+                    preferences_df = pd.DataFrame(columns=config.CSV_HEADERS)
+                    preferences_df.to_csv(config.RESULTS_CSV_FILE, index=False)
+                else: # Empty dataframe, but headers in file are correct
+                    preferences_df = pd.DataFrame(columns=config.CSV_HEADERS) # Ensure in-memory df also has columns
+
         except pd.errors.EmptyDataError:
-            print(f"Local preferences file {config.RESULTS_CSV_FILE} is empty. Starting fresh.")
+            print(f"Local CSV {config.RESULTS_CSV_FILE} is empty. Initializing file and DataFrame with headers.")
             preferences_df = pd.DataFrame(columns=config.CSV_HEADERS)
+            preferences_df.to_csv(config.RESULTS_CSV_FILE, index=False)
         except Exception as e:
-            print(f"Error loading local {config.RESULTS_CSV_FILE}: {e}. Starting fresh.")
+            print(f"Error loading local {config.RESULTS_CSV_FILE}: {e}. Initializing file and DataFrame.")
             preferences_df = pd.DataFrame(columns=config.CSV_HEADERS)
+            preferences_df.to_csv(config.RESULTS_CSV_FILE, index=False)
     else:
-        print("No existing preferences found on Hub or locally. Starting with an empty table.")
+        print(f"Local CSV {config.RESULTS_CSV_FILE} not found. Initializing file and DataFrame.")
         preferences_df = pd.DataFrame(columns=config.CSV_HEADERS)
+        preferences_df.to_csv(config.RESULTS_CSV_FILE, index=False)
+else:
+    # Successfully loaded from Hub; local file config.RESULTS_CSV_FILE should be synchronized.
+    print(f"Successfully loaded preferences from Hugging Face Hub. Local copy at {config.RESULTS_CSV_FILE} should be up-to-date.")
+    if not preferences_df.empty and list(preferences_df.columns) != config.CSV_HEADERS:
+        print(f"CRITICAL: Data from Hub has incorrect columns {list(preferences_df.columns)}. Expected {config.CSV_HEADERS}. Re-initializing local file and DataFrame to empty to prevent corruption.")
+        preferences_df = pd.DataFrame(columns=config.CSV_HEADERS)
+        preferences_df.to_csv(config.RESULTS_CSV_FILE, index=False)
+    elif preferences_df.empty:
+        # Hub data is empty. Ensure DataFrame in memory has correct columns.
+        # And ensure the local CSV (which should have been written by load_preferences_from_hf_hub) has correct headers.
+        if list(preferences_df.columns) != config.CSV_HEADERS:
+             preferences_df = pd.DataFrame(columns=config.CSV_HEADERS) # Correct columns in memory
+        
+        # Ensure the local file has correct headers if it's empty or its headers are wrong
+        needs_header_rewrite = True
+        if os.path.exists(config.RESULTS_CSV_FILE):
+            if os.path.getsize(config.RESULTS_CSV_FILE) == 0: # File is completely empty
+                 needs_header_rewrite = True
+            else:
+                try:
+                    local_headers = list(pd.read_csv(config.RESULTS_CSV_FILE, nrows=0).columns)
+                    if local_headers == config.CSV_HEADERS:
+                        needs_header_rewrite = False
+                except Exception: # Error reading headers, assume rewrite is needed
+                    pass
+        if needs_header_rewrite:
+            print(f"Local file {config.RESULTS_CSV_FILE} (after Hub sync resulted in empty data) is empty or has incorrect headers. Writing/Re-writing headers.")
+            pd.DataFrame(columns=config.CSV_HEADERS).to_csv(config.RESULTS_CSV_FILE, index=False)
+
+
+# Final safety net: ensure preferences_df is a DataFrame with correct columns.
+if not isinstance(preferences_df, pd.DataFrame) or list(preferences_df.columns) != config.CSV_HEADERS:
+    print("Critical: preferences_df is not a valid DataFrame with correct headers after initialization. Resetting to empty with correct headers.")
+    preferences_df = pd.DataFrame(columns=config.CSV_HEADERS)
+    # Ensure the CSV file reflects this state
+    preferences_df.to_csv(config.RESULTS_CSV_FILE, index=False)
+
 
 # Scan for available data
 ALL_SAMPLES_BY_DOMAIN = utils.scan_data_directory(config.DATA_FOLDER)
@@ -39,11 +105,32 @@ if not ALL_SAMPLES_BY_DOMAIN:
 # --- Scheduler for Periodic Uploads ---
 def scheduled_upload_job():
     global preferences_df
-    print(f"Running scheduled job: Saving and uploading preferences at {datetime.now()}")
-    if preferences_df is not None and not preferences_df.empty:
-        utils.save_preferences_to_hf_hub(preferences_df, config.HF_DATASET_REPO_ID, config.RESULTS_CSV_FILE, commit_message="Periodic background update")
-    else:
-        print("Scheduled job: Preferences DataFrame is empty. Nothing to upload.")
+    print(f"Running scheduled job: Preparing to upload preferences from {config.RESULTS_CSV_FILE} at {datetime.now()}")
+    
+    lock_path = config.RESULTS_CSV_FILE + ".lock"
+    with FileLock(lock_path):
+        print(f"Acquired lock for scheduled upload: {lock_path}")
+        if os.path.exists(config.RESULTS_CSV_FILE):
+            try:
+                # Read the current state of the CSV file for upload
+                df_to_upload = pd.read_csv(config.RESULTS_CSV_FILE)
+                if not df_to_upload.empty:
+                    utils.save_preferences_to_hf_hub(
+                        df_to_upload, # df_to_upload is passed for the empty check inside save_preferences_to_hf_hub
+                        config.HF_DATASET_REPO_ID,
+                        config.RESULTS_CSV_FILE, # This is the target filename on the Hub
+                        commit_message="Periodic background update"
+                    )
+                    print(f"Scheduled job: Attempted upload of data from {config.RESULTS_CSV_FILE}.")
+                else:
+                    print(f"Scheduled job: Local preferences file {config.RESULTS_CSV_FILE} is empty. Nothing to upload.")
+            except pd.errors.EmptyDataError:
+                print(f"Scheduled job: Local preferences file {config.RESULTS_CSV_FILE} is empty (read as EmptyDataError). Nothing to upload.")
+            except Exception as e:
+                print(f"Scheduled job: Error reading or uploading {config.RESULTS_CSV_FILE}: {e}")
+        else:
+            print(f"Scheduled job: Local preferences file {config.RESULTS_CSV_FILE} does not exist. Nothing to upload.")
+        print(f"Released lock for scheduled upload: {lock_path}")
 
 scheduler = BackgroundScheduler()
 scheduler.add_job(scheduled_upload_job, 'interval', hours=config.PUSH_INTERVAL_HOURS)
@@ -175,15 +262,40 @@ def process_vote(choice_index, session_id, sample_queue, current_sample_index, c
         preferred_model_key=preferred_model_key
     )
     
-    try:
-        preferences_df.to_csv(config.RESULTS_CSV_FILE, index=False)
-        print(f"Preferences saved locally to {config.RESULTS_CSV_FILE}")
-    except Exception as e:
-        print(f"Error saving preferences locally: {e}")
+    # Append the new preference to the CSV file
+    if not preferences_df.empty:
+        new_preference_df = preferences_df.iloc[-1:] # Get the last row as a new DataFrame
+        
+        lock_path = config.RESULTS_CSV_FILE + ".lock"
+        with FileLock(lock_path):
+            print(f"Acquired lock for vote processing: {lock_path}")
+            try:
+                file_exists_and_has_content = os.path.exists(config.RESULTS_CSV_FILE) and os.path.getsize(config.RESULTS_CSV_FILE) > 0
+                new_preference_df.to_csv(
+                    config.RESULTS_CSV_FILE,
+                    mode='a',
+                    header=not file_exists_and_has_content, # Write header if file is new or empty
+                    index=False
+                )
+                print(f"Appended new preference to {config.RESULTS_CSV_FILE}")
+            except Exception as e:
+                print(f"Error appending preference to local CSV {config.RESULTS_CSV_FILE}: {e}")
+            finally:
+                print(f"Released lock for vote processing: {lock_path}")
+    else:
+        print("Warning: preferences_df is empty after utils.record_preference. Cannot append to CSV.")
+
+    # Removed full CSV overwrite:
+    # try:
+    #     preferences_df.to_csv(config.RESULTS_CSV_FILE, index=False)
+    #     print(f"Preferences saved locally to {config.RESULTS_CSV_FILE}")
+    # except Exception as e:
+    #     print(f"Error saving preferences locally: {e}")
 
     current_sample_index += 1
     if current_sample_index >= len(sample_queue):
-        utils.save_preferences_to_hf_hub(preferences_df, config.HF_DATASET_REPO_ID, config.RESULTS_CSV_FILE, commit_message="Session end update")
+        # Removed session end upload:
+        # utils.save_preferences_to_hf_hub(preferences_df, config.HF_DATASET_REPO_ID, config.RESULTS_CSV_FILE, commit_message="Session end update")
         final_msg = f"# 🎉 Session Complete! 🎉\n\n### All samples have been rated. Thank you for your participation!"
         return preferences_df, current_sample_index, final_msg, None, None, None, [], [], True
     

config.py

@@ -2,7 +2,7 @@
 
 # --- File and Folder Names ---
 DATA_FOLDER = "data"  # Main folder containing domain subfolders
-BACKGROUND_IMAGE_NAME = "input_bg.jpg"  # Standard name for background input
+BACKGROUND_IMAGE_NAME = "input_bg_bb.png"  # Standard name for background input
 FOREGROUND_IMAGE_NAME = "input_fg.jpg"  # Standard name for foreground input
 PROMPT_FILE_NAME = "prompt.txt"  # Standard name for the prompt file
 # Names for the output images from different models.
@@ -11,7 +11,7 @@ MODEL_OUTPUT_IMAGE_NAMES = {
     "baseline": "cp_bg_fg.jpg",
     "kv-edit": "kvedit.jpg",
     "tf-icon": "tf-icon.png",
-    "dit-editor": "alphanoise0.05_timesteps50_QTrue_KTrue_VTrue_taua0.4_taub0.8_guidance3.0.png",
+    "dit-editor": "alphanoise0.05_timesteps50_QTrue_KTrue_VFalse_taua0.4_taub0.8_guidance3.0_all-layers.png",
 }
 # Names to display for each model in the UI (can be different from filenames)
 MODEL_DISPLAY_NAMES = {
@@ -43,7 +43,7 @@ SAMPLES_PER_DOMAIN = 3 # Number of samples to show from each domain per user ses
 # --- Hugging Face Hub ---
 HF_DATASET_REPO_ID = "matsant01/dit-editor-collected-preferences"  # Replace with your actual repo ID
 HF_TOKEN = None  # Set this if your dataset is private, or use HF_HUB_TOKEN env var
-PUSH_INTERVAL_HOURS = 1 # Interval in hours to push results to the Hub
+PUSH_INTERVAL_HOURS = 0.05 # Interval in hours to push results to the Hub
 
 # --- UI Configuration ---
 IMAGE_DISPLAY_SIZE = (300, 300) # (width, height) for displaying images