fix hf data loading for datasets>=3.1.0

stanford-crfm · Nov 4, 2024 · 7b2241f · 7b2241f
1 parent 94dcda1
commit 7b2241f
Showing 1 changed file with 4 additions and 1 deletion.
diff --git a/src/levanter/data/sharded_datasource.py b/src/levanter/data/sharded_datasource.py
@@ -197,7 +197,10 @@ def open_shard_at_row(self, shard_name: str, row: int) -> Iterator[dict]:
         dataset = self._load_dataset()
         if isinstance(dataset, datasets.IterableDataset) and shard_name != "data":
             # ex_iterable has a key that gets discarded typically
-            shard = map(lambda t: t[1], dataset._ex_iterable.shard_data_sources(int(shard_name), dataset.n_shards))
+            shard = map(
+                lambda t: t[1],
+                dataset._ex_iterable.shard_data_sources(index=int(shard_name), num_shards=dataset.n_shards),
+            )
         else:
             shard = dataset