local-llm-server/llm_server/llm/__init__.py

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