local-llm-server/llm_server/llm/vllm/tokenize.py

19 lines
617 B
Python

import requests
import tiktoken
from llm_server import opts
def tokenize(prompt: str) -> int:
tokenizer = tiktoken.get_encoding("cl100k_base")
if not prompt:
# The tokenizers have issues when the prompt is None.
return 0
try:
r = requests.post(f'{opts.backend_url}/tokenize', json={'input': prompt}, verify=opts.verify_ssl, timeout=opts.backend_generate_request_timeout)
j = r.json()
return j['length']
except Exception as e:
print(f'Failed to tokenize using VLLM -', f'{e.__class__.__name__}: {e}')
return len(tokenizer.encode(prompt)) + 10