From a4a86ad89884dea4a4056f4e8ece27a21b2a7da5 Mon Sep 17 00:00:00 2001 From: Chime Ogbuji Date: Sun, 3 Nov 2024 20:30:47 -0500 Subject: [PATCH] Fix iteration over HF dataset collection --- llms/mlx_lm/tuner/datasets.py | 5 +++-- 1 file changed, 3 insertions(+), 2 deletions(-) diff --git a/llms/mlx_lm/tuner/datasets.py b/llms/mlx_lm/tuner/datasets.py index 4599772e..cdda7abd 100644 --- a/llms/mlx_lm/tuner/datasets.py +++ b/llms/mlx_lm/tuner/datasets.py @@ -90,7 +90,8 @@ class CompletionsDatasetCollection: self.collection = data def __getitem__(self, idx: int): - item = next(self.collection) + iteration = iter(self.collection) + item = next(iteration) curr_idx = idx @@ -100,7 +101,7 @@ class CompletionsDatasetCollection: return item[curr_idx] else: curr_idx -= len(item) - item = next(self.collection) + item = next(iteration) except StopIteration: raise IndexError(idx)