From b6302b1fcc19a2134cde5e566f1392fff7a7be48 Mon Sep 17 00:00:00 2001 From: Aarav Mittal Date: Mon, 3 Aug 2026 18:45:04 -0700 Subject: [PATCH] fix: split local-tokenizer extras for SentencePiece-only installs Keep google-genai[local-tokenizer] limited to sentencepiece/protobuf so Gemini 2.x models do not pull the Hugging Face PyTorch stack. Move torch/transformers/pillow into local-tokenizer-gemma4 and point Gemma 4 import errors at that extra. Fixes #2728. --- README.md | 25 +++++++++++++++++-- google/genai/_local_tokenizer_loader.py | 13 ++++++---- google/genai/local_tokenizer.py | 5 ++++ .../test_local_tokenizer_loader.py | 18 +++++++++++++ pyproject.toml | 7 ++++++ 5 files changed, 61 insertions(+), 7 deletions(-) diff --git a/README.md b/README.md index 3ea066513..6af466fab 100644 --- a/README.md +++ b/README.md @@ -54,6 +54,23 @@ pip install google-genai uv pip install google-genai ``` +### Optional extras + +- Local text token counting for SentencePiece-backed models (for example + `gemini-2.5-flash`): + + ```sh + pip install "google-genai[local-tokenizer]" + ``` + +- Local text token counting for Gemma 4–backed models (for example + `gemini-3.5-flash`). This pulls the Hugging Face / PyTorch stack and is much + larger: + + ```sh + pip install "google-genai[local-tokenizer-gemma4]" + ``` + ## Imports ```python @@ -1212,10 +1229,14 @@ print(response) #### Local Count Tokens +Install `google-genai[local-tokenizer]` for SentencePiece-backed models such as +`gemini-2.5-flash`. For Gemma 4–backed models such as `gemini-3.5-flash`, install +`google-genai[local-tokenizer-gemma4]` instead. + ```python from google.genai import local_tokenizer -tokenizer = local_tokenizer.LocalTokenizer(model_name='gemini-3.5-flash') +tokenizer = local_tokenizer.LocalTokenizer(model_name='gemini-2.5-flash') result = tokenizer.count_tokens("What is your name?") ``` @@ -1224,7 +1245,7 @@ result = tokenizer.count_tokens("What is your name?") ```python from google.genai import local_tokenizer -tokenizer = local_tokenizer.LocalTokenizer(model_name='gemini-3.5-flash') +tokenizer = local_tokenizer.LocalTokenizer(model_name='gemini-2.5-flash') result = tokenizer.compute_tokens("What is your name?") ``` diff --git a/google/genai/_local_tokenizer_loader.py b/google/genai/_local_tokenizer_loader.py index e80926601..626b9a23f 100644 --- a/google/genai/_local_tokenizer_loader.py +++ b/google/genai/_local_tokenizer_loader.py @@ -217,14 +217,17 @@ def get_tokenizer_name(model_name: str) -> str: def get_huggingface_tokenizer(tokenizer_name: str) -> Any: """Loads huggingface tokenizer from the given tokenizer name.""" - # Load the processor which includes the tokenizer + # Load the processor which includes the tokenizer. These deps are intentionally + # not part of the lightweight ``local-tokenizer`` extra; install + # ``google-genai[local-tokenizer-gemma4]`` for Gemma 4 models. try: from transformers import AutoProcessor - except ImportError: + except ImportError as e: raise ImportError( - "Please install transformers to use huggingface tokenizer: pip install" - " transformers" - ) from ImportError + 'The Gemma 4 local tokenizer requires the Hugging Face stack ' + '(transformers, torch, etc.). Install it with: ' + 'pip install "google-genai[local-tokenizer-gemma4]"' + ) from e processor = AutoProcessor.from_pretrained( # type: ignore[no-untyped-call] GEMMA_TOKENIZER_TO_MODEL_NAMES[tokenizer_name] ) diff --git a/google/genai/local_tokenizer.py b/google/genai/local_tokenizer.py index 2d2031e55..ccb113f11 100644 --- a/google/genai/local_tokenizer.py +++ b/google/genai/local_tokenizer.py @@ -280,6 +280,11 @@ class LocalTokenizer: This class provides a local tokenizer for text only token counting. + Install ``google-genai[local-tokenizer]`` for SentencePiece-backed models + (for example ``gemini-2.5-flash``). Gemma 4–backed models (for example + ``gemini-3.5-flash``) need the heavier ``google-genai[local-tokenizer-gemma4]`` + extra. + LIMITATIONS: - Only supports text based tokenization and no multimodal tokenization. - Forward compatibility depends on the open-source tokenizer models for future diff --git a/google/genai/tests/local_tokenizer/test_local_tokenizer_loader.py b/google/genai/tests/local_tokenizer/test_local_tokenizer_loader.py index 99ed29cce..db08c9118 100644 --- a/google/genai/tests/local_tokenizer/test_local_tokenizer_loader.py +++ b/google/genai/tests/local_tokenizer/test_local_tokenizer_loader.py @@ -76,6 +76,24 @@ def test_get_tokenizer_name_unsupported(self): ): loader.get_tokenizer_name("unsupported-model") + def test_get_huggingface_tokenizer_missing_deps_message(self): + with patch.dict('sys.modules', {'transformers': None}): + # Force import failure even if transformers is installed in the env. + import builtins + + real_import = builtins.__import__ + + def _fake_import(name, *args, **kwargs): + if name == 'transformers' or name.startswith('transformers.'): + raise ImportError('No module named transformers') + return real_import(name, *args, **kwargs) + + with patch('builtins.__import__', side_effect=_fake_import): + with self.assertRaisesRegex( + ImportError, r'local-tokenizer-gemma4' + ): + loader.get_huggingface_tokenizer('gemma4') + @patch("genai._local_tokenizer_loader.os.rename") @patch("genai._local_tokenizer_loader.os.makedirs") diff --git a/pyproject.toml b/pyproject.toml index b07c2d44a..8547c63c7 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -39,9 +39,16 @@ dependencies = [ [project.optional-dependencies] aiohttp = ["aiohttp>=3.10.11, <4.0.0"] +# SentencePiece path used by Gemini 2.x / gemma3 local tokenizers. local-tokenizer = [ "sentencepiece>=0.2.0", "protobuf", +] +# Hugging Face AutoProcessor path used by Gemma 4 local tokenizers +# (e.g. gemini-3.5-flash). Pulls torch/transformers and is much larger. +local-tokenizer-gemma4 = [ + "sentencepiece>=0.2.0", + "protobuf", "pillow", "torch", "torchvision",