Чтобы разбить на более чем два класса, такие как поезд, тест и валидация, можно сделать:
probs = np.random.rand(len(df))
training_mask = probs < 0.7
test_mask = (probs>=0.7) & (probs < 0.85)
validatoin_mask = probs >= 0.85
df_training = df[training_mask]
df_test = df[test_mask]
df_validation = df[validatoin_mask]
Это поставит 70% данных на обучение, 15% в тестах и 15 % в валидации.