Modernize datasets: type hints and data splits handling

[ojeda-e]

SUMMARY:
Part of #1927. Addresses changes in datasets module -> src/llmcompressor/datasets/utils.py

TEST PLAN:

• make test
• make style
• pytest tests -m smoke

[gemini-code-assist]

Summary of Changes

Hello @ojeda-e, 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 focuses on modernizing the datasets utility module by adopting contemporary Python language features. The primary goal is to enhance code readability and maintainability through updated type hinting syntax and more expressive control flow for data split management. These changes align the codebase with newer Python standards, making it easier to understand and less prone to type-related errors.

Highlights

• Type Hint Modernization: Updated type hints across several function signatures to use the more concise X | None for optional types and built-in generics like list[T] and dict[K, V] instead of Optional[X], List[T], and Dict[K, V] from the typing module, aligning with Python 3.9+ syntax.
• Enhanced Data Split Handling: The logic for processing dataset splits in get_processed_dataset has been refactored to utilize Python 3.10's match statement, providing a cleaner and more readable way to handle different input types for the splits argument.
• Improved Error Handling: A ValueError is now explicitly raised in get_processed_dataset if an unsupported type is provided for the splits argument, making the function more robust against invalid inputs.

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

The pull request modernizes the datasets module by adding type hints and improving data splits handling. The changes include using | None for optional types, replacing Optional[Dict[str, Dataset]] with dict[str, Dataset] | None, and using a match statement for handling different types of splits. The code also includes renaming a variable for clarity.

[github-actions]

👋 Hi! Thank you for contributing to llm-compressor. Please add the ready label when the PR is ready for review.

Note: This is required to complete the testing suite, please only add the label once the PR is code complete and local testing has been performed.

[fynnsu]

Looks good! Thank you

[dsikka]

Do you mind addressing some of the test failures?

[ojeda-e]

thanks @dsikka, thanks for catching that. I added a fix for it that hopefully will address the failed tests.

[ojeda-e]

@dsikka commit b8f1c53 fixed the transformers and finetune tests. Seems like what is failing is test_kv_cache_quantization[config0] when running make test in the CI, but that passes locally for me. I also ran pytest -vv tests/llmcompressor/modifiers and it passes too.

I think this error is unrelated to the changes in this PR. Can you offer any pointers here?

[brian-dellabetta]

re-running to see if ci/cd errors persist

[ojeda-e]

thanks @brian-dellabetta all test passing now.