train model

scripts/train_model.py

@@ -7,7 +7,7 @@ from transformers import AutoConfig
 from transformers import DataCollatorForLanguageModeling
 
 
-x = input('Are you sure? [y/N]')
+x = input('Are you sure? [y/N] ')
 
 if x not in ('y', 'Y', 'yes'):
     sys.exit(0)
@@ -222,8 +222,11 @@ print(model)
 training_args = TrainingArguments(
     output_dir='./results',
     num_train_epochs=3,
-    per_device_train_batch_size=2,  # Adjust based on your GPU memory
-    per_device_eval_batch_size=2,
+    per_device_train_batch_size=4,  # Adjust based on your GPU memory
+    per_device_eval_batch_size=4,
+    optim='adamw_bnb_8bit',
+    gradient_accumulation_steps=4,
+    gradient_checkpointing=True,
     warmup_steps=500,
     weight_decay=0.01,
     logging_dir='./logs',