create_splits.py

import pandas as pd
from sklearn.model_selection import train_test_split

# Load the metadata
df = pd.read_csv("music_metadata.csv")

if df.empty:
    print("music_metadata.csv is empty. Exiting.")
    exit()

# Create integer labels for styles
df['label'] = df['style'].astype('category').cat.codes
num_labels = len(df['label'].unique())

print(f"Total songs: {len(df)}")
print(f"Total unique styles (labels): {num_labels}")

# Create label-to-ID maps for the model
label2id = dict(zip(df['style'], df['label']))
id2label = dict(zip(df['label'], df['style']))

print("\n--- Label Mappings ---")
print(label2id)
print("----------------------\n")

# Split the data (e.g., 90% train, 10% validation)
# stratify=df['label'] ensures both sets have a similar distribution of styles
train_df, val_df = train_test_split(
    df, 
    test_size=0.1, 
    random_state=42, 
    stratify=df['label']
)

# Save the splits
train_df.to_csv("train.csv", index=False)
val_df.to_csv("val.csv", index=False)

print(f"Training set size: {len(train_df)}")
print(f"Validation set size: {len(val_df)}")