Skip to content

refactor: use tiktoken for token calculation #12416

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Merged
merged 2 commits into from
Jan 7, 2025
Merged

refactor: use tiktoken for token calculation #12416

Changes from all commits
Commits
Show all changes
2 commits
Select commit Hold shift + click to select a range
477a378
refactor: replace GPT2Tokenizer implementation with tiktoken for impr…
laipz8200 Jan 7, 2025
8212125
fix: remove unnecessary verbose flag from GPT2Tokenizer encode method
laipz8200 Jan 7, 2025
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
27 changes: 15 additions & 12 deletions api/core/model_runtime/model_providers/__base/tokenizers/gpt2_tokenzier.py
View file
Open in desktop
Original file line number Diff line number Diff line change
@@ -1,13 +1,10 @@
from concurrent.futures import ProcessPoolExecutor
from os.path import abspath, dirname, join
from threading import Lock
from typing import Any, cast
from typing import Any

from transformers import GPT2Tokenizer as TransformerGPT2Tokenizer # type: ignore
import tiktoken

_tokenizer: Any = None
_lock = Lock()
_executor = ProcessPoolExecutor(max_workers=1)


class GPT2Tokenizer:
Expand All @@ -17,22 +14,28 @@ def _get_num_tokens_by_gpt2(text: str) -> int:
use gpt2 tokenizer to get num tokens
"""
_tokenizer = GPT2Tokenizer.get_encoder()
tokens = _tokenizer.encode(text, verbose=False)
tokens = _tokenizer.encode(text)
return len(tokens)

@staticmethod
def get_num_tokens(text: str) -> int:
future = _executor.submit(GPT2Tokenizer._get_num_tokens_by_gpt2, text)
result = future.result()
return cast(int, result)
# Because this process needs more cpu resource, we turn this back before we find a better way to handle it.
#
# future = _executor.submit(GPT2Tokenizer._get_num_tokens_by_gpt2, text)
# result = future.result()
# return cast(int, result)
return GPT2Tokenizer._get_num_tokens_by_gpt2(text)

@staticmethod
def get_encoder() -> Any:
global _tokenizer, _lock
with _lock:
if _tokenizer is None:
base_path = abspath(__file__)
gpt2_tokenizer_path = join(dirname(base_path), "gpt2")
_tokenizer = TransformerGPT2Tokenizer.from_pretrained(gpt2_tokenizer_path)
# Try to use tiktoken to get the tokenizer because it is faster
#
_tokenizer = tiktoken.get_encoding("gpt2")
# base_path = abspath(__file__)
# gpt2_tokenizer_path = join(dirname(base_path), "gpt2")
# _tokenizer = TransformerGPT2Tokenizer.from_pretrained(gpt2_tokenizer_path)

return _tokenizer
Loading