outman/dify
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

refactor: replace gevent threadpool with ProcessPoolExecutor in GPT2Tokenizer (#12316)

Browse Source 
Signed-off-by: -LAN- <laipz8200@outlook.com>
This commit is contained in:
-LAN- 2025-01-03 09:13:18 +08:00 committed by GitHub
parent 52b2559a14
commit 6f5a8a33d9
No known key found for this signature in database
GPG Key ID: B5690EEEBB952194
1 changed files with 4 additions and 4 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

8
api/core/model_runtime/model_providers/__base/tokenizers/gpt2_tokenzier.py
View File

@ -1,13 +1,13 @@
from concurrent.futures import ProcessPoolExecutor
from os.path import abspath, dirname, join from os.path import abspath, dirname, join
from threading import Lock from threading import Lock
from typing import Any, cast from typing import Any, cast
import gevent.threadpool # type: ignore
from transformers import GPT2Tokenizer as TransformerGPT2Tokenizer # type: ignore from transformers import GPT2Tokenizer as TransformerGPT2Tokenizer # type: ignore
_tokenizer: Any = None _tokenizer: Any = None
_lock = Lock() _lock = Lock()
_pool = gevent.threadpool.ThreadPool(1) _executor = ProcessPoolExecutor(max_workers=1)
class GPT2Tokenizer: class GPT2Tokenizer:
@ -22,8 +22,8 @@ class GPT2Tokenizer:
@staticmethod @staticmethod
def get_num_tokens(text: str) -> int: def get_num_tokens(text: str) -> int:
future = _pool.spawn(GPT2Tokenizer._get_num_tokens_by_gpt2, text) future = _executor.submit(GPT2Tokenizer._get_num_tokens_by_gpt2, text)
result = future.get(block=True) result = future.result()
return cast(int, result) return cast(int, result)
@staticmethod @staticmethod