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

mlebench/competitions/learning-agency-lab-automated-essay-scoring-2/prepare_val.py

@@ -0,0 +1,96 @@
+from pathlib import Path
+from typing import Tuple
+
+import numpy as np
+import pandas as pd
+from sklearn.model_selection import train_test_split
+
+from mlebench.utils import read_csv
+
+
+def _create_split(
+    input_df: pd.DataFrame, test_size: float, random_state: int
+) -> Tuple[pd.DataFrame, pd.DataFrame, pd.DataFrame, pd.DataFrame]:
+    """Helper function to perform a single data split and generate associated files."""
+    # Create train, test from the input dataframe
+    train_df, answers_df = train_test_split(
+        input_df, test_size=test_size, random_state=random_state
+    )
+    test_df = answers_df.drop(columns=["score"])
+
+    sample_submission_df = answers_df[["essay_id"]].copy()
+    sample_submission_df["score"] = np.random.RandomState(42).randint(
+        1, 7, size=len(sample_submission_df)
+    )
+
+    # Checks
+    assert set(train_df["essay_id"]).isdisjoint(
+        set(test_df["essay_id"])
+    ), "Essay IDs in train and test sets are not disjoint"
+    assert len(train_df) + len(test_df) == len(
+        input_df
+    ), f"Train and test sets do not sum to original train set"
+    assert len(test_df) == len(
+        sample_submission_df
+    ), f"Test and sample submission sets do not have the same length"
+    assert (
+        train_df.columns.tolist() == input_df.columns.tolist()
+    ), f"Train set columns do not match original train set, got {train_df.columns.tolist()}"
+    assert test_df.columns.tolist() == [
+        "essay_id",
+        "full_text",
+    ], f"Test set columns do not match expected columns, got {test_df.columns.tolist()}"
+    assert sample_submission_df.columns.tolist() == [
+        "essay_id",
+        "score",
+    ], f"Sample submission set columns do not match expected columns, got {sample_submission_df.columns.tolist()}"
+
+    return train_df, test_df, answers_df, sample_submission_df
+
+
+def prepare(raw: Path, public: Path, private: Path):
+    """
+    Splits the data in raw into public and private datasets with appropriate test/train splits.
+    Also creates a second, parallel split for validation purposes.
+    """
+
+    # Read the original raw data
+    old_train = read_csv(raw / "train.csv")
+
+    # --- Stage 1: Create the original train/test split for the main competition ---
+    # This block produces the exact same output as the original script.
+
+    # Original train has 17307 rows. Original hidden test has approx 8k rows. We just take 10% of the original train as the test set.
+    main_train, main_test, main_answers, main_sample_submission = _create_split(
+        input_df=old_train, test_size=0.1, random_state=0
+    )
+
+    # Write original CSVs to public/ and private/
+    main_answers.to_csv(private / "answers.csv", index=False)
+    main_train.to_csv(public / "train.csv", index=False)
+    main_test.to_csv(public / "test.csv", index=False)
+    main_sample_submission.to_csv(public / "sample_submission.csv", index=False)
+
+    # --- Stage 2: Create a new validation split from the main training data ---
+    # This block creates a new set of directories and files for validation.
+
+    # Define and create the new parallel directories for the validation set
+    public_val = public.parent / "public_val"
+    private_val = private.parent / "private_val"
+    public_val.mkdir(parents=True, exist_ok=True)
+    private_val.mkdir(parents=True, exist_ok=True)
+
+    # Calculate the test size for the second split to make the validation test set
+    # have the same number of samples as the original test set.
+    val_test_size = len(main_test) / len(main_train)
+
+    # Create the new split using the main training data as input
+    val_train, val_test, val_answers, val_sample_submission = _create_split(
+        input_df=main_train, test_size=val_test_size, random_state=0
+    )
+
+    # Write validation CSVs to public_val/ and private_val/ using identical filenames
+    val_answers.to_csv(private_val / "answers.csv", index=False)
+    val_train.to_csv(public_val / "train.csv", index=False)
+    val_test.to_csv(public_val / "test.csv", index=False)
+    val_sample_submission.to_csv(public_val / "sample_submission.csv", index=False)

mlebench/competitions/liverpool-ion-switching/grade.py

