pretrain model

scripts/prepare_pretrain_dataset.py

@@ -74,6 +74,16 @@ def batch_iterator(name=None):
             del dataset
             gc.collect()
 
+    # code
+    if name in (None, 'nampdn-ai/tiny-codes'):
+        dataset = load_dataset(name, split='train')
+        
+        for row in dataset:
+            yield row['prompt'] + '\n' + row['response']
+        
+        del dataset
+        gc.collect()
+
     # code
     if name in (None, 'bigcode/the-stack-smol-xs'):
         dataset = (
@@ -105,22 +115,42 @@ def batch_iterator(name=None):
         del dataset
         gc.collect()
 
+    # text + code
+    if name in (None, 'm-a-p/CodeFeedback-Filtered-Instruction'):
+        dataset = load_dataset(name, split='train')
+
+        for row in dataset:
+            yield row['query'] + '\n' + row['answer']
+
+        del dataset
+        gc.collect()
+
     # code
-    if name in (None, 'nampdn-ai/tiny-codes'):
+    if name in (None, 'jtatman/python-code-dataset-500k'):
         dataset = load_dataset(name, split='train')
-        
+
         for row in dataset:
-            yield row['prompt'] + '\n' + row['response']
-        
+            yield row['instruction'] + '\n' + row['output']
+
         del dataset
         gc.collect()
 
-    # text + code
-    if name in (None, 'm-a-p/CodeFeedback-Filtered-Instruction'):
+    # code
+    if name in (None, 'iamtarun/python_code_instructions_18k_alpaca'):
         dataset = load_dataset(name, split='train')
 
         for row in dataset:
-            yield row['query'] + '\n' + row['answer']
+            yield row['instruction'] + '\n' + row['input'] + '\n' + row['output']
+
+        del dataset
+        gc.collect()
+
+    # code
+    if name in (None, 'HuggingFaceH4/CodeAlpaca_20K'):
+        dataset = load_dataset(name, split='train')
+
+        for row in dataset:
+            yield row['prompt'] + '\n' + row['completion']
 
         del dataset
         gc.collect()
@@ -187,9 +217,12 @@ datasets_names = [
     'xu-song/cc100-samples',
     'ontocord/fineweb-permissive-multilingual-2m',
     'nampdn-ai/tiny-textbooks',
-    'bigcode/the-stack-smol-xs',
     'nampdn-ai/tiny-codes',
+    'bigcode/the-stack-smol-xs',
     'm-a-p/CodeFeedback-Filtered-Instruction',
+    'jtatman/python-code-dataset-500k',
+    'iamtarun/python_code_instructions_18k_alpaca',
+    '',
     'gair-prox/open-web-math-pro',
     'ajibawa-2023/Maths-College',
     'microsoft/orca-math-word-problems-200k',