diff --git a/litellm/utils.py b/litellm/utils.py index 18df5e2abf7..af22b11224b 100644 --- a/litellm/utils.py +++ b/litellm/utils.py @@ -2202,15 +2202,20 @@ def _is_streaming_request( def _select_tokenizer(model: str, custom_tokenizer: CustomHuggingfaceTokenizer | None = None): if custom_tokenizer is not None: - _tokenizer: Final = create_pretrained_tokenizer( + return _select_custom_tokenizer_helper( identifier=custom_tokenizer["identifier"], revision=custom_tokenizer["revision"], auth_token=custom_tokenizer["auth_token"], ) - return _tokenizer return _select_tokenizer_helper(model=model) +@lru_cache(maxsize=DEFAULT_MAX_LRU_CACHE_SIZE) +def _select_custom_tokenizer_helper(identifier: str, revision: str, auth_token: str | None) -> SelectTokenizerResponse: + verbose_logger.debug("Loading custom HuggingFace tokenizer %s (revision %s)", identifier, revision) + return create_pretrained_tokenizer(identifier=identifier, revision=revision, auth_token=auth_token) + + @lru_cache(maxsize=DEFAULT_MAX_LRU_CACHE_SIZE) def _select_tokenizer_helper(model: str) -> SelectTokenizerResponse: if litellm.disable_hf_tokenizer_download is True: diff --git a/tests/test_litellm/proxy/test_proxy_server.py b/tests/test_litellm/proxy/test_proxy_server.py index 04173ced776..552044e2598 100644 --- a/tests/test_litellm/proxy/test_proxy_server.py +++ b/tests/test_litellm/proxy/test_proxy_server.py @@ -13521,3 +13521,60 @@ async def test_token_counter_loads_a_custom_tokenizer_off_the_event_loop(monkeyp assert response.tokenizer_type == "huggingface_tokenizer" assert response.total_tokens > 0 assert_loop_stayed_free(took, lags) + + +async def test_token_counter_loads_a_custom_tokenizer_once_per_identifier_revision_and_token(monkeypatch): + from tokenizers import Tokenizer + + from litellm import Router + from litellm.types.router import DeploymentTypedDict + + claude_tokenizer: Final = litellm.utils._select_tokenizer("claude-fable-5")["tokenizer"] + loads: Final[list[tuple[str, str, str | None]]] = [] + + class CountingHubTokenizer: + @staticmethod + def from_pretrained(identifier: str, revision: str = "main", token: str | None = None) -> Tokenizer: + loads.append((identifier, revision, token)) + return claude_tokenizer + + def deployment(model_name: str, revision: str, auth_token: str | None) -> DeploymentTypedDict: + return { + "model_name": model_name, + "litellm_params": {"model": "openai/self-hosted-model", "api_base": "http://localhost:8080/v1"}, + "model_info": { + "custom_tokenizer": {"identifier": "my-org/tokenizer", "revision": revision, "auth_token": auth_token} + }, + } + + monkeypatch.setattr(litellm.utils, "Tokenizer", CountingHubTokenizer) + monkeypatch.setattr( + "litellm.proxy.proxy_server.llm_router", + Router( + model_list=[ + deployment("self-hosted", "main", None), + deployment("self-hosted-pinned", "v2", None), + deployment("self-hosted-private", "main", "hf_test_token"), + ] + ), + ) + litellm.utils._select_custom_tokenizer_helper.cache_clear() + try: + responses: Final = [ + await proxy_server_module.token_counter(TokenCountRequest(model="self-hosted", prompt="count me once")) + for _ in range(3) + ] + assert loads == [("my-org/tokenizer", "main", None)] + assert all(response.tokenizer_type == "huggingface_tokenizer" for response in responses) + assert len({response.total_tokens for response in responses}) == 1 + assert responses[0].total_tokens > 0 + + await proxy_server_module.token_counter(TokenCountRequest(model="self-hosted-pinned", prompt="count me once")) + await proxy_server_module.token_counter(TokenCountRequest(model="self-hosted-private", prompt="count me once")) + assert loads == [ + ("my-org/tokenizer", "main", None), + ("my-org/tokenizer", "v2", None), + ("my-org/tokenizer", "main", "hf_test_token"), + ] + finally: + litellm.utils._select_custom_tokenizer_helper.cache_clear()