@@ -0,0 +1,52 @@
+import numpy as np
+import pandas as pd
+from sklearn.metrics import f1_score
+
+from mlebench.grade_helpers import InvalidSubmissionError
+
+
+def prepare_for_metric(submission: pd.DataFrame, answers: pd.DataFrame):
+    """Prepare submission and answers for Macro F1 calculation."""
+    if len(submission) != len(answers):
+        raise InvalidSubmissionError(
+            f"Submission length ({len(submission)}) != answers length ({len(answers)})"
+        )
+
+    if "time" not in submission.columns:
+        raise InvalidSubmissionError("Submission must have a 'time' column")
+
+    if "open_channels" not in submission.columns:
+        raise InvalidSubmissionError("Submission must have an 'open_channels' column")
+
+    # Sort by time
+    submission = submission.sort_values("time").reset_index(drop=True)
+    answers = answers.sort_values("time").reset_index(drop=True)
+
+    # Check time alignment
+    if not np.allclose(submission["time"].values, answers["time"].values, atol=1e-4):
+        raise InvalidSubmissionError("Submission and answers time mismatch")
+
+    # Extract predictions and true values
+    y_pred = submission["open_channels"].values
+    y_true = answers["open_channels"].values
+
+    # Validate predictions
+    try:
+        y_pred = y_pred.astype(int)
+    except (ValueError, TypeError):
+        raise InvalidSubmissionError("Predictions must be integers")
+
+    if np.any(np.isnan(y_pred)):
+        raise InvalidSubmissionError("Predictions cannot contain NaN values")
+
+    return y_true, y_pred
+
+
+def grade(submission: pd.DataFrame, answers: pd.DataFrame) -> float:
+    """Calculate Macro F1 score."""
+    y_true, y_pred = prepare_for_metric(submission, answers)
+
+    # Macro F1: compute F1 for each class and average
+    macro_f1 = f1_score(y_true, y_pred, average='macro')
+
+    return macro_f1

mlebench/competitions/liverpool-ion-switching/prepare.py

@@ -0,0 +1,27 @@
+from pathlib import Path
+import pandas as pd
+
+
+def prepare(raw: Path, public: Path, private: Path):
+    """
+    Prepare liverpool-ion-switching 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"
+    assert "time" in test_answer.columns, "Test answer must have 'time' column"
+    assert "open_channels" in test_answer.columns, "Test answer must have 'open_channels' column"

mlebench/competitions/liverpool_ion_switching/grade.py

@@ -0,0 +1,52 @@
+import numpy as np
+import pandas as pd
+from sklearn.metrics import f1_score
+
+from mlebench.grade_helpers import InvalidSubmissionError
+
+
+def prepare_for_metric(submission: pd.DataFrame, answers: pd.DataFrame):
+    """Prepare submission and answers for Macro F1 calculation."""
+    if len(submission) != len(answers):
+        raise InvalidSubmissionError(
+            f"Submission length ({len(submission)}) != answers length ({len(answers)})"
+        )
+
+    if "time" not in submission.columns:
+        raise InvalidSubmissionError("Submission must have a 'time' column")
+
+    if "open_channels" not in submission.columns:
+        raise InvalidSubmissionError("Submission must have an 'open_channels' column")
+
+    # Sort by time
+    submission = submission.sort_values("time").reset_index(drop=True)
+    answers = answers.sort_values("time").reset_index(drop=True)
+
+    # Check time alignment
+    if not np.allclose(submission["time"].values, answers["time"].values, atol=1e-4):
+        raise InvalidSubmissionError("Submission and answers time mismatch")
+
+    # Extract predictions and true values
+    y_pred = submission["open_channels"].values
+    y_true = answers["open_channels"].values
+
+    # Validate predictions
+    try:
+        y_pred = y_pred.astype(int)
+    except (ValueError, TypeError):
+        raise InvalidSubmissionError("Predictions must be integers")
+
+    if np.any(np.isnan(y_pred)):
+        raise InvalidSubmissionError("Predictions cannot contain NaN values")
+
+    return y_true, y_pred
+
+
+def grade(submission: pd.DataFrame, answers: pd.DataFrame) -> float:
+    """Calculate Macro F1 score."""
+    y_true, y_pred = prepare_for_metric(submission, answers)
+
+    # Macro F1: compute F1 for each class and average
+    macro_f1 = f1_score(y_true, y_pred, average='macro')
+
+    return macro_f1

mlebench/competitions/liverpool_ion_switching/prepare.py

@@ -0,0 +1,27 @@
+from pathlib import Path
+import pandas as pd
+
+
+def prepare(raw: Path, public: Path, private: Path):
+    """
+    Prepare liverpool-ion-switching 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"
+    assert "time" in test_answer.columns, "Test answer must have 'time' column"
+    assert "open_channels" in test_answer.columns, "Test answer must have 'open_channels' column"

