Fix some bugs

src/run.sh

@@ -3,24 +3,25 @@
 export LC_ALL=C.UTF-8
 export LANG=C.UTF-8
 
-export OUTPUT_DIR=/to/our/path
+export OUTPUT_DIR=./
 export MODEL_NAME_OR_PATH=t5-base
 export NUM_BEAMS=3
 
-export TRAIN_FILE=/to/../train.csv
-export VALIDATION_FILE=/to/../dev.csv
-export TEST_FILE=/to/../test.csv
+export TRAIN_FILE=/home/ubuntu/code/data/train.csv
+export VALIDATION_FILE=/home/ubuntu/code/data/test.csv
+export TEST_FILE=/home/ubuntu/code/data/test.csv
 export TEXT_COLUMN=inputs
 export TARGET_COLUMN=targets
 export MAX_SOURCE_LENGTH=256
 export MAX_TARGET_LENGTH=1024
-export SOURCE_PREFIX=ingredients
+export SOURCE_PREFIX=items
+export MAX_EVAL_SAMPLES=5000
 
 export PER_DEVICE_TRAIN_BATCH_SIZE=8
 export PER_DEVICE_EVAL_BATCH_SIZE=8
 export GRADIENT_ACCUMULATION_STEPS=2
 export NUM_TRAIN_EPOCHS=5.0
-export LEARNING_RATE=1e-4
+export LEARNING_RATE=5e-4
 export WARMUP_STEPS=5000
 export LOGGING_STEPS=500
 export EVAL_STEPS=2500
@@ -30,7 +31,7 @@ python run_ed_recipe_nlg.py \
     --output_dir="$OUTPUT_DIR"  \
     --train_file="$TRAIN_FILE" \
     --validation_file="$VALIDATION_FILE" \
-    --test_file="$TEST_FILE" \
+    --max_eval_samples=$MAX_EVAL_SAMPLES \
     --text_column="$TEXT_COLUMN" \
     --target_column="$TARGET_COLUMN" \
     --source_prefix="$SOURCE_PREFIX: " \
@@ -53,4 +54,4 @@ python run_ed_recipe_nlg.py \
     --do_eval \
     --overwrite_output_dir \
     --predict_with_generate \
-    --push_to_hub
+    --push_to_hub

src/run_ed_recipe_nlg.py

@@ -345,6 +345,7 @@ def main():
 
     # Set the verbosity to info of the Transformers logger (on main process only):
     logger.info(f"Training/evaluation parameters {training_args}")
+    logger.info(f"List of TPUs {jax.devices()}")
 
     # Get the datasets: you can either provide your own CSV/JSON training and evaluation files (see below)
     # or just provide the name of one of the public datasets available on the hub at https://huggingface.co/datasets/
@@ -374,6 +375,7 @@ def main():
             data_files["test"] = data_args.test_file
             extension = data_args.test_file.split(".")[-1]
 
+        print(data_files)
         dataset = load_dataset(
             extension,
             data_files=data_files,