Fix iteration over HF dataset collection

This commit is contained in:
Chime Ogbuji 2024-11-03 20:30:47 -05:00
parent e477060a00
commit 24f40c3b8d

View File

@ -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)