Skip to content

Commit

Permalink
refactor: replace gevent threadpool with ProcessPoolExecutor in GPT2T…
Browse files Browse the repository at this point in the history
…okenizer

Signed-off-by: -LAN- <laipz8200@outlook.com>
  • Loading branch information
laipz8200 committed Jan 2, 2025
1 parent 3d150c3 commit 1886bb1
Showing 1 changed file with 4 additions and 4 deletions.
Original file line number Diff line number Diff line change
@@ -1,13 +1,13 @@
from concurrent.futures import ProcessPoolExecutor
from os.path import abspath, dirname, join
from threading import Lock
from typing import Any, cast

import gevent.threadpool # type: ignore
from transformers import GPT2Tokenizer as TransformerGPT2Tokenizer # type: ignore

_tokenizer: Any = None
_lock = Lock()
_pool = gevent.threadpool.ThreadPool(1)
_executor = ProcessPoolExecutor(max_workers=1)


class GPT2Tokenizer:
Expand All @@ -22,8 +22,8 @@ def _get_num_tokens_by_gpt2(text: str) -> int:

@staticmethod
def get_num_tokens(text: str) -> int:
future = _pool.spawn(GPT2Tokenizer._get_num_tokens_by_gpt2, text)
result = future.get(block=True)
future = _executor.submit(GPT2Tokenizer._get_num_tokens_by_gpt2, text)
result = future.result()
return cast(int, result)

@staticmethod
Expand Down

0 comments on commit 1886bb1

Please sign in to comment.