From 6a5a725f10c87bf1ccd469e0d8e67f048a225cca Mon Sep 17 00:00:00 2001 From: Davide Caroselli Date: Tue, 28 May 2024 18:00:32 +0200 Subject: [PATCH] Fix: ensure correct handling of `val_set_size` as `float` or `int` (#1655) * Fix: ensure correct handling of val_set_size as float or int * chore: lint --------- Co-authored-by: Wing Lian --- src/axolotl/utils/data/sft.py | 4 +++- 1 file changed, 3 insertions(+), 1 deletion(-) diff --git a/src/axolotl/utils/data/sft.py b/src/axolotl/utils/data/sft.py index dbc4172b4..6083e30be 100644 --- a/src/axolotl/utils/data/sft.py +++ b/src/axolotl/utils/data/sft.py @@ -494,7 +494,9 @@ def load_prepare_datasets( test_fingerprint = md5(to_hash_test) dataset = dataset.train_test_split( - test_size=cfg.val_set_size, + test_size=int(cfg.val_set_size) + if cfg.val_set_size == int(cfg.val_set_size) + else cfg.val_set_size, shuffle=False, seed=cfg.seed or 42, train_new_fingerprint=train_fingerprint,