added option for use_consecutive_dataloader to speedup IO

[abhinadduri]

this adds a new kwarg to load consecutive cells for faster IO with cells_per_set > 1

requires that the files are sorted by (context, perturbation), e.g., cells with same context and same perturbation group appear consecutively in the file

to achieve this, use the state tx sort utility from the state repository:

[gemini-code-assist]

Summary of Changes

Hello @abhinadduri, 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 introduces a substantial optimization for data loading in the cell-load library. By adding a use_consecutive_loading parameter, the system can now process pre-sorted single-cell data more efficiently, leading to faster I/O. This involves modifications to how data is sampled and how control cells are mapped, ensuring that the benefits of consecutive data organization are fully utilized throughout the data preparation process.

Highlights

• New Feature: Consecutive Data Loading: Introduced a new use_consecutive_loading option across the data loading pipeline to significantly speed up I/O operations for single-cell datasets. This feature is designed for datasets where cells with the same context and perturbation group appear consecutively in the file.
• Optimized Sampler Logic: The PerturbationSampler now includes a _process_subset_consecutive method that forms cell sets from consecutive indices, splitting sentences at code-change boundaries without shuffling, which is crucial for leveraging the sorted data.
• Mapping Strategy Enhancements: Both BatchMappingStrategy and RandomMappingStrategy have been updated to incorporate the use_consecutive_loading flag. This includes new logic for building fixed control mappings and assigning control cells sequentially (using np.tile) rather than randomly when consecutive loading is enabled.
• Version Bump: The project version in pyproject.toml has been updated from 0.8.7 to 0.9.0, indicating a minor release with new features.

🧠 New Feature in Public Preview: You can now enable Memory to help Gemini Code Assist learn from your team's feedback. This makes future code reviews more consistent and personalized to your project's style. Click here to enable Memory in your admin console.

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 a new use_consecutive_loading option to optimize IO by loading consecutive cells, which is a valuable addition for performance. The implementation is well-integrated across the PerturbationDataModule, PerturbationBatchSampler, and both BatchMappingStrategy and RandomMappingStrategy. The logic for handling consecutive loading, including filling partial sentences and sequentially assigning control cells, appears sound and correctly implemented. Backward compatibility for loading previous states is also handled appropriately. Overall, this is a solid and well-thought-out feature.