Update app.py

app.py

@@ -20,9 +20,31 @@ except Exception as e:
     print(f"❌ Model Error: {e}")
 
 # --- 2. LOAD DATA ---
+# --- 2. LOAD DATA (SMART MATCHING) ---
 print("⏳ Loading Dataset...")
-# Load exact 5k subset used in training
-dataset = load_dataset("ethz/food101", split="train").shuffle(seed=42).select(range(5000))
+
+# 1. Load the Embeddings File FIRST
+df = pd.read_parquet(DATA_FILE)
+valid_indices = df.index.tolist() # Assuming you preserved the original indices in the dataframe index
+# OR if you reset the index in the notebook, we just check the length:
+num_embeddings = len(df)
+
+print(f"   👉 Embeddings file has {num_embeddings} rows.")
+
+# 2. Load the Dataset
+dataset_full = load_dataset("ethz/food101", split="train").shuffle(seed=42).select(range(5000))
+
+# 3. CRITICAL FIX: If lengths don't match, we assume the parquet is a subset.
+# (This is a guess - if you didn't save the original indices, this might still be slightly off,
+# but it prevents the 'IndexError' crash).
+if len(dataset_full) > num_embeddings:
+    print(f"⚠️ DATA MISMATCH DETECTED: Dataset has {len(dataset_full)} but Parquet has {num_embeddings}.")
+    print("   ✂️ Truncating dataset to match Parquet length...")
+    dataset = dataset_full.select(range(num_embeddings))
+else:
+    dataset = dataset_full
+
+print(f"✅ Final Dataset Size: {len(dataset)}")
 
 # --- 3. LOAD EMBEDDINGS ---
 print(f"⏳ Loading Embeddings from {DATA_FILE}...")