mlebench/competitions/lmsys-chatbot-arena/grade.py

@@ -0,0 +1,63 @@
+import numpy as np
+import pandas as pd
+from sklearn.metrics import log_loss
+
+from mlebench.competitions.utils import one_hot_dfs_to_log_loss_inputs
+from mlebench.grade_helpers import InvalidSubmissionError
+
+
+def prepare_for_metric(submission: pd.DataFrame, answers: pd.DataFrame) -> dict:
+    """
+    The submission and answers are already one-hotted
+    """
+    classes = ["winner_model_a", "winner_model_b", "winner_tie"]
+    required_columns = ["id"] + classes
+
+    # Check if submission has the required columns
+    missing_columns = [col for col in required_columns if col not in submission.columns]
+    if missing_columns:
+        raise InvalidSubmissionError(
+            f"Submission DataFrame is missing required columns: {missing_columns}"
+        )
+
+    # Check if answers has the required columns
+    assert set(required_columns).issubset(
+        answers.columns
+    ), f"Answers DataFrame is missing required columns: {set(required_columns) - set(answers.columns)}"
+
+    # Check if submission has the correct number of rows
+    if len(submission) != len(answers):
+        raise InvalidSubmissionError(
+            f"Submission DataFrame must have {len(answers)} rows, but has {len(submission)} rows."
+        )
+
+    # Check if all values in submission are between 0 and 1
+    if (
+        not ((submission[classes] >= 0) & (submission[classes] <= 1)).all().all()
+    ):  # first all() checks if all rows are valid, second all() checks if all columns are valid
+        raise InvalidSubmissionError("All values in submission DataFrame must be between 0 and 1.")
+
+    # Check if each row in submission sums to 1
+    if not submission[classes].sum(axis=1).round(6).eq(1).all():
+        raise InvalidSubmissionError("Each row in submission DataFrame must sum to 1.")
+
+    # Use only the required columns for further processing
+    submission = submission[required_columns]
+    answers = answers[required_columns]
+
+    submission = submission.sort_values("id").reset_index(drop=True)
+    answers = answers.sort_values("id").reset_index(drop=True)
+
+    if (submission["id"].values != answers["id"].values).any():
+        raise InvalidSubmissionError("Submission and answer IDs do not match after sorting.")
+
+    log_loss_inputs = one_hot_dfs_to_log_loss_inputs(
+        submission, answers, id_column="id", apply_softmax=False
+    )
+
+    return log_loss_inputs
+
+
+def grade(submission: pd.DataFrame, answers: pd.DataFrame) -> float:
+    log_loss_inputs = prepare_for_metric(submission, answers)
+    return log_loss(**log_loss_inputs)

mlebench/competitions/lmsys-chatbot-arena/prepare.py

@@ -0,0 +1,52 @@
+import shutil
+from pathlib import Path
+
+from sklearn.model_selection import train_test_split
+
+from mlebench.utils import extract, read_csv
+
+
+def prepare(raw: Path, public: Path, private: Path):
+    """
+    Splits the data in raw into public and private datasets with appropriate test/train splits.
+    """
+
+    # Create train, test from train split
+    old_train = read_csv(raw / "train.csv")
+    # Original train has 55k rows. Original hidden test has 25k rows. We make a new test set with 10% of the original train.
+    new_train, answers = train_test_split(old_train, test_size=0.1, random_state=0)
+    new_test = answers[["id", "prompt", "response_a", "response_b"]].copy()
+
+    sample_submission = answers[["id"]].copy()
+    sample_submission["winner_model_a"] = 0.3333333333333333
+    sample_submission["winner_model_b"] = 0.3333333333333333
+    sample_submission["winner_tie"] = 0.3333333333333333
+
+    # Checks
+    assert len(new_train) + len(new_test) == len(
+        old_train
+    ), f"New train and test should have the same number of rows as the original train"
+    assert set(new_train["id"]).isdisjoint(
+        set(new_test["id"])
+    ), f"New train and test should have no overlapping ids"
+    assert new_test.columns.tolist() == [
+        "id",
+        "prompt",
+        "response_a",
+        "response_b",
+    ], f"New test should have columns id, prompt, response_a, response_b"
+    assert sample_submission.columns.tolist() == [
+        "id",
+        "winner_model_a",
+        "winner_model_b",
+        "winner_tie",
+    ], f"Sample submission should have columns id, winner_model_a, winner_model_b, winner_tie"
+    assert (
+        new_train.columns.tolist() == old_train.columns.tolist()
+    ), f"New train should have the same columns as the original train"
+
+    # Write CSVs
+    answers.to_csv(private / "answers.csv", index=False)
+    new_train.to_csv(public / "train.csv", index=False)
+    new_test.to_csv(public / "test.csv", index=False)
+    sample_submission.to_csv(public / "sample_submission.csv", index=False)

