google-deepmind / google-deepmind/gemma

Tokenizer Optimization: Global Caching & Auto-Download

Open
#532 1 comment 0 reactions 0 assignees View on GitHub
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

Open the contributing guide

Assessment

This issue has not been assessed yet.

Get new issues in your inbox

A short digest of beginner-friendly GitHub issues.