浏览代码

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

Signed-off-by: -LAN- <laipz8200@outlook.com>
-LAN- 6 月之前
父节点
当前提交
6f5a8a33d9
共有 1 个文件被更改,包括 4 次插入4 次删除
  1. 4 4
      api/core/model_runtime/model_providers/__base/tokenizers/gpt2_tokenzier.py

+ 4 - 4
api/core/model_runtime/model_providers/__base/tokenizers/gpt2_tokenzier.py

@@ -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)
-        result = future.get(block=True)
+        future = _executor.submit(GPT2Tokenizer._get_num_tokens_by_gpt2, text)
+        result = future.result()
         return cast(int, result)
         return cast(int, result)
 
 
     @staticmethod
     @staticmethod