zhangyiss/mlx-examples
1
0
Fork 0
mirror of https://github.com/ml-explore/mlx-examples.git synced 2025-08-30 02:53:41 +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
parent e477060a00
commit 24f40c3b8d
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

@ -87,7 +87,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
@ -97,7 +98,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)