mirror of
https://github.com/ml-explore/mlx-examples.git
synced 2025-08-29 18:26:37 +08:00
Fix iteration over HF dataset collection
This commit is contained in:
parent
e477060a00
commit
24f40c3b8d
@ -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)
|
||||
|
||||
|
Loading…
Reference in New Issue
Block a user