train model

scripts/train_model.py

@@ -16,16 +16,16 @@ from transformers import (
 
 
 def _batch_iterator():
-    # code
-    dataset = load_dataset('bigcode/programming-languages-keywords', split='train')
+    ## code
+    # dataset = load_dataset('bigcode/programming-languages-keywords', split='train')
     
-    for row in dataset:
-        for n in row['keywords']:
-            yield n
+    # for row in dataset:
+    #     for n in row['keywords']:
+    #         yield n
     
-    del dataset
-    gc.collect()
-    return
+    # del dataset
+    # gc.collect()
+    # return
 
     # code
     dataset = (
@@ -187,7 +187,14 @@ def batch_iterator():
     for text in _batch_iterator():
         for i in range(0, len(text), 2048):
             chunk = text[i:i + 2048]
-            yield {'text': chunk}
+            tokenized = tokenize_function(chunk)
+            yield tokenized
+
+
+def tokenize_function(text):
+    outputs = tokenizer(text, truncation=True, padding='max_length', max_length=2048)
+    outputs['labels'] = outputs['input_ids'].copy()
+    return outputs
 
 
 tokenizer = AutoTokenizer.from_pretrained('../')
@@ -241,7 +248,6 @@ training_args = TrainingArguments(
     evaluation_strategy='no',
     save_strategy='epoch',
     torch_compile=True,
-    remove_unused_columns=False,
 )
 print(training_args)