mlebench/competitions/lmsys-chatbot-arena/prepare_val.py

@@ -0,0 +1,115 @@
+import shutil
+from pathlib import Path
+
+import pandas as pd
+from sklearn.model_selection import train_test_split
+
+from mlebench.utils import extract, read_csv
+
+
+def _create_split(
+    df_in: pd.DataFrame,
+    public_path: Path,
+    private_path: Path,
+    test_size: float,
+    random_state: int,
+) -> pd.DataFrame:
+    """
+    Helper function to perform a data split and write files to specified directories.
+
+    Args:
+        df_in: The input DataFrame to split.
+        public_path: The directory for public-facing files (train set, test features).
+        private_path: The directory for private files (test answers).
+        test_size: The proportion of the dataset to allocate to the test set.
+        random_state: The seed for the random number generator.
+
+    Returns:
+        The DataFrame for the newly created training set.
+    """
+    # Create output directories if they don't exist
+    public_path.mkdir(exist_ok=True, parents=True)
+    private_path.mkdir(exist_ok=True, parents=True)
+
+    # Create train, test from the input dataframe
+    new_train, answers = train_test_split(
+        df_in, test_size=test_size, random_state=random_state
+    )
+    new_test = answers[["id", "prompt", "response_a", "response_b"]].copy()
+
+    sample_submission = answers[["id"]].copy()
+    sample_submission["winner_model_a"] = 0.3333333333333333
+    sample_submission["winner_model_b"] = 0.3333333333333333
+    sample_submission["winner_tie"] = 0.3333333333333333
+
+    # Checks
+    assert len(new_train) + len(new_test) == len(
+        df_in
+    ), f"New train and test should have the same number of rows as the original dataframe"
+    assert set(new_train["id"]).isdisjoint(
+        set(new_test["id"])
+    ), f"New train and test should have no overlapping ids"
+    assert new_test.columns.tolist() == [
+        "id",
+        "prompt",
+        "response_a",
+        "response_b",
+    ], f"New test should have columns id, prompt, response_a, response_b"
+    assert sample_submission.columns.tolist() == [
+        "id",
+        "winner_model_a",
+        "winner_model_b",
+        "winner_tie",
+    ], f"Sample submission should have columns id, winner_model_a, winner_model_b, winner_tie"
+    assert (
+        new_train.columns.tolist() == df_in.columns.tolist()
+    ), f"New train should have the same columns as the original dataframe"
+
+    # Write CSVs
+    answers.to_csv(private_path / "answers.csv", index=False)
+    new_train.to_csv(public_path / "train.csv", index=False)
+    new_test.to_csv(public_path / "test.csv", index=False)
+    sample_submission.to_csv(public_path / "sample_submission.csv", index=False)
+
+    return new_train
+
+
+def prepare(raw: Path, public: Path, private: Path):
+    """
+    Splits the data in raw into public and private datasets with appropriate test/train splits.
+    Also creates a secondary validation split (public_val, private_val) for local testing.
+    """
+
+    # --- Stage 1: Create the original competition split (train/test) ---
+    # This block generates the primary `public` and `private` directories.
+    # Its outputs MUST remain identical to the original script's outputs.
+    old_train_df = read_csv(raw / "train.csv")
+    train_for_val_split = _create_split(
+        df_in=old_train_df,
+        public_path=public,
+        private_path=private,
+        test_size=0.1,
+        random_state=0,
+    )
+
+    # --- Stage 2: Create the new validation split (train_val/test_val) ---
+    # This block takes the training set from Stage 1 and splits it again
+    # to create a new, smaller training set and a validation set.
+    # The outputs are saved to parallel `public_val` and `private_val` directories.
+    public_val = public.parent / "public_val"
+    private_val = private.parent / "private_val"
+
+    # Calculate the test_size needed to make the new validation set (`test_val`)
+    # have the same number of samples as the original test set from Stage 1.
+    # Original test size = 0.1 * total. New train size = 0.9 * total.
+    # We need a fraction `p` such that p * (0.9 * total) = 0.1 * total.
+    # p = 0.1 / 0.9 = 1/9.
+    val_test_size = 1 / 9
+
+    _create_split(
+        df_in=train_for_val_split,
+        public_path=public_val,
+        private_path=private_val,
+        test_size=val_test_size,
+        random_state=0,  # Use same random state for consistency
+    )

