fix retriever training bug

thu-coai · Aug 6, 2023 · 7020e63 · 7020e63
1 parent a263302
commit 7020e63
Show file tree

Hide file tree

Showing 3 changed files with 6 additions and 3 deletions.
diff --git a/README.md b/README.md
@@ -65,7 +65,7 @@ M H ≤ ℏ c 3 8 π G k B T u {\displaystyle M_{\mathrm {H} }\leq {\frac {\hbar
 </details>
 
 
-<details><summary><b>Question Ansering</b></summary>
+<details><summary><b>Question Answering</b></summary>
 
 ```
 ##########  Query  ##########

diff --git a/data_utils/retriever_datasets.py b/data_utils/retriever_datasets.py
@@ -139,7 +139,8 @@ def load_data(self, path):
                 "easy_neg_contexts": easy_neg_contexts,
                 "hard_neg_contexts": hard_neg_contexts,
                 "label": line["label"],
-                "neg_labels": line["neg_labels"]
+                "easy_neg_labels": line["easy_neg_labels"],
+                "hard_neg_labels": line["hard_neg_labels"],
             })
 
         return data  

diff --git a/scripts/retriever/train.sh b/scripts/retriever/train.sh
@@ -10,7 +10,7 @@ MASTER_PORT=${3-2010}
 MODEL_NAME="roberta-base"
 MODEL_DIR="${WORKING_DIR}/checkpoints/${MODEL_NAME}/"
 # data
-DATA_NAME="${4-TRAIN/p1_en1_hn1_s42}"
+DATA_NAME="${4-TRAIN/p1_en1_hn4_s42}"
 DATA_DIR="${WORKING_DIR}/retriever_data/${DATA_NAME}/merge"
 # hp
 BATCH_SIZE=64
@@ -39,6 +39,8 @@ OPTS+=" --gradient-accumulation-steps ${GRAD_ACC}"
 OPTS+=" --eval-batch-size ${EVAL_BATCH_SIZE}"
 OPTS+=" --epochs ${EPOCHS}"
 OPTS+=" --max-length 256"
+OPTS+=" --save-interval -1"
+OPTS+=" --eval-interval -1"
 # runtime
 OPTS+=" --do-train"
 OPTS+=" --save ${SAVE_PATH}"