pretrain dataset

scripts/prepare_pretrain_dataset.py

@@ -263,7 +263,7 @@ def batch_iterator(name=None):
 
     # math
     if name in (None, 'fblgit/simple-math'):
-        for split in ['train', 'test']:
+        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 row in dataset: