google-deepmind / google-deepmind/gemma
Tokenizer Optimization: Global Caching & Auto-Download
- Dominant language
- Python
- Stars
- 5.7k
- Forks
- 1k
- Avg merge
- 10h 33m
- Merged PRs (30d)
- 2
Description
## Description
This PR addresses multiple TODOs related to tokenizer performance and usability. It introduces a global LRU cache for the SentencePiece model to prevent redundant loading and parsing when instantiating multiple `Tokenizer` objects. Additionally, it upgrades the file caching utility to transparently download or copy files from remote locations (like GCS) if they are not present in the local cache.
## Changes
- **`gemma/gm/text/_tokenizer.py`**:
- Refactored `_sp` property to use a standalone, globally cached function `_load_sp_model`.
- Decorated `_load_sp_model` with `@functools.lru_cache` to ensure the underlying C++ model is loaded only once per unique path + custom tokens combination.
- Extracted custom token application logic into `_update_proto_with_custom_tokens` for better modularity.
- **`gemma/gm/utils/_file_cache.py`**:
- Enhanced `maybe_get_from_cache` to handle cache misses by attempting to copy the file from the `remote_file_path`.
- Added directory creation logic to ensure the cache path exists before writing.
- **`gemma/gm/utils/_file_cache_test.py`**:
- Added `test_cache_miss_downloads_file` to verify that a missing cache file triggers a copy operation from the source.
## Impact
- **Performance**: significantly reduces initialization time and memory usage when working with multiple tokenizer instances (e.g., in distributed training, evaluation pipelines, or tests).
- **Usability**: seamless handling of remote model paths without manual pre-downloading steps.
## Verification
- **Unit Tests**: Added coverage in `_file_cache_test.py` ensuring the download/copy logic works as expected.
- **Existing Tests**: Verified `_tokenizer_test.py` Logic remains consistent (integration tests depend on env setup but logic is unit-verified).
## Checklist
- [x] Implemented global cache for Tokenizer.
- [x] Implemented auto-download for file cache.
- [x] Added/Updated tests.
- [x] Linted code.
## #523 PR Raised for the issue. Please lmk if there's any iteration.
Contributor guide
Assessment
This issue has not been assessed yet.