udpated the model and script to load local data

pytorch_model.bin

@@ -1,3 +1,3 @@
 version https://git-lfs.github.com/spec/v1
-oid sha256:83e1f5e0a435dafa8cc6d555fd5441209ec15ff7c93cffc1079502ffa0b84b93
+oid sha256:f67e392707d4b269ea616f717bb7451e79d8cc0235449e990209b12bb74aad45
 size 1444576537

run_clm_flax.py

@@ -112,6 +112,9 @@ class DataTrainingArguments:
     dataset_config_name: Optional[str] = field(
         default=None, metadata={"help": "The configuration name of the dataset to use (via the datasets library)."}
     )
+    dataset_data_dir: Optional[str] = field(
+        default=None, metadata={"help": "The name of the data directory."}
+    )
     train_file: Optional[str] = field(default=None, metadata={"help": "The input training data file (a text file)."})
     validation_file: Optional[str] = field(
         default=None,
@@ -296,19 +299,21 @@ def main():
     if data_args.dataset_name is not None:
         # Downloading and loading a dataset from the hub.
         dataset = load_dataset(
-            data_args.dataset_name, data_args.dataset_config_name, cache_dir=model_args.cache_dir, keep_in_memory=False
+            data_args.dataset_name, data_args.dataset_config_name, data_dir=data_args.dataset_data_dir, cache_dir=model_args.cache_dir, keep_in_memory=False
         )
 
         if "validation" not in dataset.keys():
             dataset["validation"] = load_dataset(
                 data_args.dataset_name,
                 data_args.dataset_config_name,
+                data_dir=data_args.dataset_data_dir,
                 split=f"train[:{data_args.validation_split_percentage}%]",
                 cache_dir=model_args.cache_dir,
             )
             dataset["train"] = load_dataset(
                 data_args.dataset_name,
                 data_args.dataset_config_name,
+                data_dir=data_args.dataset_data_dir,
                 split=f"train[{data_args.validation_split_percentage}%:]",
                 cache_dir=model_args.cache_dir,
             )

run_pretraining.sh

@@ -9,8 +9,9 @@ export WANDB_LOG_MODEL="true"
     --model_type="gpt2" \
     --config_name="${MODEL_DIR}" \
     --tokenizer_name="${MODEL_DIR}" \
-    --dataset_name="oscar" \
-    --dataset_config_name="unshuffled_deduplicated_id" \
+    --dataset_name="./datasets/id_collection" \
+    --dataset_config_name="id_collection" \
+    --dataset_data_dir="/data/collection" \
     --do_train --do_eval \
     --block_size="512" \
     --per_device_train_batch_size="24" \