From 07c5f136f15899c9378d015d9cb7c23d26d2a3f5 Mon Sep 17 00:00:00 2001
From: Krish Dholakia <krrishdholakia@gmail.com>
Date: Tue, 7 Jan 2025 22:37:09 -0800
Subject: [PATCH] fix(utils.py): fix select tokenizer for custom tokenizer
 (#7599)

* fix(utils.py): fix select tokenizer for custom tokenizer

* fix(router.py): fix 'utils/token_counter' endpoint
---
 litellm/proxy/_new_secret_config.yaml | 14 +++++++++++---
 litellm/proxy/proxy_server.py         |  1 +
 litellm/router.py                     |  2 +-
 litellm/utils.py                      |  6 +++---
 4 files changed, 16 insertions(+), 7 deletions(-)

diff --git a/litellm/proxy/_new_secret_config.yaml b/litellm/proxy/_new_secret_config.yaml
index 91dca93acf58..280109dbc18b 100644
--- a/litellm/proxy/_new_secret_config.yaml
+++ b/litellm/proxy/_new_secret_config.yaml
@@ -10,13 +10,21 @@ model_list:
       api_key: os.environ/OPENAI_API_KEY
   - model_name: chatbot_actions
     litellm_params:
-      model: langfuse/openai-gpt-3.5-turbo
+      model: langfuse/azure/gpt-4o
+      api_base: "os.environ/AZURE_API_BASE"
+      api_key: "os.environ/AZURE_API_KEY"
       tpm: 1000000
       prompt_id: "jokes"
-  - model_name: openai-gpt-3.5-turbo
+  - model_name: openai-deepseek
     litellm_params:
-      model: openai/gpt-3.5-turbo
+      model: deepseek/deepseek-chat
       api_key: os.environ/OPENAI_API_KEY
+    model_info:
+      access_groups: ["restricted-models"]
+      custom_tokenizer: 
+        identifier: deepseek-ai/DeepSeek-V3-Base
+        revision: main
+        auth_token: os.environ/HUGGINGFACE_API_KEY
 
 
 litellm_settings:
diff --git a/litellm/proxy/proxy_server.py b/litellm/proxy/proxy_server.py
index 5890787bf17d..54b3314f833b 100644
--- a/litellm/proxy/proxy_server.py
+++ b/litellm/proxy/proxy_server.py
@@ -5606,6 +5606,7 @@ async def token_counter(request: TokenCountRequest):
     _tokenizer_used = litellm.utils._select_tokenizer(
         model=model_to_use, custom_tokenizer=custom_tokenizer
     )
+
     tokenizer_used = str(_tokenizer_used["type"])
     total_tokens = token_counter(
         model=model_to_use,
diff --git a/litellm/router.py b/litellm/router.py
index 1da1a9773041..bfa7708854e3 100644
--- a/litellm/router.py
+++ b/litellm/router.py
@@ -4379,7 +4379,7 @@ def get_router_model_info(
                         pass
 
         ## GET LITELLM MODEL INFO - raises exception, if model is not mapped
-        if not model.startswith(custom_llm_provider):
+        if not model.startswith("{}/".format(custom_llm_provider)):
             model_info_name = "{}/{}".format(custom_llm_provider, model)
         else:
             model_info_name = model
diff --git a/litellm/utils.py b/litellm/utils.py
index 00b57cef17ea..6051a8732724 100644
--- a/litellm/utils.py
+++ b/litellm/utils.py
@@ -1278,12 +1278,12 @@ def _select_tokenizer(
     model: str, custom_tokenizer: Optional[CustomHuggingfaceTokenizer] = None
 ):
     if custom_tokenizer is not None:
-        custom_tokenizer = Tokenizer.from_pretrained(
-            custom_tokenizer["identifier"],
+        _tokenizer = create_pretrained_tokenizer(
+            identifier=custom_tokenizer["identifier"],
             revision=custom_tokenizer["revision"],
             auth_token=custom_tokenizer["auth_token"],
         )
-        return {"type": "huggingface_tokenizer", "tokenizer": custom_tokenizer}
+        return _tokenizer
     return _select_tokenizer_helper(model=model)