mlebench/competitions/mcm_2024_c_test/grade.py

@@ -0,0 +1,107 @@
+import os
+from pathlib import Path
+from typing import Any
+import pandas as pd
+
+# This is a generic LLM-based grader for open-ended tasks.
+# It reads 'rubric.md' from the task directory and evaluates the submission.
+
+try:
+    from dsat.services.llm import LLMService
+    from dsat.config import LLMConfig
+except ImportError:
+    # Fallback for when running outside of dsat package context
+    import sys
+    sys.path.append(str(Path(__file__).resolve().parent.parent.parent.parent))
+    from dsat.services.llm import LLMService
+    from dsat.config import LLMConfig
+
+class Report:
+    def __init__(self, score, feedback):
+        self.score = score
+        self.feedback = feedback
+        # Standard fields expected by the framework
+        self.is_lower_better = False
+        self.submission_exists = True
+        self.valid_submission = True
+        self.gold_medal = score >= 0.9
+        self.silver_medal = score >= 0.7
+        self.bronze_medal = score >= 0.5
+        self.above_median = score >= 0.5
+        self.submission_path = ""
+        self.competition_id = "open_ended_task"
+
+def grade(submission_path: Path, competition: Any) -> Report:
+    """
+    Grades the submission using an LLM Judge against rubric.md.
+    """
+    # 1. Load the Rubric
+    task_dir = competition.raw_dir.parent
+    rubric_path = task_dir / "rubric.md"
+    
+    if not rubric_path.exists():
+        # Fallback if no rubric exists
+        print(f"Warning: Rubric not found at {rubric_path}. Returning default score.")
+        return Report(0.5, "No grading rubric defined.")
+        
+    rubric_content = rubric_path.read_text(encoding="utf-8")
+    
+    # 2. Load the Submission Content (Preview)
+    # Since it's open-ended, the 'submission_path' might be a CSV, code, or just a marker.
+    # We'll try to peek at the output artifacts if possible, or assume the agent's recent work 
+    # is what we are grading. Ideally, AIDE produces a submission file.
+    
+    submission_content = "No submission content readable."
+    if submission_path.exists():
+        try:
+            if submission_path.suffix == '.csv':
+                df = pd.read_csv(submission_path)
+                submission_content = f"CSV Submission Preview:\n{df.head().to_markdown()}"
+            else:
+                submission_content = submission_path.read_text(encoding="utf-8")[:2000]
+        except Exception as e:
+            submission_content = f"Error reading submission: {e}"
+
+    # 3. Setup LLM for Judging
+    # Note: In a real run, we might want to inject the API key securely.
+    # Here we assume environment variables are set (which they are in DSATRunner).
+    try:
+        api_key = os.getenv("API_KEY", "EMPTY")
+        base_url = os.getenv("API_BASE", "https://api.openai.com/v1")
+        model = os.getenv("LLM_MODEL", "gpt-4o")
+        
+        llm = LLMService(LLMConfig(api_key=api_key, api_base=base_url, model=model))
+        
+        prompt = f"""You are an impartial Judge. Evaluate the following submission against the provided Rubric.
+
+# RUBRIC
+{rubric_content}
+
+# SUBMISSION CONTENT
+{submission_content}
+
+# INSTRUCTION
+Assess the submission. 
+Output ONLY a float number between 0.0 and 1.0 on the first line.
+On subsequent lines, provide brief feedback.
+"""
+        # Synchronous call wrapper or direct call if possible. 
+        # Since grade() is synchronous in standard mlebench, we need a way to run async code.
+        import asyncio
+        response = asyncio.run(llm.achat([{"role": "user", "content": prompt}]))
+        
+        lines = response.strip().split('\n')
+        try:
+            score = float(lines[0].strip())
+        except ValueError:
+            # Fallback if LLM is chatty
+            import re
+            match = re.search(r"(\d+(\.\d+)?)", lines[0])
+            score = float(match.group(1)) if match else 0.5
+            
+        feedback = "\n".join(lines[1:])
+        return Report(score, feedback)
+
+    except Exception as e:
+        print(f"LLM Judging failed: {e}")
+        return Report(0.0, f"Judging failed: {e}")

