From 7b83c550abe69712ceda34279c26c3ad08d58b93 Mon Sep 17 00:00:00 2001 From: hoshi-hiyouga Date: Thu, 30 May 2024 00:20:20 +0800 Subject: [PATCH] Update loader.py Former-commit-id: ca5dd7c6c115a359e4b50e93f4ffcc9f2955ec2f --- src/llamafactory/data/loader.py | 1 + 1 file changed, 1 insertion(+) diff --git a/src/llamafactory/data/loader.py b/src/llamafactory/data/loader.py index fa5b12c5..d4a19e27 100644 --- a/src/llamafactory/data/loader.py +++ b/src/llamafactory/data/loader.py @@ -115,6 +115,7 @@ def load_single_dataset( expand_indexes = np.random.choice(len(dataset), target_num) indexes = np.concatenate((indexes, expand_indexes), axis=0) + assert len(indexes) == dataset_attr.num_samples, "Sample num mismatched." dataset = dataset.select(indexes) logger.info("Sampled {} examples from dataset {}.".format(dataset_attr.num_samples, dataset_attr))