first

flax_model.msgpack

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

run.sh

@@ -7,17 +7,16 @@ python run_mlm_flax_stream.py \
     --weight_decay="0.01" \
     --per_device_train_batch_size="62" \
     --per_device_eval_batch_size="16" \
-    --learning_rate="4e-4" \
-    --warmup_steps="1000" \
+    --learning_rate="3e-4" \
+    --warmup_steps="25000" \
     --overwrite_output_dir \
-    --num_train_steps="10000" \
+    --num_train_steps="250000" \
     --adam_beta1="0.9" \
     --adam_beta2="0.98" \
-    --logging_steps="50" \
-    --save_steps="50" \
-    --eval_steps="50" \
+    --logging_steps="1000" \
+    --save_steps="1000" \
+    --eval_steps="1000" \
     --output_dir="./" \
     --dtype="bfloat16" \
-    --push_to_hub_organization="NbAiLab" \
-    --push_to_hub_model_id="nb-roberta-base-scandi" \
+    --hub_model_id="NbAiLab/nb-roberta-base-scandi" \
     --push_to_hub

run_mlm_flax_stream.py

@@ -42,12 +42,14 @@ import optax
 from flax import jax_utils, traverse_util
 from flax.training import train_state
 from flax.training.common_utils import get_metrics, onehot, shard
+from huggingface_hub import Repository
 from transformers import (
     CONFIG_MAPPING,
     FLAX_MODEL_FOR_MASKED_LM_MAPPING,
     AutoConfig,
     AutoTokenizer,
     FlaxAutoModelForMaskedLM,
+    AutoModelForMaskedLM,
     HfArgumentParser,
     PreTrainedTokenizerBase,
     TensorType,
@@ -650,18 +652,18 @@ if __name__ == "__main__":
                     model.save_pretrained(
                         training_args.output_dir,
                         params=params,
-                        push_to_hub=training_args.push_to_hub,
-                        commit_message=f"Saving weights and logs of step {step+1}",
                     )
-                    print(f"Saving weights and logs of step {step+1}. \nThe result is saved to {training_args.output_folder} by worker {jax.process_index()}.")
+
+                    repo = Repository(local_dir=training_args.output_dir)
+                    repo.push_to_hub(commit_message=f"Saving weights and logs of step {step+1}", blocking=False)                    
+		
                 except:
-                    breakpoint()
                     model.save_pretrained(
                         training_args.output_dir,
                         params=params
                     )
                     print("Problems pushing this to the hub. The bug should be fixed.")
+                    breakpoint()
             
         # update tqdm bar
         steps.update(1)
-