shard fix
Browse files
src/axolotl/utils/data.py
CHANGED
@@ -112,7 +112,7 @@ def load_tokenized_prepared_datasets(
|
|
112 |
raise Exception("unhandled dataset load")
|
113 |
# support for using a subset of the data
|
114 |
if d.shards:
|
115 |
-
ds = ds.shuffle(seed=42)["train"].shard(num_shards=
|
116 |
d_type = d.type
|
117 |
d_type_split = d_type.split(":")
|
118 |
d_base_type = d_type_split[0]
|
|
|
112 |
raise Exception("unhandled dataset load")
|
113 |
# support for using a subset of the data
|
114 |
if d.shards:
|
115 |
+
ds = ds.shuffle(seed=42)["train"].shard(num_shards=d.shards, index=0)
|
116 |
d_type = d.type
|
117 |
d_type_split = d_type.split(":")
|
118 |
d_base_type = d_type_split[0]
|