mlebench/competitions/mcm_2024_c_test/prepare.py

@@ -0,0 +1,2 @@
+def prepare(raw, public, private):
+    pass

mlebench/competitions/ml2021spring-hw2/grade.py

@@ -0,0 +1,11 @@
+import pandas as pd
+from sklearn.metrics import accuracy_score
+
+from mlebench.competitions.utils import prepare_for_accuracy_metric
+
+
+def grade(submission: pd.DataFrame, answers: pd.DataFrame) -> float:
+    accuracy_inputs = prepare_for_accuracy_metric(
+        submission=submission, answers=answers, target_column="ClassId", id_column="Id"
+    )
+    return accuracy_score(y_true=accuracy_inputs["y_true"], y_pred=accuracy_inputs["y_pred"])

mlebench/competitions/ml2021spring-hw2/prepare.py

@@ -0,0 +1,58 @@
+from pathlib import Path
+
+import numpy as np
+import pandas as pd
+from sklearn.model_selection import train_test_split
+
+
+def prepare(raw: Path, public: Path, private: Path):
+    """
+    Create train/test split from old train set, using same train/test proportion
+    """
+    old_train = np.load(raw / "timit_11" / "timit_11" / "train_11.npy")
+    old_train_label = np.load(raw / "timit_11" / "timit_11" / "train_label_11.npy")
+    old_train_idxs = range(len(old_train))
+
+    # Create new splits
+    old_test = np.load(raw / "timit_11" / "timit_11" / "test_11.npy")
+    old_test_prop = len(old_test) / (len(old_train) + len(old_test))  # ~= 0.268
+    new_train_idxs, new_test_idxs = train_test_split(
+        old_train_idxs, test_size=old_test_prop, random_state=0
+    )
+
+    new_train = old_train[new_train_idxs]
+    new_train_label = old_train_label[new_train_idxs]
+    new_test = old_train[new_test_idxs]
+    new_test_label = old_train_label[new_test_idxs]
+
+    answers_df = pd.DataFrame({"Id": range(len(new_test)), "ClassId": new_test_label})
+
+    # Create sample submission
+    sample_submission = answers_df.copy()
+    sample_submission["ClassId"] = 0
+
+    # Save files
+    (public / "timit_11" / "timit_11").mkdir(parents=True, exist_ok=True)
+
+    np.save(public / "timit_11" / "timit_11" / "train_11.npy", new_train)
+    np.save(public / "timit_11" / "timit_11" / "train_label_11.npy", new_train_label)
+    np.save(public / "timit_11" / "timit_11" / "test_11.npy", new_test)
+    sample_submission.to_csv(public / "sampleSubmission.csv", index=False)
+    answers_df.to_csv(private / "answers.csv", index=False)
+
+    # Sanity checks
+    assert (
+        public / "timit_11" / "timit_11" / "train_11.npy"
+    ).exists(), "`train_11.npy` doesn't exist!"
+    assert (
+        public / "timit_11" / "timit_11" / "train_label_11.npy"
+    ).exists(), "`train_label_11.npy` doesn't exist!"
+    assert (
+        public / "timit_11" / "timit_11" / "test_11.npy"
+    ).exists(), "`test_11.npy` doesn't exist!"
+    assert (public / "sampleSubmission.csv").exists(), "`sampleSubmission.csv` doesn't exist!"
+    assert (private / "answers.csv").exists(), "`answers.csv` doesn't exist!"
+
+    assert len(new_train) + len(new_test) == len(
+        old_train
+    ), f"Expected {len(old_train)} samples in combined new train and test splits, got {len(new_train) + len(new_test)}!"