20 lines
606 B
Python
20 lines
606 B
Python
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)
|
|
|
|
backend_mode = cluster_config.get_backend(backend_url)['mode']
|
|
if backend_mode == 'vllm':
|
|
return vllm.tokenize(prompt, backend_url)
|
|
elif backend_mode == 'ooba':
|
|
return oobabooga.tokenize(prompt)
|
|
else:
|
|
raise Exception(backend_mode)
|