update short form script

AkariAsai · Dec 22, 2023 · ed22968 · ed22968
1 parent 2c20ff1
commit ed22968
Showing 1 changed file with 2 additions and 2 deletions.
diff --git a/retrieval_lm/run_short_form.py b/retrieval_lm/run_short_form.py
@@ -50,7 +50,7 @@ def postprocess_answer_option_conditioned(answer):
 
 def call_model_rerank_w_scores_batch(prompt, evidences, model, max_new_tokens=15,
                                      ret_tokens=None, rel_tokens=None, grd_tokens=None, ut_tokens=None,
-                                     use_seqscore=False, threshold=0.5, 
+                                     use_seqscore=False, threshold=0.5,
                                      w_rel=1.0, w_sup=1.0, w_use=0.5, mode="adaptive_retrieval", closed=False):
     results = {}
     if mode != "always_retrieve":
@@ -312,7 +312,7 @@ def main():
     def generate(prompt, evidences, max_new_tokens):
         return call_model_rerank_w_scores_batch(prompt, evidences=evidences, model=model, max_new_tokens=max_new_tokens,
                                                 rel_tokens=rel_tokens, ret_tokens=ret_tokens, grd_tokens=grd_tokens, ut_tokens=ut_tokens,
-                                                threshold=args.threshold, beam_width=args.beam_width, max_depth=args.max_depth, use_seqscore=args.use_seqscore,
+                                                threshold=args.threshold, max_depth=args.max_depth, use_seqscore=args.use_seqscore,
                                                 w_rel=args.w_rel, w_sup=args.w_sup, w_use=args.w_use, mode=args.mode, closed=args.task in ["fever", "arc_c"])
 
     preds = []