feat: Gemma 4 local tokenizer support

PiperOrigin-RevId: 929258534

Author: yinghsienwu

google/genai/_local_tokenizer_loader.py

@@ -24,6 +24,7 @@
 import requests  # type: ignore
 import sentencepiece as spm
 from sentencepiece import sentencepiece_model_pb2
+from transformers import AutoProcessor
 
 
 # Source of truth: https://cloud.google.com/vertex-ai/generative-ai/docs/learn/models
@@ -47,6 +48,18 @@
     "gemini-3-pro-preview": "gemma3",
 }
 
+# https://github.com/google/gemma_pytorch stop supporting gemma 4 moving forward.
+_GEMINI_MODELS_TO_HUGGINGFACE_TOKENIZER_NAMES = {
+    "gemini-3.5-flash": "gemma4",
+    "gemini-3.1-flash-lite": "gemma4",
+    "gemini-3.1-pro-preview": "gemma4",
+    "gemini-4-flash-preview": "gemma4",
+}
+
+GEMMA_TOKENIZER_TO_MODEL_NAMES = {
+    "gemma4": "google/gemma-4-E4B-it",
+}
+
 
 @dataclasses.dataclass(frozen=True)
 class _TokenizerConfig:
@@ -202,11 +215,23 @@ def get_tokenizer_name(model_name: str) -> str:
     return _GEMINI_MODELS_TO_TOKENIZER_NAMES[model_name]
   if model_name in _GEMINI_STABLE_MODELS_TO_TOKENIZER_NAMES.keys():
     return _GEMINI_STABLE_MODELS_TO_TOKENIZER_NAMES[model_name]
+  if model_name in _GEMINI_MODELS_TO_HUGGINGFACE_TOKENIZER_NAMES.keys():
+    return _GEMINI_MODELS_TO_HUGGINGFACE_TOKENIZER_NAMES[model_name]
   raise ValueError(
       f"Model {model_name} is not supported. Supported models: {', '.join(_GEMINI_MODELS_TO_TOKENIZER_NAMES.keys())}, {', '.join(_GEMINI_STABLE_MODELS_TO_TOKENIZER_NAMES.keys())}.\n"  # pylint: disable=line-too-long
   )
 
 
+def get_huggingface_tokenizer(tokenizer_name: str):
+  """Loads huggingface tokenizer from the given tokenizer name."""
+  # Load the processor which includes the tokenizer
+  processor = AutoProcessor.from_pretrained(
+      GEMMA_TOKENIZER_TO_MODEL_NAMES[tokenizer_name]
+  )
+  # Access the underlying tokenizer if needed
+  return processor.tokenizer
+
+
 @functools.lru_cache()
 def get_sentencepiece(tokenizer_name: str) -> spm.SentencePieceProcessor:
   """Loads sentencepiece tokenizer from the given tokenizer name."""

google/genai/local_tokenizer.py

@@ -292,7 +292,10 @@ class LocalTokenizer:
   def __init__(self, model_name: str):
     self._tokenizer_name = loader.get_tokenizer_name(model_name)
     self._model_proto = loader.load_model_proto(self._tokenizer_name)
