[BREAKING] refactor reranker template

[0russwest0]

No description provided.

[gemini-code-assist]

Summary of Changes

Hello @0russwest0, 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 implements a breaking refactor of the reranker template and its associated data handling mechanisms. The core objective is to improve the flexibility and efficiency of reranker training by standardizing the input data format to accommodate multiple positive and negative examples per query, and by introducing configurable sampling parameters. This overhaul simplifies the underlying logic across the training pipeline, from data preparation to loss computation, contributing to a more robust and scalable framework.

Highlights

• Reranker Data Format Update: The dataset format for reranker training has been updated to a more flexible structure, now using 'messages', 'positive_messages', and 'negative_messages' fields to explicitly support multiple positive and negative examples per query.
• Configurable Data Sampling: New environment variables, 'MAX_POSITIVE_SAMPLES' and 'MAX_NEGATIVE_SAMPLES', have been introduced to allow users to control the maximum number of positive and negative samples extracted per query during training, enhancing data augmentation flexibility.
• Streamlined Reranker Logic: The internal processing for reranker templates, including data preprocessing, encoding, and collation, has been significantly refactored and simplified. This includes a direct calculation of logits for generative reranker loss, removing intermediate probability steps.
• Documentation and Link Updates: Both Chinese and English documentation for reranker best practices have been updated to reflect the new dataset format, field descriptions, and the newly added environment variables. Repository links for training scripts have also been corrected.

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 in your pull request via creating an issue comment (i.e. comment on the pull request page) using either /gemini <command> or @gemini-code-assist <command>. Below is a summary of the supported commands.

Feature	Command	Description
Code Review	/gemini review	Performs a code review for the current pull request in its current state.
Pull Request Summary	/gemini summary	Provides a summary of the current pull request in its current state.
Comment	@gemini-code-assist	Responds in comments when explicitly tagged, both in issue comments and review comments.
Help	/gemini help	Displays a list of available commands.

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 introduces a significant and beneficial refactoring of the reranker template. The changes include a new, more flexible data format, and corresponding updates to data preprocessing, encoding, and loss calculation. This simplifies the implementation, especially for generative and listwise rerankers, and improves code clarity. My review identifies a few minor issues: a couple of incorrect type hints and return values in the Python code, and a discrepancy between the implementation and the documentation regarding data expansion logic. The core logic of the refactoring is sound and these changes are a good improvement.

[0russwest0]

/gemini review

[gemini-code-assist]

Code Review

This pull request introduces a significant and breaking refactoring of the reranker template. The dataset format is now more flexible, supporting multiple positive and negative examples in a structured manner. The implementation of _reranker_encode and _reranker_data_collator has been greatly simplified and improved. The loss calculation for generative rerankers has also been corrected and simplified. The documentation has been updated to reflect these changes, including new advanced features for Qwen3-Reranker. My review focuses on improving the clarity of the documentation regarding batch size calculation, fixing a type hint mismatch, and improving code style by moving local imports to the top of the file.

[gemini-code-assist]

It's a good practice to place all imports at the top of the file for better readability and to avoid potential issues with module resolution. Please move from collections import defaultdict to the top of the file and remove this line.

[gemini-code-assist]

For better code style and maintainability, imports should be at the top of the file. Please move import os and import random to the top. Note that import os is already present at the top, so you only need to add import random and remove these lines.