mirror of
https://github.com/BerriAI/litellm.git
synced 2026-09-29 01:42:19 +00:00
130 lines
5.5 KiB
Python
130 lines
5.5 KiB
Python
import json
|
|
from typing import Final
|
|
|
|
import pytest
|
|
import tiktoken
|
|
from tokenizers import Tokenizer as ReferenceTokenizer
|
|
|
|
from litellm.rust_bridge import _native
|
|
from litellm.utils import claude_json_str
|
|
from tests.test_litellm.litellm_core_utils.test_decode_special_tokens import TOKENIZER_JSON
|
|
|
|
pytestmark = pytest.mark.requires_rust_extension
|
|
|
|
|
|
def test_tiktoken_codec_round_trips_and_counts() -> None:
|
|
tokenizer: Final = _native.Tokenizer.from_tiktoken("cl100k_base")
|
|
encoded: Final = tokenizer.encode("hello world")
|
|
|
|
assert tokenizer.name == "cl100k_base"
|
|
assert tokenizer.count("hello world") == len(encoded)
|
|
assert tokenizer.decode(encoded) == "hello world"
|
|
|
|
|
|
def test_huggingface_codec_skips_special_tokens() -> None:
|
|
tokenizer: Final = _native.Tokenizer.from_json(claude_json_str)
|
|
encoded: Final = tokenizer.encode("<SOS>hello<EOT>")
|
|
|
|
assert "<SOS>" in tokenizer.decode(encoded, skip_special_tokens=False)
|
|
assert tokenizer.decode(encoded, skip_special_tokens=True) == "hello"
|
|
|
|
|
|
def test_tiktoken_codec_keeps_the_requested_encoding_name() -> None:
|
|
assert _native.Tokenizer.from_tiktoken("gpt2").name == "gpt2"
|
|
assert _native.Tokenizer.from_tiktoken("r50k_base").name == "r50k_base"
|
|
assert _native.Tokenizer.from_tiktoken("gpt2").encode("hi") == _native.Tokenizer.from_tiktoken("r50k_base").encode(
|
|
"hi"
|
|
)
|
|
|
|
|
|
def test_tiktoken_codec_exposes_its_vocabulary() -> None:
|
|
reference: Final = tiktoken.get_encoding("cl100k_base")
|
|
tokenizer: Final = _native.Tokenizer.from_tiktoken("cl100k_base")
|
|
|
|
assert tokenizer.special_tokens() == reference._special_tokens
|
|
assert tokenizer.max_token_value() == reference.max_token_value
|
|
assert tokenizer.token_byte_values() == reference.token_byte_values()
|
|
assert tokenizer.encode_single_token(b"hello") == reference.encode_single_token("hello")
|
|
assert tokenizer.is_special_token(reference.eot_token) and not tokenizer.is_special_token(0)
|
|
with pytest.raises(KeyError):
|
|
tokenizer.encode_single_token(b"<|not-a-token|>")
|
|
|
|
|
|
def test_huggingface_codec_rejects_tiktoken_only_calls() -> None:
|
|
tokenizer: Final = _native.Tokenizer.from_json(claude_json_str)
|
|
with pytest.raises(ValueError, match="requires a tiktoken encoding"):
|
|
tokenizer.token_byte_values()
|
|
with pytest.raises(ValueError, match="requires a Hugging Face tokenizer"):
|
|
_native.Tokenizer.from_tiktoken("cl100k_base").get_vocab()
|
|
|
|
|
|
def test_unknown_tiktoken_encoding_raises_value_error() -> None:
|
|
with pytest.raises(ValueError, match="unsupported tokenizer"):
|
|
_native.Tokenizer.from_tiktoken("unknown-encoding")
|
|
|
|
|
|
def test_tiktoken_codec_decodes_truncated_unicode_like_python() -> None:
|
|
reference: Final = tiktoken.get_encoding("cl100k_base")
|
|
tokenizer: Final = _native.Tokenizer.from_tiktoken(reference.name)
|
|
encoded: Final = reference.encode("🙂漢字")
|
|
|
|
assert tuple(tokenizer.decode(encoded[:end]) for end in range(1, len(encoded) + 1)) == tuple(
|
|
reference.decode(encoded[:end]) for end in range(1, len(encoded) + 1)
|
|
)
|
|
|
|
|
|
FAST_TEXTS: Final = (
|
|
"",
|
|
"hello world <|endoftext|>",
|
|
"café 漢字 ع 🙂 line\r\n indented 123456789",
|
|
"<SOS>x<EOT> a\u0301 fi",
|
|
)
|
|
|
|
|
|
def test_fast_counting_is_an_opt_in_over_the_same_loaded_tokenizer() -> None:
|
|
for tokenizer in (
|
|
_native.Tokenizer.from_tiktoken("cl100k_base"),
|
|
_native.Tokenizer.from_tiktoken("o200k_base"),
|
|
_native.Tokenizer.from_json(claude_json_str),
|
|
):
|
|
assert [tokenizer.count(text, fast=True) for text in FAST_TEXTS] == [
|
|
tokenizer.count(text) for text in FAST_TEXTS
|
|
]
|
|
|
|
|
|
@pytest.mark.parametrize(
|
|
"name", ("cl100k_base", "o200k_base", "o200k_harmony", "p50k_base", "p50k_edit", "r50k_base", "gpt2")
|
|
)
|
|
@pytest.mark.asyncio
|
|
async def test_token_counter_counts_over_a_shared_tokenizer(name: str) -> None:
|
|
messages: Final = [{"role": "user", "content": "hello wide world"}, {"role": "assistant", "content": "ok"}]
|
|
body: Final = json.dumps({"model": "gpt-4", "messages": messages}).encode()
|
|
tokenizer: Final = _native.Tokenizer.from_tiktoken(name)
|
|
reference: Final = tiktoken.get_encoding(name)
|
|
for text in FAST_TEXTS:
|
|
assert tokenizer.count(text, fast=True) == tokenizer.count(text) == len(reference.encode_ordinary(text))
|
|
|
|
exact: Final = await _native.TokenCounter.from_tokenizer(tokenizer).acount_request(body)
|
|
fast: Final = await _native.TokenCounter.from_tokenizer(tokenizer, fast=True).acount_request(body)
|
|
|
|
assert exact == fast
|
|
assert exact["input_tokens"] == 3 + sum(
|
|
3 + len(reference.encode_ordinary(message["role"])) + len(reference.encode_ordinary(message["content"]))
|
|
for message in messages
|
|
)
|
|
|
|
|
|
@pytest.mark.parametrize("configured", (False, True))
|
|
@pytest.mark.asyncio
|
|
async def test_fast_count_preserves_huggingface_configuration(configured: bool) -> None:
|
|
reference: Final = ReferenceTokenizer.from_str(TOKENIZER_JSON)
|
|
if configured:
|
|
reference.enable_truncation(max_length=3)
|
|
reference.enable_padding(pad_id=0, pad_token="[UNK]", length=5)
|
|
tokenizer: Final = _native.Tokenizer.from_json(reference.to_str())
|
|
counter: Final = _native.TokenCounter.from_tokenizer(tokenizer, fast=True)
|
|
for text in ("", "Hello", "Hello World Hello World", "[BOS] Hello"):
|
|
expected: Final = len(reference.encode(text))
|
|
assert tokenizer.count(text, fast=True) == tokenizer.count(text) == expected
|
|
result: Final = await counter.acount_request(json.dumps({"prompt": text}).encode())
|
|
assert result["input_tokens"] == expected
|