Преглед изворни кода

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

Signed-off-by: -LAN- <laipz8200@outlook.com>
-LAN- пре 3 месеци
родитељ
комит
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 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:
@@ -22,8 +22,8 @@ class GPT2Tokenizer:
 
     @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