mirror of
				https://github.com/langgenius/dify.git
				synced 2025-11-03 20:33:00 +00:00 
			
		
		
		
	refactor: replace gevent threadpool with ProcessPoolExecutor in GPT2Tokenizer (#12316)
Signed-off-by: -LAN- <laipz8200@outlook.com>
This commit is contained in:
		
							parent
							
								
									52b2559a14
								
							
						
					
					
						commit
						6f5a8a33d9
					
				@ -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
 | 
			
		||||
 | 
			
		||||
		Loading…
	
	
			
			x
			
			
		
	
		Reference in New Issue
	
	Block a user