-    self._tokenizer = loader.get_sentencepiece(self._tokenizer_name)
+    if self._tokenizer_name != "gemma3":
+      self._tokenizer = loader.get_huggingface_tokenizer(self._tokenizer_name)
+    else:
+      self._tokenizer = loader.get_sentencepiece(self._tokenizer_name)
 
   @_common.experimental_warning(
       "The SDK's local tokenizer implementation is experimental and may change"
@@ -365,20 +368,39 @@ def compute_tokens(
       # tokens_info=[TokensInfo(token_ids=[279, 329, 1313, 2508, 13], tokens=[b' What', b' is', b' your', b' name', b'?'], role='user')]
     """
     processed_contents = t.t_contents(contents)
+    roles = []
+
     text_accumulator = _TextsAccumulator()
     for content in processed_contents:
       text_accumulator.add_content(content)
-    tokens_protos = self._tokenizer.EncodeAsImmutableProto(
-        text_accumulator.get_texts()
-    )
-
-    roles = []
-    for content in processed_contents:
       if content.parts:
         for _ in content.parts:
           roles.append(content.role)
 
     token_infos = []
+    if self._tokenizer_name != "gemma3":
+      # Use the HuggingFace tokenizer since gemma_pytorch is not available for
+      # gemma 4+.
+      token_ids = self._tokenizer.encode(list(text_accumulator.get_texts()))
+      for token_id, role in zip(token_ids, roles):
+        token_infos.append(
+            types.TokensInfo(
+                token_ids=token_id,
+                tokens=[
+                    token.replace("_", " ")
+                    .encode("utf-8")
+                    .replace(b"\xe2\x96\x81", b" ")
+                    for token in self._tokenizer.convert_ids_to_tokens(token_id)
+                ],
+                role=role,
+            )
+        )
+      return types.ComputeTokensResult(tokens_info=token_infos)
+
+    tokens_protos = self._tokenizer.EncodeAsImmutableProto(
+        text_accumulator.get_texts()
+    )
+
     for tokens_proto, role in zip(tokens_protos, roles):
       token_infos.append(
           types.TokensInfo(

google/genai/tests/local_tokenizer/test_local_tokenizer.py

@@ -341,3 +341,79 @@ def test_invalid_format(self):
   def test_invalid_hex_value(self):
     with self.assertRaisesRegex(ValueError, 'Invalid hex value'):
       local_tokenizer._parse_hex_byte('<0xFG>')
+
+
+class TestLocalTokenizerHuggingFace(unittest.TestCase):
+
+  def setUp(self):
+    self.mock_load_model_proto = patch(
+        'genai._local_tokenizer_loader.load_model_proto'
+    ).start()
+    self.mock_get_huggingface_tokenizer = patch(
+        'genai._local_tokenizer_loader.get_huggingface_tokenizer'
+    ).start()
+
+    self.mock_load_model_proto.return_value = MagicMock()
+    self.mock_tokenizer = MagicMock()
+    self.mock_get_huggingface_tokenizer.return_value = self.mock_tokenizer
+
+    # gemini-3.5-flash maps to gemma4 (HuggingFace)
+    self.tokenizer = local_tokenizer.LocalTokenizer(model_name='gemini-3.5-flash')
+
+  def tearDown(self):
+    patch.stopall()
+
+  def test_count_tokens_simple_string(self):
+    self.mock_tokenizer.encode.return_value = [[1, 2, 3]]
+    result = self.tokenizer.count_tokens('Hello world')
+    self.assertEqual(result.total_tokens, 3)
+    self.mock_tokenizer.encode.assert_called_once_with(['Hello world'])
+
+  def test_compute_tokens_simple_string(self):
+    self.mock_tokenizer.encode.return_value = [[1, 2, 3]]
+    self.mock_tokenizer.convert_ids_to_tokens.return_value = ['He', 'llo', ' world']
+
+    result = self.tokenizer.compute_tokens('Hello world')
+
+    self.assertEqual(len(result.tokens_info), 1)
+    self.assertEqual(result.tokens_info[0].token_ids, [1, 2, 3])
+    self.assertEqual(result.tokens_info[0].tokens, [b'He', b'llo', b' world'])
+    self.assertEqual(result.tokens_info[0].role, 'user')
+
+    self.mock_tokenizer.encode.assert_called_once_with(['Hello world'])
+    self.mock_tokenizer.convert_ids_to_tokens.assert_called_once_with([1, 2, 3])
+
+  def test_compute_tokens_special_characters(self):
+    self.mock_tokenizer.encode.return_value = [[1, 2]]
+    # Use U+2581 (lower one eighth block) and underscore
+    self.mock_tokenizer.convert_ids_to_tokens.return_value = ['_world', '\u2581hello']
+
+    result = self.tokenizer.compute_tokens('dummy')
+
+    self.assertEqual(result.tokens_info[0].tokens, [b' world', b' hello'])
+
+  def test_compute_tokens_with_chat_history(self):
+    self.mock_tokenizer.encode.return_value = [[1], [2, 3]]
+    self.mock_tokenizer.convert_ids_to_tokens.side_effect = [
+        ['Hello'],
+        ['Hi', ' there!']
+    ]
+    history = [
+        types.Content(role='user', parts=[types.Part(text='Hello')]),
+        types.Content(role='model', parts=[types.Part(text='Hi there!')]),
+    ]
+    result = self.tokenizer.compute_tokens(history)
+    self.assertEqual(len(result.tokens_info), 2)
+    self.assertEqual(result.tokens_info[0].token_ids, [1])
+    self.assertEqual(result.tokens_info[0].tokens, [b'Hello'])
+    self.assertEqual(result.tokens_info[0].role, 'user')
+    self.assertEqual(result.tokens_info[1].token_ids, [2, 3])
+    self.assertEqual(result.tokens_info[1].tokens, [b'Hi', b' there!'])
+    self.assertEqual(result.tokens_info[1].role, 'model')
+
+    self.mock_tokenizer.encode.assert_called_once_with(['Hello', 'Hi there!'])
+    self.mock_tokenizer.convert_ids_to_tokens.assert_has_calls([
+        unittest.mock.call([1]),
+        unittest.mock.call([2, 3])
+    ])
+

google/genai/tests/local_tokenizer/test_local_tokenizer_loader.py

@@ -58,6 +58,18 @@ def test_get_tokenizer_name_success(self):
         loader.get_tokenizer_name("gemini-2.5-pro-preview-06-05"), "gemma3"
     )
 
+  def test_get_tokenizer_name_huggingface(self):
+    self.assertEqual(loader.get_tokenizer_name("gemini-3.5-flash"), "gemma4")
+    self.assertEqual(
+        loader.get_tokenizer_name("gemini-3.1-flash-lite"), "gemma4"
+    )
+    self.assertEqual(
+        loader.get_tokenizer_name("gemini-3.1-pro-preview"), "gemma4"
+    )
+    self.assertEqual(
+        loader.get_tokenizer_name("gemini-4-flash-preview"), "gemma4"
+    )
+
   def test_get_tokenizer_name_unsupported(self):
     with self.assertRaisesRegex(
         ValueError, "Model unsupported-model is not supported"
@@ -233,3 +245,24 @@ def test_get_sentencepiece_caching(
 
     # Should only be loaded once due to lru_cache
     mock_get.assert_called_once()
+
+
+class TestGetHuggingFaceTokenizer(unittest.TestCase):
+
+  @patch("genai._local_tokenizer_loader.AutoProcessor")
+  def test_get_huggingface_tokenizer_success(self, mock_auto_processor):
+    mock_processor = MagicMock()
+    mock_tokenizer = MagicMock()
+    mock_processor.tokenizer = mock_tokenizer
+    mock_auto_processor.from_pretrained.return_value = mock_processor
+
+    tokenizer = loader.get_huggingface_tokenizer("gemma4")
+
+    self.assertEqual(tokenizer, mock_tokenizer)
+    mock_auto_processor.from_pretrained.assert_called_once_with(
+        "google/gemma-4-E4B-it"
+    )
+
+  def test_get_huggingface_tokenizer_unsupported(self):
+    with self.assertRaises(KeyError):
+      loader.get_huggingface_tokenizer("unsupported")

pyproject.toml

@@ -39,7 +39,7 @@ dependencies = [
 
 [project.optional-dependencies]
 aiohttp = ["aiohttp>=3.10.11, <4.0.0"]
-local-tokenizer = ["sentencepiece>=0.2.0", "protobuf"]
+local-tokenizer = ["sentencepiece>=0.2.0", "protobuf", "transformers"]
 pyopenssl = ["pyopenssl"]
 
 [project.urls]

requirements.txt

@@ -32,3 +32,4 @@ websockets==16.0
 mcp>=1.14.0; python_version > '3.9'
 sentencepiece>=0.2.0
 protobuf
+transformers>=5.10.1