pretrain dataset

scripts/prepare_pretrain_dataset.py

@@ -263,8 +263,8 @@ def batch_iterator(name=None):
 
     # math
     if name in (None, 'fblgit/simple-math'):
-        for split in ['arithmetic.float3_valid', 'arithmetic.float34_train', 'arithmetic.float2_train', 'arithmetic.float3_train', 'arithmetic.float34_valid', 'arithmetic.float4_train', 'arithmetic.float4_valid', 'arithmetic.float2_valid']:
-            dataset = load_dataset(name, split=split)
+        for split in ['train', 'test']:
+            dataset = load_dataset(name, revision='refs/convert/parquet', split=split)
 
             for row in dataset:
                 yield (