Abdelrahman Abdallah
68 lines
2.7 KiB
Python
68 lines
2.7 KiB
Python
from __future__ import annotations
|
|
|
|
import unittest
|
|
from unittest.mock import patch
|
|
|
|
from src.tokenizer_runtime import (
|
|
ResolvedTokenCounter,
|
|
TokenCounterInfo,
|
|
clear_token_counter_cache,
|
|
count_tokens,
|
|
describe_token_counter,
|
|
resolve_token_counter,
|
|
)
|
|
|
|
class TokenizerRuntimeTests(unittest.TestCase):
|
|
def tearDown(self) -> None:
|
|
clear_token_counter_cache()
|
|
|
|
def test_gpt_models_prefer_tiktoken_backend_when_available(self) -> None:
|
|
fake_counter = ResolvedTokenCounter(
|
|
info=TokenCounterInfo(
|
|
backend='tiktoken',
|
|
source='o200k_base',
|
|
accurate=True,
|
|
),
|
|
count_text=lambda text: len(text.split()),
|
|
)
|
|
with patch('src.tokenizer_runtime._try_build_tiktoken_counter', return_value=fake_counter):
|
|
with patch('src.tokenizer_runtime._try_build_transformers_counter', return_value=None):
|
|
info = describe_token_counter('gpt-4o-mini')
|
|
token_count = count_tokens('hello world from claw code', 'gpt-4o-mini')
|
|
|
|
self.assertEqual(info.backend, 'tiktoken')
|
|
self.assertTrue(info.accurate)
|
|
self.assertEqual(token_count, 5)
|
|
|
|
def test_transformers_backend_can_be_selected_with_env_override(self) -> None:
|
|
fake_counter = ResolvedTokenCounter(
|
|
info=TokenCounterInfo(
|
|
backend='transformers',
|
|
source='/tmp/fake-tokenizer (local_files_only)',
|
|
accurate=True,
|
|
),
|
|
count_text=lambda text: len(text.split()),
|
|
)
|
|
with patch.dict(
|
|
'os.environ',
|
|
{'CLAW_CODE_TOKENIZER_PATH': '/tmp/fake-tokenizer'},
|
|
clear=False,
|
|
):
|
|
with patch('src.tokenizer_runtime._try_build_transformers_counter', return_value=fake_counter):
|
|
info = describe_token_counter('Qwen/Qwen3-Coder-30B-A3B-Instruct')
|
|
token_count = count_tokens('one two three', 'Qwen/Qwen3-Coder-30B-A3B-Instruct')
|
|
|
|
self.assertEqual(info.backend, 'transformers')
|
|
self.assertTrue(info.accurate)
|
|
self.assertEqual(token_count, 3)
|
|
|
|
def test_fallback_backend_is_used_when_all_tokenizers_fail(self) -> None:
|
|
with patch('src.tokenizer_runtime._try_build_tiktoken_counter', return_value=None):
|
|
with patch('src.tokenizer_runtime._try_build_transformers_counter', return_value=None):
|
|
counter = resolve_token_counter('unknown-model')
|
|
token_count = count_tokens('abcd' * 5, 'unknown-model')
|
|
|
|
self.assertEqual(counter.info.backend, 'heuristic')
|
|
self.assertFalse(counter.info.accurate)
|
|
self.assertGreater(token_count, 0)
|