diff --git a/google/genai/local_tokenizer.py b/google/genai/local_tokenizer.py index 2d2031e55..f563fe173 100644 --- a/google/genai/local_tokenizer.py +++ b/google/genai/local_tokenizer.py @@ -59,6 +59,9 @@ class _TextsAccumulator: def __init__(self) -> None: self._texts: list[str] = [] + def __len__(self) -> int: + return len(self._texts) + def get_texts(self) -> Iterable[str]: return self._texts @@ -369,14 +372,19 @@ 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 = [] + roles: list[Optional[str]] = [] text_accumulator = _TextsAccumulator() for content in processed_contents: + texts_before = len(text_accumulator) text_accumulator.add_content(content) - if content.parts: - for _ in content.parts: - roles.append(content.role) + # A part does not map to exactly one tokenized text: a function_call + # or function_response part contributes the function name plus every + # key and string value of its args/response as separate texts, and a + # thought_signature-only part contributes none. Extend `roles` by the + # number of texts the accumulator actually added, so the zip() below + # stays aligned with `text_accumulator.get_texts()`. + roles.extend([content.role] * (len(text_accumulator) - texts_before)) token_infos = [] if self._tokenizer_name in loader.GEMMA_TOKENIZER_TO_MODEL_NAMES: diff --git a/google/genai/tests/local_tokenizer/test_roles_alignment.py b/google/genai/tests/local_tokenizer/test_roles_alignment.py new file mode 100644 index 000000000..1fb6cba13 --- /dev/null +++ b/google/genai/tests/local_tokenizer/test_roles_alignment.py @@ -0,0 +1,73 @@ +# Copyright 2025 Google LLC +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +# + +import unittest +from unittest.mock import MagicMock, patch + +from sentencepiece import sentencepiece_model_pb2 + +from ... import local_tokenizer +from ... import types + +_FC_CONTENT = types.Content( + role='model', + parts=[types.Part(function_call=types.FunctionCall( + name='get_weather', args={'location': 'Boston'}))], +) +_USER_CONTENT = types.Content(role='user', parts=[types.Part(text='thanks')]) +# accumulator yields: 'get_weather', 'location', 'Boston', 'thanks' +_EXPECTED_ROLES = ['model', 'model', 'model', 'user'] + + +class TestSentencePieceBranch(unittest.TestCase): + + def setUp(self): + patch('genai._local_tokenizer_loader.load_model_proto').start() + m = patch('genai._local_tokenizer_loader.get_sentencepiece').start() + self.addCleanup(patch.stopall) + self.mock_tokenizer = MagicMock() + m.return_value = self.mock_tokenizer + self.tokenizer = local_tokenizer.LocalTokenizer(model_name='gemini-2.5-flash') + self.tokenizer._model_proto = sentencepiece_model_pb2.ModelProto( + pieces=[sentencepiece_model_pb2.ModelProto.SentencePiece(piece='x')] * 10) + + def test_roles_align_with_texts(self): + def proto(i): + p = MagicMock(); p.pieces = [MagicMock(id=1, piece=f't{i}')]; return p + self.mock_tokenizer.EncodeAsImmutableProto.side_effect = ( + lambda texts: [proto(i) for i in range(len(texts))]) + r = self.tokenizer.compute_tokens([_FC_CONTENT, _USER_CONTENT]) + self.assertEqual([t.role for t in r.tokens_info], _EXPECTED_ROLES) + + +class TestHuggingFaceBranch(unittest.TestCase): + + def setUp(self): + m = patch('genai._local_tokenizer_loader.get_huggingface_tokenizer').start() + self.addCleanup(patch.stopall) + self.mock_tokenizer = MagicMock() + m.return_value = self.mock_tokenizer + self.tokenizer = local_tokenizer.LocalTokenizer(model_name='gemini-3.5-flash') + + def test_roles_align_with_texts(self): + self.mock_tokenizer.encode.side_effect = ( + lambda texts: [[i] for i in range(len(texts))]) + self.mock_tokenizer.convert_ids_to_tokens.side_effect = lambda ids: ['tok'] + r = self.tokenizer.compute_tokens([_FC_CONTENT, _USER_CONTENT]) + self.assertEqual([t.role for t in r.tokens_info], _EXPECTED_ROLES) + + +if __name__ == '__main__': + unittest.main()