zhangyiss/mlx-examples
1
0
Fork 0
mirror of https://github.com/ml-explore/mlx-examples.git synced 2025-08-28 15:08:37 +08:00
Code Issues Actions Packages Projects Releases Wiki Activity

Fix iteration over HF dataset collection

Browse Source
This commit is contained in:
Chime Ogbuji 2024-11-03 20:30:47 -05:00 committed by Awni Hannun
parent 78c33e5037
commit a4a86ad898
1 changed files with 3 additions and 2 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

5
llms/mlx_lm/tuner/datasets.py
View File

@ -90,7 +90,8 @@ class CompletionsDatasetCollection:
self.collection = data self.collection = data
def __getitem__(self, idx: int): def __getitem__(self, idx: int):
item = next(self.collection) iteration = iter(self.collection)
item = next(iteration)
curr_idx = idx curr_idx = idx
@ -100,7 +101,7 @@ class CompletionsDatasetCollection:
return item[curr_idx] return item[curr_idx]
else: else:
curr_idx -= len(item) curr_idx -= len(item)
item = next(self.collection) item = next(iteration)
except StopIteration: except StopIteration:
raise IndexError(idx) raise IndexError(idx)