llm_server/llm/__init__.py

@@ -1,15 +1,15 @@
 from llm_server.cluster.cluster_config import cluster_config
 from llm_server.llm import oobabooga, vllm
-from llm_server.custom_redis import redis
 
 
 def get_token_count(prompt: str, backend_url: str):
-    if not isinstance(prompt, str):
-        print(prompt)
-
-    assert isinstance(prompt, str)
     assert isinstance(backend_url, str)
 
+    if not prompt:
+        # The tokenizers have issues when the prompt is None.
+        return 0
+    assert isinstance(prompt, str)
+
     backend_mode = cluster_config.get_backend(backend_url)['mode']
     if backend_mode == 'vllm':
         return vllm.tokenize(prompt, backend_url)

llm_server/llm/vllm/tokenize.py

@@ -9,16 +9,17 @@ from llm_server.cluster.cluster_config import cluster_config
 
 def tokenize(prompt: str, backend_url: str) -> int:
     assert backend_url
-    assert isinstance(prompt, str)
     assert isinstance(backend_url, str)
 
+    if not prompt:
+        # The tokenizers have issues when the prompt is None.
+        return 0
+    assert isinstance(prompt, str)
+
     # The backend could have died between when the request was
     # submitted and now, so let's double check it's still online.
     backend_url = cluster_config.validate_backend(backend_url)
 
-    if not prompt:
-        # The tokenizers have issues when the prompt is None.
-        return 0
     tokenizer = tiktoken.get_encoding("cl100k_base")
 
     # Split the prompt into 300 character chunks