Skip to content

[FEAT] : Introduce FP16 inference optimization and batching parameters to Reranker #813

Description

@suhaniiz

Is your feature request related to a problem? Please describe.

Yes. When processing a large volume of retrieved documents (e.g., 100+ chunks), passing the entire payload to model.predict(pairs) all at once risks exhausting GPU VRAM, leading to Out-Of-Memory (OOM) crashes. Furthermore, running full float32 precision inference on CUDA devices is slower and consumes twice as much memory as necessary for the BAAI/bge-reranker-v2-m3 model.

Describe the solution you'd like

I would love to introduce two performance optimizations to the Reranker class:

Auto-Precision (float16 / bfloat16): Detect if a CUDA device is available and automatically initialize the CrossEncoder with torch_dtype=torch.float16 to drastically accelerate inference speeds and slash memory footprints.

Batching Control: Expose a batch_size parameter (defaulting to 32) in the rerank() signature and forward it directly to SentenceTransformers' predict() function to cleanly throttle data chunks and ensure predictable memory consumption.

Describe alternatives you've considered

No response

Additional Context

No response

GSSoC '26

  • Yes, I am participating in GirlScript Summer of Code and would like to build this.

Metadata

Metadata

Assignees

Labels

enhancementNew feature or improvementgssocGirlScript Summer of Code 2026 issue/PR

Projects

No projects

Milestone

No milestone

Relationships

None yet

Development

No branches or pull requests

Issue actions