Add batch size parameter to SentenceTransformersReranker

castorini / pygaggle

a gaggle of deep neural architectures for text ranking and question answering, designed for Pyserini

http://pygaggle.ai/

Apache License 2.0

339 stars 99 forks source link

Closed hugoabonizio closed 2 years ago

hugoabonizio commented 2 years ago

With this parameter we can use different batch sizes, as already possible with other rerankers (e.g. MonoT5), example:

reranker = SentenceTransformersReranker(
    pretrained_model_name_or_path='cross-encoder/ms-marco-MiniLM-L-6-v2',
    batch_size=50,
)

rodrigonogueira4 commented 2 years ago

@lintool, could you please CR this PR?