← Back to API reference

neuralbench.transforms.SklearnSplit

pydantic model neuralbench.transforms.SklearnSplit

Perform train/val/test split using sklearn’s train_test_split.

Parameters:

• split_by (str) – Column name to use for splitting by (e.g., ‘timeline’ or ‘subject’). If set to “_index” and “_index” is not in the events dataframe, the events dataframe will be reset to have a new column with row indices, named “_index”.

• valid_split_ratio (float) – Ratio of the full dataset to use for validation.

• test_split_ratio (float) – Ratio of the full dataset to use for testing.

• valid_random_state (int) – Random state for validation split.

• test_random_state (int) – Random state for test split.

• stratify_by (str | None) – Column name to use for stratified splitting. If None, no stratification is applied.

Fields:

• split_by (str)

• stratify_by (str | None)

• test_random_state (int)

• test_split_ratio (float)

• valid_random_state (int)

• valid_split_ratio (float)

field split_by: str = 'timeline'

field valid_split_ratio: float = 0.2

field test_split_ratio: float = 0.2

field valid_random_state: int = 33

field test_random_state: int = 33

field stratify_by: str | None = None

requirements: tp.ClassVar[tuple[str, ...]] = ()

← Back to API reference