updates

run.sh

@@ -6,8 +6,8 @@ python run_mlm_flax_stream.py \
     --dataset_name="NbAiLab/scandinavian" \
     --max_seq_length="512" \
     --weight_decay="0.01" \
-    --per_device_train_batch_size="62" \
-    --per_device_eval_batch_size="62" \
+    --per_device_train_batch_size="12" \
+    --per_device_eval_batch_size="12" \
     --learning_rate="1e-4" \
     --warmup_steps="10000" \
     --overwrite_output_dir \

run_mlm_flax_stream.py

@@ -395,11 +395,11 @@ if __name__ == "__main__":
 
     if model_args.tokenizer_name:
         tokenizer = AutoTokenizer.from_pretrained(
-            model_args.tokenizer_name, cache_dir=model_args.cache_dir, use_fast=model_args.use_fast_tokenizer
+            model_args.tokenizer_name, cache_dir=model_args.cache_dir, use_fast=model_args.use_fast_tokenizer,model_max_length=512
         )
     elif model_args.model_name_or_path:
         tokenizer = AutoTokenizer.from_pretrained(
-            model_args.model_name_or_path, cache_dir=model_args.cache_dir, use_fast=model_args.use_fast_tokenizer
+            model_args.model_name_or_path, cache_dir=model_args.cache_dir, use_fast=model_args.use_fast_tokenizer,model_max_length=512
         )
     else:
         raise ValueError(
@@ -451,6 +451,10 @@ if __name__ == "__main__":
     num_epochs = int(training_args.num_train_epochs)
     train_batch_size = int(training_args.per_device_train_batch_size) * jax.device_count()
     eval_batch_size = int(training_args.per_device_eval_batch_size) * jax.device_count()
+    
+    print("***************************")
+    print(f"Train Batch Size: {train_batch_size}")
+    print("***************************")
 
     # define number steps per stream epoch
     num_train_steps = data_args.num_train_steps