optimized setup.sh

setup.sh

@@ -7,7 +7,11 @@ echo "Copying files..."
 cp debug/utils.py .venv/lib/python3.10/site-packages/transformers/generation/utils.py
 cp debug/modeling_llama.py .venv/lib/python3.10/site-packages/transformers/models/llama/modeling_llama.py
 
-export LLM_PATH=meta-llama/Llama-2-7b-hf
+if [ -d "$HOME/Llama-2-7b-hf" ]; then
+    export LLM_PATH="$HOME/Llama-2-7b-hf"
+else
+    export LLM_PATH="meta-llama/Llama-2-7b-hf"
+fi
 echo "LLM_PATH: $LLM_PATH"
 
 if [ -n "$LD_LIBRARY_PATH" ]; then

test_movielens.sh

@@ -2,7 +2,7 @@ python main.py \
 --mode test \
 --router share \
 --gating Dense \
---batch_size 8 \
+--batch_size 4 \
 --accumulate_grad_batches 16 \
 --dataset movielens_data \
 --data_dir data/ref/movielens \

train_movielens.sh

@@ -2,8 +2,8 @@ python main.py \
 --mode train \
 --router share \
 --gating Dense \
---batch_size 1 \
---accumulate_grad_batches 4 \
+--batch_size 4 \
+--accumulate_grad_batches 16 \
 --dataset movielens_data \
 --data_dir data/ref/movielens \
 --cans_num 20 \