From 24f40c3b8d3fbf2d8da76d69ce1c853bf8449058 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 c75171e5..1a86fb9f 100644 --- a/llms/mlx_lm/tuner/datasets.py +++ b/llms/mlx_lm/tuner/datasets.py @@ -87,7 +87,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 @@ -97,7 +98,7 @@ class CompletionsDatasetCollection: return item[curr_idx] else: curr_idx -= len(item) - item = next(self.collection) + item = next(iteration) except StopIteration: raise IndexError(idx)