Selaa lähdekoodia

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

Signed-off-by: -LAN- <laipz8200@outlook.com>
-LAN- 3 kuukautta sitten
vanhempi
commit
6f5a8a33d9

+ 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