Saving weights and logs of step 10000

config.json

@@ -21,7 +21,7 @@
   "pad_token_id": 1,
   "position_embedding_type": "absolute",
   "torch_dtype": "bfloat16",
-  "transformers_version": "4.14.0.dev0",
+  "transformers_version": "4.16.0.dev0",
   "type_vocab_size": 1,
   "use_cache": true,
   "vocab_size": 50265

events.out.tfevents.1640883381.t1v-n-e1a08808-w-0.315965.0.v2

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:de870d4511ea6c2936e5dc3234370e389dd77a065fc9c1db21fa0da61126a2f8
+size 40

events.out.tfevents.1640906290.t1v-n-e1a08808-w-0.356422.0.v2

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:5b6066f9af1f00c3dc400093acd04238a3fe4aad6ba2bc8e926969521ba07e3e
+size 1470136

flax_model.msgpack

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:4d70201bc16c7d6ca23e7a2c885f8a3e434db1d78ed3426f592f9904bda35c4b
+size 498796983

run_mlm_flax.py

@@ -622,6 +622,7 @@ if __name__ == "__main__":
 
         # Generate an epoch by shuffling sampling indices from the train dataset
         num_train_samples = len(tokenized_datasets["train"])
+        print(f'Total number of training samples: {num_train_samples}')
         train_samples_idx = jax.random.permutation(input_rng, jnp.arange(num_train_samples))
         train_batch_idx = generate_batch_splits(train_samples_idx, train_batch_size)
 

run_step1.sh

@@ -3,12 +3,13 @@
     --model_type="roberta" \
     --config_name="./" \
     --tokenizer_name="./" \
-    --dataset_name="NbAiLab/nbailab_extended" \
+    --train_file /mnt/disks/flaxdisk/corpus/train_1_4.json \
+    --validation_file /mnt/disks/flaxdisk/corpus/validation.json \
     --cache_dir="/mnt/disks/flaxdisk/cache/" \
     --max_seq_length="128" \
     --weight_decay="0.01" \
-    --per_device_train_batch_size="256" \
-    --per_device_eval_batch_size="256" \
+    --per_device_train_batch_size="200" \
+    --per_device_eval_batch_size="200" \
     --learning_rate="4e-4" \
     --warmup_steps="10000" \
     --overwrite_output_dir \