Included contriever encoder-class type for proper document encoding (#…

…1912)
castorini · Jun 6, 2024 · d814290 · d814290
1 parent cf029cd
commit d814290
Showing 1 changed file with 2 additions and 1 deletion.
diff --git 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)