feat(transformers): add ut for generation #1210
Conversation
![gemini-code-assist[bot]](https://avatars.githubusercontent.com/in/956858?s=60&v=4){kind=small}
There was a problem hiding this comment.
Summary of Changes
Hello @wcrzlh, I'm Gemini Code Assist1! 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 PR significantly enhances the transformers generation module by introducing comprehensive unit tests for beam search, candidate generation, logits processing, and stopping criteria. It also refines existing tensor operations to align with MindSpore best practices and adds new test utilities to improve testing robustness and efficiency. A new speculative sampling function is also introduced.
Highlights
- Expanded Unit Test Coverage: Introduced new test suites for BeamSearch, CandidateGenerator, LogitsProcessors, and StoppingCriteria, greatly improving the reliability of the generation module.
- MindSpore Tensor API Alignment: Updated tensor creation and manipulation calls across the generation module to leverage mint.empty, mint.tensor, and ops.full, promoting consistency and potentially better performance with MindSpore.
- Enhanced Test Utilities: Added new decorators (@slow, @is_flaky) and helper functions in testing_utils.py to manage slow and flaky tests, and to configure models/configs for more stable testing.
- Speculative Sampling Implementation: Integrated a new _speculative_sampling function within mindone/transformers/generation/utils.py, laying the groundwork for advanced decoding strategies.
- Simplified Module Imports: Streamlined imports in mindone/transformers/generation/init.py for easier access to generation components.
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
-
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. ↩
There was a problem hiding this comment.
Code Review
This pull request adds a comprehensive suite of unit tests for the generation functionality, which is a great addition for ensuring code quality and correctness. It also includes several fixes to adapt the codebase from PyTorch to MindSpore, such as replacing PyTorch-specific API calls with their MindSpore equivalents. My review focuses on improving code style and correctness in type hints. Overall, this is a valuable contribution.
| from .beam_search import * | ||
| from .candidate_generator import * | ||
| from .logits_process import * | ||
| from .stopping_criteria import * | ||
| from .utils import * |
There was a problem hiding this comment.
Using from .module import * is generally discouraged as it can lead to namespace pollution and makes it difficult to determine which names are part of the public API. Consider using __all__ in the submodules to explicitly declare the public API, or import names explicitly to improve clarity and maintainability.
| def heal_tokens(self, input_ids: ms.tensor, tokenizer: Optional["PreTrainedTokenizerBase"] = None) -> ms.tensor: | ||
| r""" | ||
| Generates sequences of token ids for models with a language modeling head. | ||
| Parameters: | ||
| input_ids (`ms.Tensor`): The sequence used as a prompt for the generation. | ||
| input_ids (`ms.tensor`): The sequence used as a prompt for the generation. | ||
| tokenizer (`PreTrainedTokenizerBase`, *optional*): The tokenizer used to decode the input ids. | ||
| Return: | ||
| `ms.Tensor` where each sequence has its tail token replaced with its appropriate extension. | ||
| `ms.tensor` where each sequence has its tail token replaced with its appropriate extension. | ||
| """ |
There was a problem hiding this comment.
The type hint ms.tensor is incorrect. ms.tensor is a factory function for creating tensors, while ms.Tensor is the type. Please use ms.Tensor for type hints in the function signature and docstrings for correctness and clarity.
| def heal_tokens(self, input_ids: ms.tensor, tokenizer: Optional["PreTrainedTokenizerBase"] = None) -> ms.tensor: | |
| r""" | |
| Generates sequences of token ids for models with a language modeling head. | |
| Parameters: | |
| input_ids (`ms.Tensor`): The sequence used as a prompt for the generation. | |
| input_ids (`ms.tensor`): The sequence used as a prompt for the generation. | |
| tokenizer (`PreTrainedTokenizerBase`, *optional*): The tokenizer used to decode the input ids. | |
| Return: | |
| `ms.Tensor` where each sequence has its tail token replaced with its appropriate extension. | |
| `ms.tensor` where each sequence has its tail token replaced with its appropriate extension. | |
| """ | |
| def heal_tokens(self, input_ids: ms.Tensor, tokenizer: Optional["PreTrainedTokenizerBase"] = None) -> ms.Tensor: | |
| r""" | |
| Generates sequences of token ids for models with a language modeling head. | |
| Parameters: | |
| input_ids (`ms.Tensor`): The sequence used as a prompt for the generation. | |
| tokenizer (`PreTrainedTokenizerBase`, *optional*): The tokenizer used to decode the input ids. | |
| Return: | |
| `ms.Tensor` where each sequence has its tail token replaced with its appropriate extension. | |
| """ |
| """ | ||
| the latter code assumes the input_ids is not empty, | ||
| input_id has to be checked if contains elements | ||
| """ |
There was a problem hiding this comment.
This multiline string is used as a comment. For better readability and to avoid potential runtime effects of docstrings, it's recommended to use standard # comments.
# The following code assumes `input_ids` is not empty,
# so it has to be checked if it contains elements.| """ | ||
| seq_bias key has to be tuple with int so have to use | ||
| tokenizer function to convert str to int | ||
| """ |
There was a problem hiding this comment.
This multiline string is used as a comment. For better readability and to avoid potential runtime effects of docstrings, it's recommended to use standard # comments.
# The `seq_bias` key must be a tuple of integers (token IDs), so we use the
# tokenizer function to convert the string token to an integer ID.
2 participants
What does this PR do?
Fixes # (issue)
Adds # (feature)
Before submitting
What's New. Here are thedocumentation guidelines
Who can review?
Anyone in the community is free to review the PR once the tests have passed. Feel free to tag
members/contributors who may be interested in your PR.
@xxx