Update train.sh

train.sh

@@ -1,9 +1,12 @@
+export HF_DATASETS_CACHE="/workspace/.cache/huggingface/datasets"
+
 python run_mlm.py \
     --model_name_or_path Bingsu/my_mobilebert_untrained \
-    --dataset_name Bingsu/my-korean-training-corpus-half \
+    --train_file train.txt \
+    --validation_file val.txt \
     --per_device_train_batch_size 16 \
     --per_device_eval_batch_size 16 \
-    --gradient_accumulation_steps 16 \
+    --gradient_accumulation_steps 2 \
     --max_steps 1000000 \
     --lr_scheduler_type cosine_with_restarts \
     --warmup_ratio 0.05 \
@@ -17,6 +20,4 @@ python run_mlm.py \
     --push_to_hub \
     --hub_strategy checkpoint \
     --output_dir mobilebert_ko \
-    --overwrite_output_dir \
-    --gradient_checkpointing \
-    --use_auth_token
+    --overwrite_output_dir