test2

run_mlm_flax.py

@@ -126,6 +126,12 @@ class DataTrainingArguments:
     overwrite_cache: bool = field(
         default=False, metadata={"help": "Overwrite the cached training and evaluation sets"}
     )
+    static_learning_rate: bool = field(
+        default=False, metadata={"help": "Use a non decaying learning rate"}
+    )
+    auth_token: bool = field(
+        default=False, metadata={"help": "Use authorisation token"}
+    )
     validation_split_percentage: Optional[int] = field(
         default=5,
         metadata={
@@ -327,12 +333,14 @@ if __name__ == "__main__":
                 data_args.dataset_config_name,
                 split=f"train[:{data_args.validation_split_percentage}%]",
                 cache_dir=model_args.cache_dir,
+                use_auth_token=data_args.auth_token,
             )
             datasets["train"] = load_dataset(
                 data_args.dataset_name,
                 data_args.dataset_config_name,
                 split=f"train[{data_args.validation_split_percentage}%:]",
                 cache_dir=model_args.cache_dir,
+                use_auth_token=data_args.auth_token,
             )
     else:
         data_files = {}
@@ -481,7 +489,7 @@ if __name__ == "__main__":
 
     if model_args.model_name_or_path:
         model = FlaxAutoModelForMaskedLM.from_pretrained(
-            model_args.model_name_or_path, config=config, seed=training_args.seed, dtype=getattr(jnp, model_args.dtype)
+        model_args.model_name_or_path, config=config, seed=training_args.seed, dtype=getattr(jnp, model_args.dtype)
         )
     else:
         model = FlaxAutoModelForMaskedLM.from_config(
@@ -499,9 +507,15 @@ if __name__ == "__main__":
     warmup_fn = optax.linear_schedule(
         init_value=0.0, end_value=training_args.learning_rate, transition_steps=training_args.warmup_steps
     )
+
+    if data_argsdata_argtatic_learning_rate:
+        end_lr_value = training_args.learning_rate
+    else:
+        end_lr_value = 0
+
     decay_fn = optax.linear_schedule(
         init_value=training_args.learning_rate,
-        end_value=0,
+        end_value=end_lr_value,
         transition_steps=num_train_steps - training_args.warmup_steps,
     )
     linear_decay_lr_schedule_fn = optax.join_schedules(

run_step1.sh

@@ -10,7 +10,6 @@
     --per_device_train_batch_size="256" \
     --per_device_eval_batch_size="256" \
     --learning_rate="2e-4" \
-    --end_learning_rate="2e-4" \
     --warmup_steps="5000" \
     --overwrite_output_dir \
     --num_train_epochs="1000" \
@@ -20,5 +19,6 @@
     --save_steps="5000" \
     --eval_steps="5000" \
     --preprocessing_num_workers="64" \
-    --use_auth_token
+    --use_auth_token="True" \
+    --static_learning_rale="True" \
     --push_to_hub