predict.sh

#!/usr/bin/env bash

tpu_name=$1  # default-dgdw2, jzb

# model parameter
model_type=$2  # 3B, 11B
model_dir=$3  # models/*, ft_models/*
model_dir=gs://neulab-qa/unifiedqa/${model_dir}
step=$4  # 1100500
output=$5

# dataset parameter
mix=$6
split=$7
search_alg=greedy_decode
beam_size=1
temp=1.0
max_decode_length=128

if [[ $model_type == '3B' ]]; then
    tpb=131072
elif [[ $model_type == '11B' ]]; then
    tpb=16384
fi
inp_len=512
tgt_len=128
model_gin_file=t5_${model_type}_operative_config.gin
model_parallelism=8

mkdir -p $(dirname "${output}")

./run_test.py \
    --tpu=${tpu_name} \
    --gcp_project="${PROJECT}" \
    --tpu_zone="${ZONE}" \
    --model_dir="${model_dir}" \
    --gin_file="${model_gin_file}" \
    --t5_tfds_data_dir="${DATA_DIR}" \
    --gin_file="infer.gin" \
    --gin_file="${search_alg}.gin" \
    --gin_param="mesh_eval_dataset_fn.mixture_or_task_name = '${mix}'" \
    --gin_param="mesh_eval_dataset_fn.num_eval_examples = None" \
    --gin_param="run.dataset_split = '${split}'" \
    --gin_param="output_filename = '${output}'" \
    --gin_param="utils.run.sequence_length = {'inputs': ${inp_len}, 'targets': ${tgt_len}}" \
    --gin_param="utils.tpu_mesh_shape.model_parallelism = ${model_parallelism}" \
    --gin_param="utils.tpu_mesh_shape.tpu_topology = '${TPU_SIZE}'" \
    --gin_param="infer_checkpoint_step = ${step}" \
    --gin_param="run.batch_size = ('tokens_per_batch', ${tpb})" \
    --gin_param="Bitransformer.decode.max_decode_length = ${max_decode_length}" \
    --gin_param="Bitransformer.decode.beam_size = ${beam_size}" \
    --gin_param="Bitransformer.decode.temperature = ${temp}"