dslighting 1.7.1__py3-none-any.whl → 1.7.6__py3-none-any.whl

mlebench/competitions/inaturalist-2019-fgvc6/prepare.py

@@ -0,0 +1,259 @@
+import json
+import random
+import shutil
+import tarfile
+from pathlib import Path
+
+import pandas as pd
+from sklearn.model_selection import train_test_split
+from tqdm import tqdm
+
+from mlebench.utils import get_logger
+
+logger = get_logger(__name__)
+
+
+def add_to_tar(src: Path, out: Path):
+
+    assert src.exists(), f"Source directory `{src}` does not exist."
+    assert src.is_dir(), f"Expected a directory, but got `{src}`."
+
+    tqdm_desc = f"Taring {src.name} to {out.name}"
+    file_paths = [path for path in src.rglob("*") if path.is_file()]
+    total_files = len(file_paths)
+
+    with tarfile.open(out, "w") as tar:
+        for file_path in tqdm(file_paths, desc=tqdm_desc, unit="file", total=total_files):
+            tar.add(file_path, arcname=file_path.relative_to(src))
+
+
+def prepare(raw: Path, public: Path, private: Path):
+
+    dev_mode = False
+    image_count = 2 if dev_mode else float("inf")  # We copy over 2 images per category for dev mode
+
+    # Extract train_val2019.tar.gz which contains images
+    train_tar_path = raw / "train_val2019.tar.gz"
+    train_extract_path = raw
+    if not (raw / "train_val2019").exists():
+        shutil.unpack_archive(train_tar_path, train_extract_path)
+
+    # Create train, test from train split
+    json_path = raw / "train2019.json"
+    with open(json_path, "r", encoding="utf-8") as f:
+        old_train_metadata = json.load(f)
+
+    # Organize data by category so that we can split per-category later
+    annotation_image_metadata_by_category = {}  # We'll collect both `annotations` and `images` here
+    for annotation_info, image_info in list(
+        zip(old_train_metadata["annotations"], old_train_metadata["images"])
+    ):
+        assert (
+            annotation_info["image_id"] == image_info["id"]
+        ), f"Mismatching image_id in annotation and image: {annotation_info['image_id']} vs {image_info['id']}"
+        category_id = annotation_info["category_id"]
+        if category_id not in annotation_image_metadata_by_category:
+            annotation_image_metadata_by_category[category_id] = []
+        annotation_image_metadata_by_category[category_id].append(
+            {
+                "annotation": annotation_info,
+                "image": image_info,
+            }
+        )
+
+    # Split train/test
+    train_sample_count = 0  # Useful for tqdm later
+    train_annotation_image_metadata_by_category = {}
+    test_annotation_image_metadata_by_category = {}
+
+    for category_id, annotation_image_metadata in tqdm(
+        annotation_image_metadata_by_category.items(), desc="Assigning train/test splits"
+    ):
+        # Create split by "category" (class)
+        # Original train+val has 268,243 images, test has 35,400 images, 0.12 ratio
+        test_size = 0.12
+        n_samples = len(annotation_image_metadata)
+        if n_samples == 1:
+            # If only one sample, put it in train
+            train_annotations_images = annotation_image_metadata
+            test_annotations_images = []
+        elif n_samples < 5:  # Minimum 5 samples to ensure at least 1 in test
+            num_test_samples = max(1, int(n_samples * test_size))
+            train_annotations_images = annotation_image_metadata[:-num_test_samples]
+            test_annotations_images = annotation_image_metadata[-num_test_samples:]
+        else:
+            train_annotations_images, test_annotations_images = train_test_split(
+                annotation_image_metadata, test_size=test_size, random_state=0
+            )
+
+        train_annotation_image_metadata_by_category[category_id] = train_annotations_images
+        test_annotation_image_metadata_by_category[category_id] = test_annotations_images
+        train_sample_count += len(train_annotations_images)
+
+    # Create new train2019.json
+    new_train_metadata = (
+        old_train_metadata.copy()
+    )  # Keep 'info', 'categories', 'licenses' unchanged
+    new_train_metadata.update(
+        {
+            "annotations": [],
+            "images": [],
+        }
+    )
+    for category_id, annotation_image_metadata in tqdm(
+        train_annotation_image_metadata_by_category.items(),
+        desc="Creating new train2019.json",
+        total=len(train_annotation_image_metadata_by_category),
+    ):
+        for annotation_image in annotation_image_metadata:
+            new_annotation = annotation_image["annotation"].copy()
+            new_train_metadata["annotations"].append(new_annotation)
+            new_image = annotation_image["image"].copy()
+            new_train_metadata["images"].append(new_image)
+
+    with open(public / "train2019.json", "w") as f:
+        json.dump(new_train_metadata, f, indent=4, sort_keys=True)
+
+    # Copy over val2019.json
+    shutil.copy(raw / "val2019.json", public / "val2019.json")
+    logger.info(f"Copied {raw / 'val2019.json'} to {public / 'val2019.json'}")
+
+    # Create new test2019.json
+    new_to_old_file_name = {}
+    new_test_metadata = old_train_metadata.copy()
+    del new_test_metadata["categories"]
+    new_test_metadata.update(
+        {
+            "annotations": [],
+            "images": [],
+        }
+    )
+    # Flatten and shuffle test set so that we don't have all the same catedgories in a row
+    test_annotations_images = [
+        item for sublist in test_annotation_image_metadata_by_category.values() for item in sublist
+    ]
+    random.Random(0).shuffle(test_annotations_images)
+    for idx, annotation_image in tqdm(
+        enumerate(test_annotations_images),
+        desc="Creating new test2019.json",
+        total=len(test_annotations_images),
+    ):
+
+        new_annotation = annotation_image["annotation"].copy()
+        new_test_metadata["annotations"].append(new_annotation)
+
+        new_image = annotation_image["image"].copy()
+        old_file_name = new_image["file_name"]
+        # go from e.g. "train_val2019/Plants/400/d1322d13ccd856eb4236c8b888546c79.jpg" to "test2019/d1322d13ccd856eb4236c8b888546c79.jpg"
+        new_file_name = "test2019/" + old_file_name.split("/")[-1]
+        # keep track of things so we know what to copy later
+        new_to_old_file_name[new_file_name] = old_file_name
+        new_image["file_name"] = new_file_name
+        new_test_metadata["images"].append(new_image)
+    with open(public / "test2019.json", "w") as f:
+        # The public test data, of course, doesn't have annotations
+        public_new_test = new_test_metadata.copy()
+        del public_new_test["annotations"]
+        assert public_new_test.keys() == {
+            "images",
+            "info",
+            "licenses",
+        }, f"Public test metadata keys should be 'images', 'info', 'licenses', but found {public_new_test.keys()}"
+        json.dump(public_new_test, f, indent=4, sort_keys=True)
+
+    (public / "train_val2019").mkdir(parents=True, exist_ok=True)
+    (public / "test2019").mkdir(parents=True, exist_ok=True)
+
+    # Save private test answers
+    answers_rows = []
+    for image_info, annotation_info in zip(
+        new_test_metadata["images"], new_test_metadata["annotations"]
+    ):
+        assert (
+            image_info["id"] == annotation_info["image_id"]
+        ), f"Mismatching image_id in image and annotation: {image_info['id']} vs {annotation_info['image_id']}"
+        answers_rows.append(
+            {
+                "id": image_info["id"],
+                "predicted": annotation_info["category_id"],
+            }
+        )
+    answers_df = pd.DataFrame(answers_rows)
+    answers_df.to_csv(private / "answers.csv", index=False)
+
+    # Create new sample submission based on answers_df
+    sample_df = answers_df.copy()
+    sample_df["predicted"] = [random.Random(42).randint(0, 1009) for _ in range(len(sample_df))]
+    sample_df.to_csv(public / "kaggle_sample_submission.csv", index=False)
+
+    assert len(answers_df) == len(
+        new_test_metadata["images"]
+    ), f"Expected {len(new_test_metadata['images'])} rows in answers, but found {len(answers_df)}"
+    assert len(sample_df) == len(
+        answers_df
+    ), f"Expected {len(answers_df)} rows in sample submission, but found {len(sample_df)}"
+    assert answers_df["id"].equals(
+        sample_df["id"]
+    ), "Mismatched 'id' columns between answers and sample submission"
+
+    # Copy train images
+    train_images_copied = 0
+    for category_id, annotation_image_metadata in tqdm(
+        train_annotation_image_metadata_by_category.items(),
+        desc="Copying train images grouped by category",
+    ):
+        for idx, annotation_image in enumerate(annotation_image_metadata):
+            if dev_mode and idx >= image_count:
+                break
+            old_path = raw / annotation_image["image"]["file_name"]
+            new_path = public / annotation_image["image"]["file_name"]
+            new_path.parent.mkdir(parents=True, exist_ok=True)
+            shutil.copy(old_path, new_path)
+            train_images_copied += 1
+
+    # Copy test images
+    test_images_copied = 0
+    for image_info in tqdm(new_test_metadata["images"], desc="Copying test images"):
+        if dev_mode and test_images_copied >= image_count:
+            break
+        old_path = raw / new_to_old_file_name[image_info["file_name"]]
+        new_path = public / image_info["file_name"]
+        new_path.parent.mkdir(parents=True, exist_ok=True)
+        shutil.copy(old_path, new_path)
+        test_images_copied += 1
+
+    logger.info(f"Copied {train_images_copied} train images and {test_images_copied} test images")
+
+    if not dev_mode:
+        assert len(list((public / "train_val2019").glob("**/*.jpg"))) == len(
+            new_train_metadata["images"]
+        ), f"Mismatching number of images in train_images, got {len(list((public / 'train_val2019').glob('**/*.jpg')))}"
+
+        assert len(new_train_metadata["annotations"]) == len(
+            new_train_metadata["images"]
+        ), f"Mismatching number of annotations in train_metadata, got {len(new_train_metadata['annotations'])}"
+
+    if not dev_mode:
+        assert len(list((public / "test2019").glob("**/*.jpg"))) == len(
+            new_test_metadata["images"]
+        ), f"Mismatching number of images in test_images, got {len(list((public / 'test2019').glob('**/*.jpg')))}"
+
+        assert len(new_test_metadata["annotations"]) == len(
+            new_test_metadata["images"]
+        ), f"Mismatching number of annotations in test_metadata, got {len(new_test_metadata['annotations'])}"
+
+        assert len(new_train_metadata["annotations"]) + len(
+            new_test_metadata["annotations"]
+        ) == len(old_train_metadata["annotations"]), (
+            f"Expected {len(old_train_metadata['annotations'])} annotations in total, but found"
+            f"{len(new_train_metadata['annotations'])} in train and {len(new_test_metadata['annotations'])} in test"
+        )
+
+    # Re-compress images
+    add_to_tar(
+        src=public / "test2019", out=public / "test2019.tar.gz"
+    )  # Add to tar but don't actually compress with gzip to save time
+    add_to_tar(src=public / "train_val2019", out=public / "train_val2019.tar.gz")
+    # Remove uncompressed directories
+    shutil.rmtree(public / "test2019")
+    shutil.rmtree(public / "train_val2019")

