pretrain_ict.sh 1.3 KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445
  1. #! /bin/bash
  2. # Runs the "217M" parameter biencoder model for ICT retriever
  3. RANK=0
  4. WORLD_SIZE=1
  5. PRETRAINED_BERT_PATH=<Specify path of pretrained BERT model>
  6. TEXT_DATA_PATH=<Specify path and file prefix of the text data>
  7. TITLE_DATA_PATH=<Specify path and file prefix od the titles>
  8. CHECKPOINT_PATH=<Specify path>
  9. python pretrain_ict.py \
  10. --num-layers 12 \
  11. --hidden-size 768 \
  12. --num-attention-heads 12 \
  13. --tensor-model-parallel-size 1 \
  14. --micro-batch-size 32 \
  15. --seq-length 256 \
  16. --max-position-embeddings 512 \
  17. --train-iters 100000 \
  18. --vocab-file bert-vocab.txt \
  19. --tokenizer-type BertWordPieceLowerCase \
  20. --DDP-impl torch \
  21. --bert-load ${PRETRAINED_BERT_PATH} \
  22. --log-interval 100 \
  23. --eval-interval 1000 \
  24. --eval-iters 10 \
  25. --retriever-report-topk-accuracies 1 5 10 20 100 \
  26. --retriever-score-scaling \
  27. --load $CHECKPOINT_PATH \
  28. --save $CHECKPOINT_PATH \
  29. --data-path ${TEXT_DATA_PATH} \
  30. --titles-data-path ${TITLE_DATA_PATH} \
  31. --lr 0.0001 \
  32. --lr-decay-style linear \
  33. --weight-decay 1e-2 \
  34. --clip-grad 1.0 \
  35. --lr-warmup-fraction 0.01 \
  36. --save-interval 4000 \
  37. --exit-interval 8000 \
  38. --query-in-block-prob 0.1 \
  39. --fp16