add n-gram support

Guest400123064 · Apr 25, 2024 · 22804ce · 22804ce
1 parent 0ff99e2
commit 22804ce
Showing 1 changed file with 8 additions and 0 deletions.
diff --git a/scripts/benchmark_beir.py b/scripts/benchmark_beir.py
@@ -146,6 +146,7 @@ def evaluate_retriever(args: argparse.Namespace) -> None:
             b=args.bm25_b,
             delta=args.bm25_delta,
             sp_file=args.sp_file,
+            n_grams=args.n_grams,
         )
         model = BEIRWrapper(store)
         retriever = EvaluateRetrieval(model)
@@ -221,6 +222,13 @@ def get_args() -> argparse.Namespace:
         required=False,
         help="Path to the SentencePiece model file; default to None (LLaMA2)",
     )
+    parser.add_argument(
+        "--n-grams",
+        type=int,
+        default=1,
+        required=False,
+        help="The n-gram size up to n for tokenizations; default to 1",
+    )
     parser.add_argument(
         "--split",
         type=str,