mlebench/competitions/inaturalist-2019-fgvc6/prepare_val.py

@@ -0,0 +1,304 @@
+import json
+import random
+import shutil
+import tarfile
+from pathlib import Path
+from typing import Dict, List
+
+import pandas as pd
+from sklearn.model_selection import train_test_split
+from tqdm import tqdm
+
+from mlebench.utils import get_logger
+
+logger = get_logger(__name__)
+
+
+def add_to_tar(src: Path, out: Path):
+
+    assert src.exists(), f"Source directory `{src}` does not exist."
+    assert src.is_dir(), f"Expected a directory, but got `{src}`."
+
+    tqdm_desc = f"Taring {src.name} to {out.name}"
+    file_paths = [path for path in src.rglob("*") if path.is_file()]
+    total_files = len(file_paths)
+
+    with tarfile.open(out, "w") as tar:
+        for file_path in tqdm(file_paths, desc=tqdm_desc, unit="file", total=total_files):
+            tar.add(file_path, arcname=file_path.relative_to(src))
+
+
+def _split_by_category(
+    data_by_category: Dict, test_size: float, random_state: int
+) -> tuple[Dict, Dict]:
+    """Splits data for each category into train and test sets using the original script's logic."""
+    train_split_by_category = {}
+    test_split_by_category = {}
+
+    for category_id, annotation_image_metadata in tqdm(
+        data_by_category.items(),
+        desc=f"Assigning train/test splits (test_size={test_size:.3f})",
+    ):
+        n_samples = len(annotation_image_metadata)
+        if n_samples == 1:
+            # If only one sample, put it in train
+            train_annotations_images = annotation_image_metadata
+            test_annotations_images = []
+        elif n_samples < 5:  # Minimum 5 samples to ensure at least 1 in test
+            num_test_samples = max(1, int(n_samples * test_size))
+            train_annotations_images = annotation_image_metadata[:-num_test_samples]
+            test_annotations_images = annotation_image_metadata[-num_test_samples:]
+        else:
+            train_annotations_images, test_annotations_images = train_test_split(
+                annotation_image_metadata, test_size=test_size, random_state=random_state
+            )
+
+        train_split_by_category[category_id] = train_annotations_images
+        test_split_by_category[category_id] = test_annotations_images
+
+    return train_split_by_category, test_split_by_category
+
+
+def _generate_split_files(
+    train_annotation_image_metadata_by_category: Dict,
+    test_annotation_image_metadata_by_category: Dict,
+    old_train_metadata: Dict,
+    raw_path: Path,
+    public_path: Path,
+    private_path: Path,
+    dev_mode: bool,
+    image_count: int,
+):
+    """
+    Processes a given train/test split and saves all corresponding files
+    (metadata, images, private answers, etc.) to the specified output directories.
+    """
+    public_path.mkdir(parents=True, exist_ok=True)
+    private_path.mkdir(parents=True, exist_ok=True)
+
+    # Create new train2019.json
+    new_train_metadata = (
+        old_train_metadata.copy()
+    )  # Keep 'info', 'categories', 'licenses' unchanged
+    new_train_metadata.update(
+        {
+            "annotations": [],
+            "images": [],
+        }
+    )
+    for category_id, annotation_image_metadata in tqdm(
+        train_annotation_image_metadata_by_category.items(),
+        desc=f"[{public_path.name}] Creating new train2019.json",
+        total=len(train_annotation_image_metadata_by_category),
+    ):
+        for annotation_image in annotation_image_metadata:
+            new_annotation = annotation_image["annotation"].copy()
+            new_train_metadata["annotations"].append(new_annotation)
+            new_image = annotation_image["image"].copy()
+            new_train_metadata["images"].append(new_image)
+
+    with open(public_path / "train2019.json", "w") as f:
+        json.dump(new_train_metadata, f, indent=4, sort_keys=True)
+
+    # Copy over val2019.json
+    shutil.copy(raw_path / "val2019.json", public_path / "val2019.json")
+
+    # Create new test2019.json
+    new_to_old_file_name = {}
+    new_test_metadata = old_train_metadata.copy()
+    del new_test_metadata["categories"]
+    new_test_metadata.update(
+        {
+            "annotations": [],
+            "images": [],
+        }
+    )
+    # Flatten and shuffle test set so that we don't have all the same categories in a row
+    test_annotations_images = [
+        item
+        for sublist in test_annotation_image_metadata_by_category.values()
+        for item in sublist
+    ]
+    random.Random(0).shuffle(test_annotations_images)
+    for idx, annotation_image in tqdm(
+        enumerate(test_annotations_images),
+        desc=f"[{public_path.name}] Creating new test2019.json",
+        total=len(test_annotations_images),
+    ):
+        new_annotation = annotation_image["annotation"].copy()
+        new_test_metadata["annotations"].append(new_annotation)
+
+        new_image = annotation_image["image"].copy()
+        old_file_name = new_image["file_name"]
+        # go from e.g. "train_val2019/Plants/400/d1322d13ccd856eb4236c8b888546c79.jpg" to "test2019/d1322d13ccd856eb4236c8b888546c79.jpg"
+        new_file_name = "test2019/" + old_file_name.split("/")[-1]
+        # keep track of things so we know what to copy later
+        new_to_old_file_name[new_file_name] = old_file_name
+        new_image["file_name"] = new_file_name
+        new_test_metadata["images"].append(new_image)
+
+    with open(public_path / "test2019.json", "w") as f:
+        # The public test data, of course, doesn't have annotations
+        public_new_test = new_test_metadata.copy()
+        del public_new_test["annotations"]
+        assert public_new_test.keys() == {
+            "images",
+            "info",
+            "licenses",
+        }, f"Public test metadata keys should be 'images', 'info', 'licenses', but found {public_new_test.keys()}"
+        json.dump(public_new_test, f, indent=4, sort_keys=True)
+
+    (public_path / "train_val2019").mkdir(parents=True, exist_ok=True)
+    (public_path / "test2019").mkdir(parents=True, exist_ok=True)
+
+    # Save private test answers
+    answers_rows = []
+    for image_info, annotation_info in zip(
+        new_test_metadata["images"], new_test_metadata["annotations"]
+    ):
+        assert (
+            image_info["id"] == annotation_info["image_id"]
+        ), f"Mismatching image_id in image and annotation: {image_info['id']} vs {annotation_info['image_id']}"
+        answers_rows.append(
+            {
+                "id": image_info["id"],
+                "predicted": annotation_info["category_id"],
+            }
+        )
+    answers_df = pd.DataFrame(answers_rows)
+    answers_df.to_csv(private_path / "answers.csv", index=False)
+
+    # Create new sample submission based on answers_df
+    sample_df = answers_df.copy()
+    sample_df["predicted"] = [random.Random(42).randint(0, 1009) for _ in range(len(sample_df))]
+    sample_df.to_csv(public_path / "kaggle_sample_submission.csv", index=False)
+
+    # Copy train images
+    for annotation_image_metadata in tqdm(
+        train_annotation_image_metadata_by_category.values(),
+        desc=f"[{public_path.name}] Copying train images",
+    ):
+        for idx, annotation_image in enumerate(annotation_image_metadata):
+            if dev_mode and idx >= image_count:
+                break
+            old_path = raw_path / annotation_image["image"]["file_name"]
+            new_path = public_path / annotation_image["image"]["file_name"]
+            new_path.parent.mkdir(parents=True, exist_ok=True)
+            shutil.copy(old_path, new_path)
+
+    # Copy test images
+    for image_info in tqdm(
+        new_test_metadata["images"], desc=f"[{public_path.name}] Copying test images"
+    ):
+        if dev_mode and len(new_to_old_file_name) >= image_count:
+            break
+        old_path = raw_path / new_to_old_file_name[image_info["file_name"]]
+        new_path = public_path / image_info["file_name"]
+        new_path.parent.mkdir(parents=True, exist_ok=True)
+        shutil.copy(old_path, new_path)
+
+    # Re-compress images
+    add_to_tar(
+        src=public_path / "test2019", out=public_path / "test2019.tar.gz"
+    )  # Add to tar but don't actually compress with gzip to save time
+    add_to_tar(src=public_path / "train_val2019", out=public_path / "train_val2019.tar.gz")
+    # Remove uncompressed directories
+    shutil.rmtree(public_path / "test2019")
+    shutil.rmtree(public_path / "train_val2019")
+    logger.info(f"Finished generating files for {public_path.name}")
+
+
+def prepare(raw: Path, public: Path, private: Path):
+
+    dev_mode = False
+    image_count = 2 if dev_mode else float("inf")  # We copy over 2 images per category for dev mode
+
+    # Extract train_val2019.tar.gz which contains images
+    train_tar_path = raw / "train_val2019.tar.gz"
+    train_extract_path = raw
+    if not (raw / "train_val2019").exists():
+        logger.info("Extracting raw image data...")
+        shutil.unpack_archive(train_tar_path, train_extract_path)
+
+    # Create train, test from train split
+    json_path = raw / "train2019.json"
+    with open(json_path, "r", encoding="utf-8") as f:
+        old_train_metadata = json.load(f)
+
+    # Organize data by category so that we can split per-category later
+    annotation_image_metadata_by_category = {}  # We'll collect both `annotations` and `images` here
+    for annotation_info, image_info in list(
+        zip(old_train_metadata["annotations"], old_train_metadata["images"])
+    ):
+        assert (
+            annotation_info["image_id"] == image_info["id"]
+        ), f"Mismatching image_id in annotation and image: {annotation_info['image_id']} vs {image_info['id']}"
+        category_id = annotation_info["category_id"]
+        if category_id not in annotation_image_metadata_by_category:
+            annotation_image_metadata_by_category[category_id] = []
+        annotation_image_metadata_by_category[category_id].append(
+            {
+                "annotation": annotation_info,
+                "image": image_info,
+            }
+        )
+
+    # --- 1. Original Data Split (Train/Test) ---
+    logger.info("--- Generating Original Train/Test Split ---")
+    # Original train+val has 268,243 images, test has 35,400 images, ~0.12 ratio
+    original_test_size = 0.12
+    (
+        original_train_split,
+        original_test_split,
+    ) = _split_by_category(
+        annotation_image_metadata_by_category,
+        test_size=original_test_size,
+        random_state=0,
+    )
+
+    _generate_split_files(
+        original_train_split,
+        original_test_split,
+        old_train_metadata,
+        raw,
+        public,
+        private,
+        dev_mode,
+        image_count,
+    )
+    logger.info(f"Original split saved to {public.name} and {private.name}")
+
+    # --- 2. New Validation Data Split (Train/Val) ---
+    logger.info("--- Generating New Train/Validation Split ---")
+    # Define new output directories
+    public_val = public.parent / "public_val"
+    private_val = private.parent / "private_val"
+
+    # We want the new validation set ('test_val') to have the same size as the
+    # original test set. We are splitting the *original_train_split* to get it.
+    # test_val_size = new_test_size * train_size
+    # We want: test_val_size ≈ test_size
+    # So: new_test_size * (N * (1-0.12)) ≈ N * 0.12
+    # new_test_size ≈ 0.12 / (1 - 0.12)
+    val_split_test_size = original_test_size / (1.0 - original_test_size)
+
+    (
+        validation_train_split,
+        validation_test_split,
+    ) = _split_by_category(
+        original_train_split,  # Split the TRAIN set from the first split
+        test_size=val_split_test_size,
+        random_state=0,  # Use same random state for consistency
+    )
+
+    _generate_split_files(
+        validation_train_split,
+        validation_test_split,
+        old_train_metadata,
+        raw,
+        public_val,
+        private_val,
+        dev_mode,
+        image_count,
+    )
+    logger.info(f"Validation split saved to {public_val.name} and {private_val.name}")

