Adding base config and organizing configs

configs/base/config.json

@@ -0,0 +1,25 @@
+{
+    "architectures": [
+      "RobertaForMaskedLM"
+    ],
+    "attention_probs_dropout_prob": 0.1,
+    "bos_token_id": 0,
+    "eos_token_id": 2,
+    "gradient_checkpointing": false,
+    "hidden_act": "gelu",
+    "hidden_dropout_prob": 0.1,
+    "hidden_size": 768,
+    "initializer_range": 0.02,
+    "intermediate_size": 3072,
+    "layer_norm_eps": 1e-05,
+    "max_position_embeddings": 514,
+    "model_type": "roberta",
+    "num_attention_heads": 12,
+    "num_hidden_layers": 12,
+    "pad_token_id": 1,
+    "position_embedding_type": "absolute",
+    "transformers_version": "4.9.0.dev0",
+    "type_vocab_size": 1,
+    "use_cache": true,
+    "vocab_size": 50265
+  }

run.sh

@@ -1,9 +1,9 @@
 # From https://arxiv.org/pdf/1907.11692.pdf
-HUB_TOKEN=`cat $HOME/.huggingface/token`
+python -c "import jax; print('TPUs', jax.device_count())"
 ./run_mlm_flax.py \
-    --output_dir="./" \
+    --output_dir="./outputs" \
     --model_type="roberta" \
-    --config_name="./" \
+    --config_name="./configs/large" \
     --tokenizer_name="./" \
     --dataset_name="mc4" \
     --dataset_config_name="es" \
@@ -25,7 +25,4 @@ HUB_TOKEN=`cat $HOME/.huggingface/token`
     --num_train_steps="500000" \
     --eval_steps="10000" \
     --logging_steps="500" \
-    --dtype="bfloat16" \
-    --push_to_hub_model_id="flax-community/bertin-roberta-large-spanish" \
-    --push_to_hub_token="$HUB_TOKEN"
-    --push_to_hub 2>&1 | tee run.log
+    --dtype="bfloat16" 2>&1 | tee run.log

run_stream.sh

@@ -3,7 +3,7 @@ python -c "import jax; print('TPUs', jax.device_count())"
 python ./run_mlm_flax_stream.py \
     --output_dir="./outputs" \
     --model_type="roberta" \
-    --config_name="./config-base.json" \
+    --config_name="./configs/base" \
     --tokenizer_name="./" \
     --dataset_name="./mc4" \
     --dataset_config_name="es" \