This commit is contained in:
Cyberes 2023-10-05 18:07:59 -06:00
parent 3d0a5cf0a2
commit 5a61bdccd4
2 changed files with 10 additions and 9 deletions

View File

@ -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)

View File

@ -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