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)