From d814290e846d94ff4d9083afb5da73a491a10a0d Mon Sep 17 00:00:00 2001 From: Ramraj Chandradevan Date: Wed, 5 Jun 2024 20:24:45 -0500 Subject: [PATCH] Included contriever encoder-class type for proper document encoding (#1912) --- pyserini/encode/__main__.py | 3 ++- 1 file changed, 2 insertions(+), 1 deletion(-) diff --git a/pyserini/encode/__main__.py b/pyserini/encode/__main__.py index b294444f6..e7c8b3d78 100644 --- a/pyserini/encode/__main__.py +++ b/pyserini/encode/__main__.py @@ -34,6 +34,7 @@ "cosdpr": CosDprDocumentEncoder, "auto": AutoDocumentEncoder, "clip": ClipDocumentEncoder, + "contriever": AutoDocumentEncoder, } def init_encoder(encoder, encoder_class, device, pooling, l2_norm, prefix, multimodal): @@ -113,7 +114,7 @@ def parse_args(parser, commands): encoder_parser = commands.add_parser('encoder') encoder_parser.add_argument('--encoder', type=str, help='encoder name or path', required=True) encoder_parser.add_argument('--encoder-class', type=str, required=False, default=None, - choices=["dpr", "bpr", "tct_colbert", "ance", "sentence-transformers", "openai-api", "auto"], + choices=["dpr", "bpr", "tct_colbert", "ance", "sentence-transformers", "openai-api", "auto", "contriever"], help='which query encoder class to use. `default` would infer from the args.encoder') encoder_parser.add_argument('--fields', help='fields to encode', nargs='+', default=['text'], required=False) encoder_parser.add_argument('--multimodal', action='store_true', default=False)