-
Notifications
You must be signed in to change notification settings - Fork 4
/
art-nq-T0-3B.sh
119 lines (103 loc) · 4.23 KB
/
art-nq-T0-3B.sh
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
# CHANGE THIS PATH TO POINT TO DATA / CHECKPOINTS
BASE_DIR="/mnt/disks/project"
BERT_VOCAB_FILE="${BASE_DIR}/bert-vocab/bert-large-uncased-vocab.txt"
DATA_DIR="${BASE_DIR}/data"
TRAIN_DATA="${DATA_DIR}/qas/nq-train.csv"
QA_FILE_DEV="${DATA_DIR}/qas/nq-dev.csv"
QA_FILE_TEST="${DATA_DIR}/qas/nq-test.csv"
EVIDENCE_DATA_PATH="${DATA_DIR}/wikipedia-split/psgs_w100.tsv"
# Please specify the model name from Huggingface models
CROSS_ATTENTION_SCORER="bigscience/T0_3B"
TOPK_DOCUMENTS=32
# SET THIS TO NUMBER OF GPUs
WORLD_SIZE=16
EPOCHS=10
RETRIEVER_CONFIG="base"
PRETRAINED_RETRIEVER_CHECKPOINT="${BASE_DIR}/checkpoints/mss-retriever-${RETRIEVER_CONFIG}"
PRETRAINED_RETRIEVER_EMBEDDING_PATH="${BASE_DIR}/embedding-path/psgs_w100-mss-retriever-full-wikipedia-${RETRIEVER_CONFIG}.pkl"
mkdir -p ${BASE_DIR}"/embedding-path/art-finetuning-embedding"
CHECKPOINT_PATH="${BASE_DIR}/checkpoints/nq-mss-${RETRIEVER_CONFIG}-init"
EMBEDDING_PATH="${BASE_DIR}/embedding-path/art-finetuning-embedding/psgs_w100-retriever-nq-${RETRIEVER_CONFIG}-topk${TOPK_DOCUMENTS}-epochs${EPOCHS}-bsize64-indexer.pkl"
# Copy the path
if [ -f ${EMBEDDING_PATH} ]; then
echo "${EMBEDDING_PATH} exists. Not copying file"
else
echo "Copying file"
cp ${PRETRAINED_RETRIEVER_EMBEDDING_PATH} ${EMBEDDING_PATH}
fi
DISTRIBUTED_ARGS="-m torch.distributed.launch --nproc_per_node ${WORLD_SIZE} --nnodes 1 --node_rank 0 --master_addr localhost --master_port 6000"
function config_base() {
export CONFIG_ARGS="--num-layers 12 \
--hidden-size 768 \
--num-attention-heads 12 \
--kv-channels 64 \
--ffn-hidden-size 3072 \
--model-parallel-size 1"
}
function config_large() {
export CONFIG_ARGS="--num-layers 24 \
--hidden-size 1024 \
--num-attention-heads 16 \
--kv-channels 64 \
--ffn-hidden-size 4096 \
--model-parallel-size 1"
}
if [ ${RETRIEVER_CONFIG} == "base" ]; then
config_base
elif [ ${RETRIEVER_CONFIG} == "large" ]; then
config_large
else
echo "Invalid model configuration"
exit 1
fi
OPTIONS=" \
--train-data ${TRAIN_DATA} \
--qa-file-dev ${QA_FILE_DEV} \
--qa-file-test ${QA_FILE_TEST} \
--evidence-data-path ${EVIDENCE_DATA_PATH} \
--indexed-evidence-bert-tokenized-data-path ${DATA_DIR}/evidence-wikipedia-indexed-mmap/bert/wikipedia-evidence-bert_text_document \
--indexed-title-bert-tokenized-data-path ${DATA_DIR}/evidence-wikipedia-indexed-mmap/bert/wikipedia-evidence-bert_title_document \
--indexed-evidence-t0-tokenized-data-path ${DATA_DIR}/evidence-wikipedia-indexed-mmap/t0/wikipedia-evidence-t0_text_document \
--indexed-title-t0-tokenized-data-path ${DATA_DIR}/evidence-wikipedia-indexed-mmap/t0/wikipedia-evidence-t0_title_document \
--save-interval 500 \
--save ${CHECKPOINT_PATH} \
--load ${CHECKPOINT_PATH} \
--pretrained-dualencoder-load ${PRETRAINED_RETRIEVER_CHECKPOINT} \
--embedding-path ${EMBEDDING_PATH} \
--vocab-file ${BERT_VOCAB_FILE} \
--log-interval 20 \
--eval-interval 500 \
--weight-decay 1.0e-1 \
--seq-length 512 \
--seq-length-retriever 256 \
--max-position-embeddings 512 \
--fp16 \
--num-workers 2 \
--distributed-backend nccl \
--checkpoint-activations \
--task ZERO-SHOT-RETRIEVER \
--tokenizer-type BertWordPieceLowerCase \
--epochs ${EPOCHS} \
--sample-rate 1.0 \
--batch-size 4 \
--eval-batch-size 1 \
--lr 2e-5 \
--warmup 0.01 \
--DDP-impl local \
--lr-decay-style linear \
--max-training-rank ${WORLD_SIZE} \
--topk-retrievals ${TOPK_DOCUMENTS} \
--report-topk-accuracies 1 5 20 50 100 \
--art-training \
--retriever-score-scaling \
--update-retriever \
--allow-trivial-doc \
--shard-size 16 \
--hf-model-name ${CROSS_ATTENTION_SCORER} \
--initialize-t0-model-tokenizer-evidence \
--t0-model-in-bf16 \
--index-reload-interval 500 \
--compute-fresh-evidence-embeddings "
COMMAND="WORLD_SIZE=${WORLD_SIZE} python ${DISTRIBUTED_ARGS} tasks/run.py ${OPTIONS} ${CONFIG_ARGS}"
eval "${COMMAND}"
exit