Update src/test_set.py

src/test_set.py

@@ -68,24 +68,31 @@ def generate_test_set(max_samples_per_pair: int = MAX_TEST_SAMPLES) -> pd.DataFr
     
     return test_df
 
+
 def get_public_test_set() -> pd.DataFrame:
     """Get public test set (sources only, no targets)."""
     
     try:
         # Try to load existing test set
+        print(f"Loading test set from: {TEST_SET_DATASET}")
         dataset = load_dataset(TEST_SET_DATASET, split='train')
         test_df = dataset.to_pandas()
         print(f"Loaded existing test set with {len(test_df)} samples")
         
     except Exception as e:
         print(f"Could not load existing test set: {e}")
-        print("Generating new test set...")
+        print("This is expected for first run. Generating new test set...")
         
         # Generate new test set
         test_df = generate_test_set()
         
         # Save complete test set (with targets) privately
-        save_complete_test_set(test_df)
+        print("Saving test set for future use...")
+        try:
+            save_complete_test_set(test_df)
+        except Exception as save_error:
+            print(f"Warning: Could not save test set: {save_error}")
+            print("Continuing with generated test set...")
     
     # Return public version (without targets)
     public_columns = [