mlebench/competitions/instant-gratification/grade.py

@@ -0,0 +1,55 @@
+import numpy as np
+import pandas as pd
+from sklearn.metrics import roc_auc_score
+
+from mlebench.grade_helpers import InvalidSubmissionError
+
+
+def prepare_for_metric(submission: pd.DataFrame, answers: pd.DataFrame):
+    """Prepare submission and answers for AUC-ROC calculation."""
+    if len(submission) != len(answers):
+        raise InvalidSubmissionError(
+            f"Submission length ({len(submission)}) != answers length ({len(answers)})"
+        )
+
+    if "id" not in submission.columns:
+        raise InvalidSubmissionError("Submission must have an 'id' column")
+
+    if "target" not in submission.columns:
+        raise InvalidSubmissionError("Submission must have a 'target' column")
+
+    # Sort by id
+    submission = submission.sort_values("id").reset_index(drop=True)
+    answers = answers.sort_values("id").reset_index(drop=True)
+
+    # Check id alignment
+    if not (submission["id"] == answers["id"]).all():
+        raise InvalidSubmissionError("Submission and answers id mismatch")
+
+    # Extract predictions and true values
+    y_pred = submission["target"].values
+    y_true = answers["target"].values
+
+    # Validate predictions
+    try:
+        y_pred = y_pred.astype(float)
+    except (ValueError, TypeError):
+        raise InvalidSubmissionError("Predictions must be numeric")
+
+    if np.any(np.isnan(y_pred)):
+        raise InvalidSubmissionError("Predictions cannot contain NaN values")
+
+    if np.any((y_pred < 0) | (y_pred > 1)):
+        raise InvalidSubmissionError("Predictions must be probabilities between 0 and 1")
+
+    return y_true, y_pred
+
+
+def grade(submission: pd.DataFrame, answers: pd.DataFrame) -> float:
+    """Calculate AUC-ROC score."""
+    y_true, y_pred = prepare_for_metric(submission, answers)
+
+    # AUC-ROC
+    auc = roc_auc_score(y_true, y_pred)
+
+    return auc

mlebench/competitions/instant-gratification/prepare.py

@@ -0,0 +1,25 @@
+from pathlib import Path
+import pandas as pd
+
+
+def prepare(raw: Path, public: Path, private: Path):
+    """
+    Prepare instant-gratification dataset.
+    """
+    # Read data
+    train = pd.read_csv(raw / "train.csv")
+    test = pd.read_csv(raw / "test.csv")
+    sample_submission = pd.read_csv(raw / "sample_submission.csv")
+    test_answer = pd.read_csv(raw / "test_answer.csv")
+
+    # Public files (visible to agents)
+    train.to_csv(public / "train.csv", index=False)
+    test.to_csv(public / "test.csv", index=False)
+    sample_submission.to_csv(public / "sample_submission.csv", index=False)
+
+    # Private files (for grading)
+    test_answer.to_csv(private / "test.csv", index=False)
+
+    # Validation checks
+    assert len(test_answer) == len(sample_submission), \
+        f"Test answer ({len(test_answer)}) and sample submission ({len(sample_submission)}) must have same length"