reorder "import jax" to avoid the stuck on importing it, change the dataset

run_clm_flax.py

@@ -30,11 +30,11 @@ from dataclasses import dataclass, field
 from pathlib import Path
 from typing import Callable, Optional
 
+import jax
 import datasets
 from datasets import Dataset, load_dataset
 from tqdm import tqdm
 
-import jax
 import jax.numpy as jnp
 import optax
 import transformers

run_finetuning.sh

@@ -5,18 +5,18 @@ export WANDB_LOG_MODEL="true"
 
 ./run_clm_flax.py \
     --model_name_or_path="./flax_model.msgpack" \
-    --output_dir="${MODEL_DIR}/finetuning2" \
+    --output_dir="${MODEL_DIR}/finetuning" \
     --model_type="gpt2" \
     --config_name="${MODEL_DIR}" \
     --tokenizer_name="${MODEL_DIR}" \
     --dataset_name="./text_collection" \
     --dataset_config_name="text_collection" \
-    --dataset_data_dir="/dataset/fiction/story_all" \
+    --dataset_data_dir="/media/storage/datasets/storial/books_txt" \
     --do_train --do_eval \
     --block_size="512" \
     --per_device_train_batch_size="8" \
     --per_device_eval_batch_size="8" \
-    --learning_rate="0.0000001" --warmup_steps="1000" \
+    --learning_rate="0.00005" --warmup_steps="1000" \
     --adam_beta1="0.9" --adam_beta2="0.98" --weight_decay="0.01" \
     --overwrite_output_dir \
     --num_train_epochs="20" \