update jax converter

jax2torch.py

@@ -1,8 +1,12 @@
-from transformers import GPT2Config, GPT2LMHeadModel
+from transformers import AutoTokenizer, GPT2LMHeadModel
 
 '''
-This is a script to convert the Jax model to Pytorch model
+This is a script to convert the Jax model and the tokenizer to Pytorch model
 '''
 
 model = GPT2LMHeadModel.from_pretrained(".", from_flax=True)
 model.save_pretrained(".")
+
+tokenizer = AutoTokenizer.from_pretrained(".")
+tokenizer.save_pretrained(".")
+

run_pretraining.sh

@@ -4,6 +4,7 @@ export WANDB_PROJECT="hf-flax-gpt2-indonesian"
 export WANDB_LOG_MODEL="true"
 
 ./run_clm_flax.py \
+    --model_name_or_path="flax_model.msgpack" \
     --output_dir="${MODEL_DIR}" \
     --model_type="gpt2" \
     --config_name="${MODEL_DIR}" \