Update loader.py

2024-05-30 00:20:20 +08:00 · 2024-05-30 00:20:20 +08:00 · ca5dd7c6c1
parent f9a88b89ca
commit ca5dd7c6c1
1 changed files with 1 additions and 0 deletions
--- 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))