add config for training

main.py

@@ -94,15 +94,15 @@ if __name__ == "__main__":
         os.makedirs(cache_processing_dataset_folder)
     num_train_shards = len(glob.glob(os.path.join(train_dataset_root_folder, 'shard_*')))
     num_test_shards = len(glob.glob(os.path.join(test_dataset_root_folder, 'shard_*')))
-    num_epochs = 20
+    num_epochs = 5000
 
     training_args = TrainingArguments(
         output_dir=checkpoint_path,
         # fp16=True,
         group_by_length=True,
-        per_device_train_batch_size=2,
-        per_device_eval_batch_size=2,
-        gradient_accumulation_steps=1,
+        per_device_train_batch_size=16,
+        per_device_eval_batch_size=16,
+        gradient_accumulation_steps=8,
         num_train_epochs=1,  # each epoch per shard data
         logging_steps=1,
         learning_rate=1e-4,
@@ -146,7 +146,7 @@ if __name__ == "__main__":
                                               cache_file_name=os.path.join(cache_processing_dataset_folder,
                                                                            'cache-train-shard-{}.arrow'.format(
                                                                                train_dataset_shard_idx))
-                                              ).shard(1000, 0)  # Remove shard split when train
+                                              )  # .shard(1000, 0)  # Remove shard split when train
         # load test shard subset
         test_dataset = load_prepared_dataset(os.path.join(test_dataset_root_folder,
                                                           'shard_{}'.format(test_dataset_shard_idx)),