Skip to content

Commit 1886bb1

Browse files
committed
refactor: replace gevent threadpool with ProcessPoolExecutor in GPT2Tokenizer
Signed-off-by: -LAN- <[email protected]>
1 parent 3d150c3 commit 1886bb1

File tree

1 file changed

+4
-4
lines changed

1 file changed

+4
-4
lines changed

api/core/model_runtime/model_providers/__base/tokenizers/gpt2_tokenzier.py

Lines changed: 4 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -1,13 +1,13 @@
1+
from concurrent.futures import ProcessPoolExecutor
12
from os.path import abspath, dirname, join
23
from threading import Lock
34
from typing import Any, cast
45

5-
import gevent.threadpool # type: ignore
66
from transformers import GPT2Tokenizer as TransformerGPT2Tokenizer # type: ignore
77

88
_tokenizer: Any = None
99
_lock = Lock()
10-
_pool = gevent.threadpool.ThreadPool(1)
10+
_executor = ProcessPoolExecutor(max_workers=1)
1111

1212

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

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

2929
@staticmethod

0 commit comments

Comments
 (0)