05_split_dataframe.py

from pathlib import Path
import pandas as pd
from sklearn.model_selection import train_test_split

PROJECT_DIR = Path(__file__).resolve().parent
INPUT_PATH = PROJECT_DIR / "SBAN_clean.parquet"
RANDOM_STATE = 42

df = pd.read_parquet(INPUT_PATH)

train, temp = train_test_split(
    df,
    test_size=0.30,
    stratify=df["dataset_name"],
    random_state=RANDOM_STATE,
)
val, test = train_test_split(
    temp,
    test_size=2 / 3,
    stratify=temp["dataset_name"],
    random_state=RANDOM_STATE,
)

for name, part in [("train", train), ("val", val), ("test", test)]:
    output_path = PROJECT_DIR / f"SBAN_{name}.parquet"
    part = part.reset_index(drop=True)
    part.to_parquet(output_path, index=False)
    print(f"{name}: {len(part):,} satır")
    print(part["dataset_name"].value_counts().sort_index().to_string())
    print()