Support reranker inference

[tastelikefeet]

PR type

• Bug Fix
• New Feature
• Document Updates
• More Models or Datasets Support

PR information

Write the detail information belongs to this PR.

Experiment results

Paste your experiment result here(if needed).

[gemini-code-assist]

Summary of Changes

Hello @tastelikefeet, I'm Gemini Code Assist¹! I'm currently reviewing this pull request and will post my feedback shortly. In the meantime, here's a summary to help you and other reviewers quickly get up to speed!

This pull request significantly enhances the system's capability to perform reranker inference by integrating support across both PyTorch and vLLM inference engines. It introduces a new argument for controlling activation post-scoring, adds specific model and template support for BGE rerankers, and refines data handling for reranker tasks, thereby broadening the range of supported models and improving flexibility for sequence classification and scoring applications.

Highlights

• Reranker Activation Control: Introduced a new command-line argument reranker_use_activation to control whether a sigmoid activation function is applied after the reranker score calculation during inference. This provides flexibility for different reranker models.
• PyTorch Engine Reranker Support: The PyTorch inference engine (PtEngine) now fully supports reranker models, including handling reranker and generative_reranker task types and applying the specified activation.
• vLLM Engine Reranker Integration: Extended the vLLM inference engine (VllmEngine) to support reranker models. This includes passing relevant task types ('classify', 'score') and activation parameters to vLLM's PoolingParams for efficient reranker inference.
• BGE Reranker Model and Template: Added bge_reranker as a new LLMModelType and a corresponding template, enabling direct support for BAAI's BGE reranker models within the system.
• Qwen Reranker vLLM Overrides: Updated the Qwen template to provide specific hf_overrides for vLLM when using Qwen reranker models, ensuring correct architecture and classifier token handling.
• Unified Reranker Data Handling: Modified the _reranker_encode and _reranker_data_collator methods in the base template to differentiate between training and inference modes, optimizing data preparation for reranker tasks.

Using Gemini Code Assist

The full guide for Gemini Code Assist can be found on our documentation page, here are some quick tips.

Invoking Gemini

You can request assistance from Gemini at any point by creating a comment using either /gemini <command> or @gemini-code-assist <command>. Below is a summary of the supported commands on the current page.

Feature	Command	Description

Customization

To customize Gemini Code Assist for GitHub experience, repository maintainers can create a configuration file and/or provide a custom code review style guide (such as PEP-8 for Python) by creating and adding files to a .gemini/ folder in the base of the repository. Detailed instructions can be found here.

Limitations & Feedback

Gemini Code Assist may make mistakes. Please leave feedback on any instances where its feedback is incorrect or counter productive. You can react with 👍 and 👎 on @gemini-code-assist comments. If you're interested in giving your feedback about your experience with Gemini Code Assist for Github and other Google products, sign up here.

You can also get AI-powered code generation, chat, as well as code reviews directly in the IDE at no cost with the Gemini Code Assist IDE Extension.

Footnotes

1. Review the Privacy Notices, Generative AI Prohibited Use Policy, Terms of Service, and learn how to configure Gemini Code Assist in GitHub here. Gemini can make mistakes, so double check it and use code with caution. ↩

[gemini-code-assist]

Code Review

This pull request adds support for reranker inference to both the PyTorch and vLLM backends. This includes new command-line arguments, model types (bge_reranker), and templates. The changes are quite extensive and touch upon inference engines, model registration, and data processing logic.

My review focuses on a few key areas:

• Configuration Management: I've pointed out several places where os.environ.get is used for configuration. This is not a good practice as it relies on implicit global state. I've recommended moving these to more explicit configuration mechanisms.
• Correctness: I found a potential batch processing bug in the PyTorch engine for generative_reranker models and suggested a fix.
• Maintainability: I've suggested refactoring some repetitive code blocks in the vLLM engine and argument parsing to improve code clarity and reduce duplication.

Overall, this is a solid contribution that significantly expands the framework's capabilities. The suggested changes should improve the robustness and maintainability of the new functionality.

[Jintao-Huang]

\gemini review