[train] support embeding/reranker packing & support reranker/embedding cache_dataset

[Jintao-Huang]

No description provided.

[gemini-code-assist]

Summary of Changes

Hello @Jintao-Huang, 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 enhances the data packing mechanism within the training pipeline, specifically targeting improved support for embedding and reranker models. The changes introduce more flexible handling of sequence lengths, allowing for both single and list-based length representations to be correctly processed and aggregated during data packing. This ensures that complex multi-sequence inputs are properly prepared for model training, while also updating the template version to maintain compatibility.

Highlights

• Flexible Length Handling: The data processing utilities have been updated to support more flexible handling of sequence lengths, allowing lengths to be either a single value or a list of values, which are then appropriately summed or retained as a list during packing.
• Template Version Update: The internal template version has been incremented from v4 to v5 to ensure compatibility and prevent caching issues, reflecting the changes in how data is processed and packed.
• Refined Packing Logic: The packing_row function now correctly aggregates lengths from individual rows, summing them if they are provided as a list, which is crucial for supporting embedding and reranker models that may involve multiple sequences per packed example.

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
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 pull request 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 support for packing in embedding and reranker tasks. The changes are well-implemented and consistent. The core idea is to allow the length of a sample to be a list of integers, which is handled correctly during data processing and packing. I have one suggestion to improve the robustness of data validation by replacing an assert statement with a ValueError. Also, there is a small typo in the pull request title ('embeding' should be 'embedding').

[gemini-code-assist]

Using assert for data validation is not recommended, as assertions can be disabled with the -O (optimize) flag in Python, which might lead to unexpected behavior in production. It's better to use an explicit if check and raise a ValueError for robustness.

Suggested change

	assert len(lengths) != 0, f'batched: {batched}'
	if not lengths:
	raise ValueError(f'lengths should not be empty. batched: {batched}')

[Jintao-Huang]

/gemini review

[gemini-code-assist]

Code Review

This pull request adds support for packing and cached datasets for embedding, reranker, and sequence classification tasks. The changes include updates to documentation, new example scripts, and modifications to data processing logic to handle list-based lengths for samples, which is necessary for tasks like reranking. The code has also been refactored to move padding_side and padding_free arguments to a more appropriate location. Overall, the changes are well-implemented and improve the framework's capabilities.

[gemini-code-assist]

The code assumes that inputs.positive or inputs.negative will contain at least one document, and accesses (inputs.positive + inputs.negative)[0] without checking if the list is empty. This could lead to an IndexError if a reranker inference is attempted with a query but no documents. It would be more robust to add a check to ensure the list of documents is not empty before accessing its first element.

Suggested change

	if len(anchor.messages) == 1:
	# Ensure that load_data_args true runs through inference successfully
	assistant_messages = (inputs.positive + inputs.negative)[0].messages
	assert anchor.messages[0]['role'] == 'user' and assistant_messages[0]['role'] == 'assistant'
	anchor.messages = anchor.messages + assistant_messages
	if len(anchor.messages) == 1:
	# Ensure that load_data_args true runs through inference successfully
	docs = inputs.positive + inputs.negative
	if docs:
	assistant_messages = docs[0].messages
	assert anchor.messages[0]['role'] == 'user' and assistant_messages[0]['role'] == 'assistant'
	anchor.messages = anchor.messages + assistant_messages