first vommit
This commit is contained in:
Till-Immer committed
commit
a04e68bc95
4926 files changed
+764440
No files matched your search
@@ -0,0 +1,17 @@
|
||||
# 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.
|
||||
#
|
||||
|
||||
|
||||
"""Tests for the Google GenAI SDK's local_tokenizer module."""
|
||||
BIN
Binary file not shown.
BIN
Binary file not shown.
BIN
Binary file not shown.
+410
@@ -0,0 +1,410 @@
|
||||
# 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
|
||||
|
||||
|
||||
class TestLocalTokenizer(unittest.TestCase):
|
||||
|
||||
def setUp(self):
|
||||
# This setup will be used by all tests
|
||||
self.mock_load_model_proto = patch(
|
||||
'genai._local_tokenizer_loader.load_model_proto'
|
||||
).start()
|
||||
self.addCleanup(patch.stopall)
|
||||
self.mock_get_sentencepiece = patch(
|
||||
'genai._local_tokenizer_loader.get_sentencepiece'
|
||||
).start()
|
||||
|
||||
self.mock_load_model_proto.return_value = MagicMock()
|
||||
self.mock_tokenizer = MagicMock()
|
||||
self.mock_get_sentencepiece.return_value = self.mock_tokenizer
|
||||
|
||||
self.tokenizer = local_tokenizer.LocalTokenizer(model_name='gemini-3-pro-preview')
|
||||
|
||||
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_count_tokens_list_of_strings(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_count_tokens_with_content_object(self):
|
||||
self.mock_tokenizer.encode.return_value = [[1, 2, 3]]
|
||||
content = types.Content(parts=[types.Part(text='Hello world')])
|
||||
result = self.tokenizer.count_tokens(content)
|
||||
self.assertEqual(result.total_tokens, 3)
|
||||
self.mock_tokenizer.encode.assert_called_once_with(['Hello world'])
|
||||
|
||||
def test_count_tokens_with_chat_history(self):
|
||||
self.mock_tokenizer.encode.return_value = [[1, 2], [3, 4, 5]]
|
||||
history = [
|
||||
types.Content(role='user', parts=[types.Part(text='Hello')]),
|
||||
types.Content(role='model', parts=[types.Part(text='Hi there!')]),
|
||||
]
|
||||
result = self.tokenizer.count_tokens(history)
|
||||
self.assertEqual(result.total_tokens, 5)
|
||||
self.mock_tokenizer.encode.assert_called_once_with(['Hello', 'Hi there!'])
|
||||
|
||||
def test_count_tokens_with_tools(self):
|
||||
self.mock_tokenizer.encode.return_value = [
|
||||
[1],
|
||||
[1, 2],
|
||||
[1, 2, 3],
|
||||
[1, 2, 3, 4],
|
||||
[1, 2, 3, 4, 5],
|
||||
[1, 2, 3, 4, 5, 6],
|
||||
]
|
||||
tool = types.Tool(
|
||||
function_declarations=[
|
||||
types.FunctionDeclaration(
|
||||
name='get_weather',
|
||||
description='Get the weather for a location',
|
||||
parameters=types.Schema(
|
||||
type=types.Type.OBJECT,
|
||||
properties={
|
||||
'location': types.Schema(
|
||||
type=types.Type.STRING, description='The location'
|
||||
)
|
||||
},
|
||||
required=['location'],
|
||||
),
|
||||
)
|
||||
]
|
||||
)
|
||||
config = types.CountTokensConfig(tools=[tool])
|
||||
result = self.tokenizer.count_tokens(
|
||||
'What is the weather in Boston?', config=config
|
||||
)
|
||||
self.assertEqual(result.total_tokens, 21)
|
||||
self.mock_tokenizer.encode.assert_called_once_with([
|
||||
'What is the weather in Boston?',
|
||||
'get_weather',
|
||||
'Get the weather for a location',
|
||||
'location',
|
||||
'location',
|
||||
'The location',
|
||||
])
|
||||
|
||||
def test_count_tokens_with_function_call(self):
|
||||
self.mock_tokenizer.encode.return_value = [[1, 2], [3], [4, 5]]
|
||||
content = types.Content(
|
||||
role='model',
|
||||
parts=[
|
||||
types.Part(
|
||||
function_call=types.FunctionCall(
|
||||
name='get_weather', args={'location': 'Boston'}
|
||||
)
|
||||
)
|
||||
],
|
||||
)
|
||||
result = self.tokenizer.count_tokens(content)
|
||||
self.assertEqual(result.total_tokens, 5)
|
||||
self.mock_tokenizer.encode.assert_called_once_with(
|
||||
['get_weather', 'location', 'Boston']
|
||||
)
|
||||
|
||||
def test_count_tokens_with_function_response(self):
|
||||
self.mock_tokenizer.encode.return_value = [[1, 2], [3], [4, 5]]
|
||||
content = types.Content(
|
||||
role='user',
|
||||
parts=[
|
||||
types.Part(
|
||||
function_response=types.FunctionResponse(
|
||||
name='get_weather', response={'weather': 'sunny'}
|
||||
)
|
||||
)
|
||||
],
|
||||
)
|
||||
result = self.tokenizer.count_tokens(content)
|
||||
self.assertEqual(result.total_tokens, 5)
|
||||
self.mock_tokenizer.encode.assert_called_once_with(
|
||||
['get_weather', 'weather', 'sunny']
|
||||
)
|
||||
|
||||
def test_count_tokens_with_unsupported_content(self):
|
||||
with self.assertRaises(ValueError):
|
||||
self.tokenizer.count_tokens(
|
||||
[
|
||||
types.Content(
|
||||
parts=[
|
||||
types.Part(
|
||||
inline_data=types.Blob(
|
||||
data=b'test', mime_type='image/png'
|
||||
)
|
||||
)
|
||||
]
|
||||
)
|
||||
]
|
||||
)
|
||||
|
||||
def test_count_tokens_with_system_instruction(self):
|
||||
self.mock_tokenizer.encode.return_value = [[1, 2, 3], [4, 5]]
|
||||
config = types.CountTokensConfig(
|
||||
system_instruction=types.Content(
|
||||
parts=[types.Part(text='You are a helpful assistant.')]
|
||||
)
|
||||
)
|
||||
result = self.tokenizer.count_tokens('Hello', config=config)
|
||||
self.assertEqual(result.total_tokens, 5)
|
||||
self.mock_tokenizer.encode.assert_called_once_with(
|
||||
['Hello', 'You are a helpful assistant.']
|
||||
)
|
||||
|
||||
def test_count_tokens_with_response_schema(self):
|
||||
self.mock_tokenizer.encode.return_value = [
|
||||
[1],
|
||||
[1, 2],
|
||||
[1, 2, 3],
|
||||
[1, 2, 3, 4],
|
||||
[1, 2, 3, 4, 5],
|
||||
]
|
||||
schema = types.Schema(
|
||||
type=types.Type.OBJECT,
|
||||
format='schema_format',
|
||||
description='Recipe schema',
|
||||
enum=['schema_enum1', 'schema_enum2'],
|
||||
properties={
|
||||
'recipe_name': types.Schema(
|
||||
type=types.Type.STRING,
|
||||
description='Name of the recipe',
|
||||
)
|
||||
},
|
||||
items=types.Schema(
|
||||
type=types.Type.STRING,
|
||||
description='Item in the recipe',
|
||||
),
|
||||
example={
|
||||
'recipe_example': types.Schema(
|
||||
type=types.Type.STRING,
|
||||
description='example in the recipe',
|
||||
)
|
||||
},
|
||||
required=['recipe_name'],
|
||||
)
|
||||
config = types.CountTokensConfig(
|
||||
generation_config=types.GenerationConfig(response_schema=schema)
|
||||
)
|
||||
result = self.tokenizer.count_tokens(
|
||||
'Generate a recipe for chocolate chip cookies.', config=config
|
||||
)
|
||||
self.assertEqual(result.total_tokens, 15)
|
||||
self.mock_tokenizer.encode.assert_called_once_with([
|
||||
'Generate a recipe for chocolate chip cookies.',
|
||||
'schema_format',
|
||||
'Recipe schema',
|
||||
'schema_enum1',
|
||||
'schema_enum2',
|
||||
'recipe_name',
|
||||
'Item in the recipe',
|
||||
'recipe_name',
|
||||
'Name of the recipe',
|
||||
'recipe_example',
|
||||
])
|
||||
|
||||
def test_count_tokens_with_unsupported_fields_logs_warning(self):
|
||||
self.mock_tokenizer.encode.return_value = [[1, 2, 3]]
|
||||
content_with_unsupported = types.Content(
|
||||
role='user',
|
||||
parts=[
|
||||
types.Part(text='hello'),
|
||||
# executable_code is test_count_tokens_with_system_instruction by _TextsAccumulator
|
||||
types.Part(
|
||||
executable_code=types.ExecutableCode(
|
||||
language='PYTHON', code='print(1)'
|
||||
)
|
||||
),
|
||||
],
|
||||
)
|
||||
with self.assertLogs('google_genai.local_tokenizer', level='WARNING') as cm:
|
||||
self.tokenizer.count_tokens(content_with_unsupported)
|
||||
self.assertIn(
|
||||
'Content contains unsupported types for token counting', cm.output[0]
|
||||
)
|
||||
|
||||
def test_compute_tokens_simple_string(self):
|
||||
mock_spt = MagicMock()
|
||||
mock_spt.pieces = [
|
||||
MagicMock(id=1, piece='He'),
|
||||
MagicMock(id=2, piece='llo'),
|
||||
MagicMock(id=3, piece=' world'),
|
||||
]
|
||||
self.mock_tokenizer.EncodeAsImmutableProto.return_value = [mock_spt]
|
||||
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.EncodeAsImmutableProto.assert_called_once_with(
|
||||
['Hello world']
|
||||
)
|
||||
|
||||
def test_compute_tokens_with_chat_history(self):
|
||||
mock_spt1 = MagicMock()
|
||||
mock_spt1.pieces = [MagicMock(id=1, piece='Hello')]
|
||||
mock_spt2 = MagicMock()
|
||||
mock_spt2.pieces = [
|
||||
MagicMock(id=2, piece='Hi'),
|
||||
MagicMock(id=3, piece=' there!'),
|
||||
]
|
||||
self.mock_tokenizer.EncodeAsImmutableProto.return_value = [
|
||||
mock_spt1,
|
||||
mock_spt2,
|
||||
]
|
||||
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.EncodeAsImmutableProto.assert_called_once_with(
|
||||
['Hello', 'Hi there!']
|
||||
)
|
||||
|
||||
def test_compute_tokens_with_byte_tokens(self):
|
||||
mock_spt = MagicMock()
|
||||
mock_spt.pieces = [
|
||||
MagicMock(id=1, piece='<0x48>'),
|
||||
MagicMock(id=2, piece='ello'),
|
||||
]
|
||||
self.mock_tokenizer.EncodeAsImmutableProto.return_value = [mock_spt]
|
||||
self.tokenizer._model_proto = sentencepiece_model_pb2.ModelProto(
|
||||
pieces=[
|
||||
sentencepiece_model_pb2.ModelProto.SentencePiece(),
|
||||
sentencepiece_model_pb2.ModelProto.SentencePiece(
|
||||
type=sentencepiece_model_pb2.ModelProto.SentencePiece.Type.BYTE
|
||||
),
|
||||
sentencepiece_model_pb2.ModelProto.SentencePiece(
|
||||
type=sentencepiece_model_pb2.ModelProto.SentencePiece.Type.NORMAL
|
||||
),
|
||||
]
|
||||
)
|
||||
result = self.tokenizer.compute_tokens('Hello')
|
||||
self.assertEqual(len(result.tokens_info), 1)
|
||||
self.assertEqual(result.tokens_info[0].token_ids, [1, 2])
|
||||
self.assertEqual(result.tokens_info[0].tokens, [b'H', b'ello'])
|
||||
self.mock_tokenizer.EncodeAsImmutableProto.assert_called_once_with(
|
||||
['Hello']
|
||||
)
|
||||
|
||||
|
||||
class TestParseHexByte(unittest.TestCase):
|
||||
|
||||
def test_valid_hex(self):
|
||||
self.assertEqual(local_tokenizer._parse_hex_byte('<0x41>'), 65)
|
||||
self.assertEqual(local_tokenizer._parse_hex_byte('<0xFF>'), 255)
|
||||
self.assertEqual(local_tokenizer._parse_hex_byte('<0x00>'), 0)
|
||||
|
||||
def test_invalid_length(self):
|
||||
with self.assertRaisesRegex(ValueError, 'Invalid byte length'):
|
||||
local_tokenizer._parse_hex_byte('<0x41')
|
||||
with self.assertRaisesRegex(ValueError, 'Invalid byte length'):
|
||||
local_tokenizer._parse_hex_byte('<0x411>')
|
||||
|
||||
def test_invalid_format(self):
|
||||
with self.assertRaisesRegex(ValueError, 'Invalid byte format'):
|
||||
local_tokenizer._parse_hex_byte(' 0x41>')
|
||||
with self.assertRaisesRegex(ValueError, 'Invalid byte format'):
|
||||
local_tokenizer._parse_hex_byte('<0x41 ')
|
||||
|
||||
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_get_huggingface_tokenizer = patch(
|
||||
'genai._local_tokenizer_loader.get_huggingface_tokenizer'
|
||||
).start()
|
||||
self.addCleanup(patch.stopall)
|
||||
|
||||
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 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])
|
||||
])
|
||||
+249
@@ -0,0 +1,249 @@
|
||||
# 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, mock_open, patch
|
||||
|
||||
import sentencepiece as spm
|
||||
from sentencepiece import sentencepiece_model_pb2
|
||||
|
||||
from ... import _local_tokenizer_loader as loader
|
||||
|
||||
# A minimal valid sentencepiece model proto
|
||||
FAKE_MODEL_CONTENT = sentencepiece_model_pb2.ModelProto(
|
||||
pieces=[
|
||||
sentencepiece_model_pb2.ModelProto.SentencePiece(
|
||||
piece="<unk>",
|
||||
score=0,
|
||||
type=sentencepiece_model_pb2.ModelProto.SentencePiece.Type.UNKNOWN,
|
||||
),
|
||||
sentencepiece_model_pb2.ModelProto.SentencePiece(
|
||||
piece="<s>",
|
||||
score=0,
|
||||
type=sentencepiece_model_pb2.ModelProto.SentencePiece.Type.CONTROL,
|
||||
),
|
||||
sentencepiece_model_pb2.ModelProto.SentencePiece(
|
||||
piece="</s>",
|
||||
score=0,
|
||||
type=sentencepiece_model_pb2.ModelProto.SentencePiece.Type.CONTROL,
|
||||
),
|
||||
sentencepiece_model_pb2.ModelProto.SentencePiece(
|
||||
piece="a",
|
||||
score=0,
|
||||
type=sentencepiece_model_pb2.ModelProto.SentencePiece.Type.NORMAL,
|
||||
),
|
||||
]
|
||||
).SerializeToString()
|
||||
|
||||
GEMMA3_HASH = "1299c11d7cf632ef3b4e11937501358ada021bbdf7c47638d13c0ee982f2e79c"
|
||||
|
||||
|
||||
class TestGetTokenizerName(unittest.TestCase):
|
||||
|
||||
def test_get_tokenizer_name_success(self):
|
||||
self.assertEqual(loader.get_tokenizer_name("gemini-2.5-pro"), "gemma3")
|
||||
self.assertEqual(
|
||||
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"
|
||||
):
|
||||
loader.get_tokenizer_name("unsupported-model")
|
||||
|
||||
|
||||
@patch("genai._local_tokenizer_loader.os.rename")
|
||||
@patch("genai._local_tokenizer_loader.os.makedirs")
|
||||
@patch("genai._local_tokenizer_loader.os.remove")
|
||||
@patch("genai._local_tokenizer_loader.open", new_callable=mock_open)
|
||||
@patch("genai._local_tokenizer_loader.os.path.exists")
|
||||
@patch("genai._local_tokenizer_loader.requests.get")
|
||||
@patch("genai._local_tokenizer_loader.hashlib.sha256")
|
||||
class TestLoaderFunctions(unittest.TestCase):
|
||||
|
||||
def setUp(self):
|
||||
# Clear caches before each test
|
||||
loader.load_model_proto.cache_clear()
|
||||
loader.get_sentencepiece.cache_clear()
|
||||
# Patch tempfile.gettempdir to control cache location
|
||||
self.tempdir_patcher = patch(
|
||||
"tempfile.gettempdir", return_value="/tmp/fake_temp_dir"
|
||||
)
|
||||
self.mock_tempdir = self.tempdir_patcher.start()
|
||||
|
||||
def tearDown(self):
|
||||
self.tempdir_patcher.stop()
|
||||
|
||||
def _setup_get_mock(self, mock_get):
|
||||
mock_response = MagicMock()
|
||||
mock_response.content = FAKE_MODEL_CONTENT
|
||||
mock_response.raise_for_status.return_value = None
|
||||
mock_get.return_value = mock_response
|
||||
|
||||
def test_load_model_proto_from_url(
|
||||
self,
|
||||
mock_sha256,
|
||||
mock_get,
|
||||
mock_exists,
|
||||
mock_open_func,
|
||||
mock_remove,
|
||||
mock_makedirs,
|
||||
mock_rename,
|
||||
):
|
||||
mock_exists.return_value = False # Don't use cache
|
||||
self._setup_get_mock(mock_get)
|
||||
mock_sha256.return_value.hexdigest.return_value = GEMMA3_HASH
|
||||
|
||||
proto = loader.load_model_proto("gemma3")
|
||||
|
||||
self.assertIsInstance(proto, sentencepiece_model_pb2.ModelProto)
|
||||
self.assertEqual(len(proto.pieces), 4)
|
||||
mock_get.assert_called_once()
|
||||
mock_makedirs.assert_called_once()
|
||||
mock_open_func.assert_called()
|
||||
mock_rename.assert_called_once()
|
||||
|
||||
def test_load_model_proto_from_cache(
|
||||
self,
|
||||
mock_sha256,
|
||||
mock_get,
|
||||
mock_exists,
|
||||
mock_open_func,
|
||||
mock_remove,
|
||||
mock_makedirs,
|
||||
mock_rename,
|
||||
):
|
||||
mock_exists.return_value = True # Use cache
|
||||
mock_open_func.return_value.read.return_value = FAKE_MODEL_CONTENT
|
||||
mock_sha256.return_value.hexdigest.return_value = GEMMA3_HASH
|
||||
|
||||
proto = loader.load_model_proto("gemma3")
|
||||
|
||||
self.assertIsInstance(proto, sentencepiece_model_pb2.ModelProto)
|
||||
mock_get.assert_not_called()
|
||||
|
||||
def test_load_model_proto_corrupted_cache(
|
||||
self,
|
||||
mock_sha256,
|
||||
mock_get,
|
||||
mock_exists,
|
||||
mock_open_func,
|
||||
mock_remove,
|
||||
mock_makedirs,
|
||||
mock_rename,
|
||||
):
|
||||
mock_exists.return_value = True # Use cache initially
|
||||
self._setup_get_mock(mock_get)
|
||||
mock_open_func.return_value.__enter__.return_value.read.return_value = (
|
||||
b"corrupted"
|
||||
)
|
||||
|
||||
# First hash for corrupted cache, second for good download
|
||||
mock_sha256.side_effect = [
|
||||
MagicMock(hexdigest=MagicMock(return_value="wrong_hash")),
|
||||
MagicMock(hexdigest=MagicMock(return_value=GEMMA3_HASH)),
|
||||
]
|
||||
|
||||
proto = loader.load_model_proto("gemma3")
|
||||
|
||||
self.assertIsInstance(proto, sentencepiece_model_pb2.ModelProto)
|
||||
mock_remove.assert_called_once()
|
||||
mock_get.assert_called_once()
|
||||
|
||||
def test_load_model_proto_bad_hash_from_url(
|
||||
self,
|
||||
mock_sha256,
|
||||
mock_get,
|
||||
mock_exists,
|
||||
mock_open_func,
|
||||
mock_remove,
|
||||
mock_makedirs,
|
||||
mock_rename,
|
||||
):
|
||||
mock_exists.return_value = False
|
||||
self._setup_get_mock(mock_get)
|
||||
mock_sha256.return_value.hexdigest.return_value = "wrong_hash"
|
||||
|
||||
with self.assertRaisesRegex(
|
||||
ValueError, "Downloaded model file is corrupted"
|
||||
):
|
||||
loader.load_model_proto("gemma3")
|
||||
|
||||
def test_load_model_proto_unsupported(self, *args):
|
||||
with self.assertRaisesRegex(
|
||||
ValueError, "Tokenizer unsupported is not supported"
|
||||
):
|
||||
loader.load_model_proto("unsupported")
|
||||
|
||||
def test_get_sentencepiece_success(
|
||||
self,
|
||||
mock_sha256,
|
||||
mock_get,
|
||||
mock_exists,
|
||||
mock_open_func,
|
||||
mock_remove,
|
||||
mock_makedirs,
|
||||
mock_rename,
|
||||
):
|
||||
mock_exists.return_value = False
|
||||
self._setup_get_mock(mock_get)
|
||||
mock_sha256.return_value.hexdigest.return_value = GEMMA3_HASH
|
||||
|
||||
processor = loader.get_sentencepiece("gemma3")
|
||||
|
||||
self.assertIsInstance(processor, spm.SentencePieceProcessor)
|
||||
mock_get.assert_called_once()
|
||||
|
||||
def test_get_sentencepiece_unsupported(self, *args):
|
||||
with self.assertRaisesRegex(
|
||||
ValueError, "Tokenizer unsupported is not supported"
|
||||
):
|
||||
loader.get_sentencepiece("unsupported")
|
||||
|
||||
def test_get_sentencepiece_caching(
|
||||
self,
|
||||
mock_sha256,
|
||||
mock_get,
|
||||
mock_exists,
|
||||
mock_open_func,
|
||||
mock_remove,
|
||||
mock_makedirs,
|
||||
mock_rename,
|
||||
):
|
||||
mock_exists.return_value = False
|
||||
self._setup_get_mock(mock_get)
|
||||
mock_sha256.return_value.hexdigest.return_value = GEMMA3_HASH
|
||||
|
||||
# Call twice
|
||||
loader.get_sentencepiece("gemma3")
|
||||
loader.get_sentencepiece("gemma3")
|
||||
|
||||
# Should only be loaded once due to lru_cache
|
||||
mock_get.assert_called_once()
|
||||
|
||||
|
||||
Reference in new issue
Block a user