@@ -0,0 +1,354 @@
import re
import threading
from toolbox import update_ui , get_conf
from multiprocessing import Process , Pipe
import numpy as np
from onnxruntime import InferenceSession , SessionOptions
from sentencepiece import SentencePieceProcessor
# 模型来源 K024/ChatGLM-6b-onnx-u8s8
global glm_onnx_handle
glm_onnx_handle = None
load_message = " ChatGLM_onnx尚未加载，加载需要一段时间。注意，取决于`config.py`的配置，ChatGLM_onnx消耗大量的内存（CPU）或显存（GPU），也许会导致低配(内存<8GB）计算机卡死 …… "
# Default paths
tokenizer_path = " YOUR/TOKENIZER_PATH/sentencepiece.model "
onnx_model_path = " YOUR/TOKENIZER_PATH/chatglm-6b-int8.onnx "
# Currently `MatMulInteger` and `DynamicQuantizeLinear` are only supported on CPU,
# although they are documented as supported on CUDA.
providers = [ " CPUExecutionProvider " ]
# if torch.cuda.is_available():
# providers = ["CUDAExecutionProvider"] + providers
#################################################################################
class GetGLMHandle ( Process ) :
def __init__ ( self ) :
super ( ) . __init__ ( daemon = True )
self . parent , self . child = Pipe ( )
self . ChatGLM_onnx_model = None # tokenizer_path
self . ChatGLM_onnx_tokenizer = None # onnx_model_path
self . info = " "
self . success = True
self . check_dependency ( )
self . start ( )
self . threadLock = threading . Lock ( )
def check_dependency ( self ) :
try :
import sentencepiece
self . info = " 依赖检测通过 "
self . success = True
except :
self . info = " 缺少ChatGLM_onnx的依赖，如果要使用ChatGLM_onnx，除了基础的pip依赖以外，您还需要运行`pip install -r request_llm/requirements_ChatGLM_onnx.txt`安装ChatGLM_onnx的依赖。 "
self . success = False
def ready ( self ) :
return self . ChatGLM_onnx_model is not None
def run ( self ) :
# 子进程执行
# 第一次运行，加载参数
retry = 0
while True :
try :
if self . ChatGLM_onnx_model is None :
# Initialize the ChatGLMModel and ChatGLMTokenizer
self . ChatGLM_onnx_model = ChatGLMModel ( )
self . ChatGLM_onnx_tokenizer = ChatGLMTokenizer ( )
break
else :
break
except :
retry + = 1
if retry > 3 :
self . child . send ( ' [Local Message] Call ChatGLM_onnx fail 不能正常加载ChatGLM_onnx的参数。 ' )
raise RuntimeError ( " 不能正常加载ChatGLM_onnx的参数！ " )
while True :
# 进入任务等待状态
kwargs = self . child . recv ( )
# 收到消息，开始请求
try :
# Use the ChatGLMModel and ChatGLMTokenizer to generate a response
response = tuple ( self . ChatGLM_onnx_model . generate_iterate ( kwargs [ ' query ' ] ) )
# Send the output data
self . child . send ( response [ - 1 ] )
except :
from toolbox import trimmed_format_exc
self . child . send ( ' [Local Message] Call ChatGLM_onnx fail. ' + ' \n ``` \n ' + trimmed_format_exc ( ) + ' \n ``` \n ' )
# 请求处理结束，开始下一个循环
self . child . send ( ' [Finish] ' )
def stream_chat ( self , * * kwargs ) :
# 主进程执行
self . threadLock . acquire ( )
self . parent . send ( kwargs )
while True :
res = self . parent . recv ( )
if res != ' [Finish] ' :
yield res
else :
break
self . threadLock . release ( )
#################################################################################
class ChatGLMModel ( ) :
def __init__ ( self , onnx_model_path = onnx_model_path , tokenizer_path = tokenizer_path , profile = False ) - > None :
self . tokenizer = ChatGLMTokenizer ( tokenizer_path )
options = SessionOptions ( )
options . enable_profiling = profile
self . session = InferenceSession ( onnx_model_path , options , providers = providers )
self . eop_token_id = self . tokenizer [ " <eop> " ]
# input & output names
self . past_names = [ f " past_ { name } _ { i } " for i in range ( 28 ) for name in [ " key " , " value " ] ]
self . present_names = [ f " present_ { name } _ { i } " for i in range ( 28 ) for name in [ " key " , " value " ] ]
self . output_names = [ " logits " ] + self . present_names
# default kv_cache for first inference
self . default_past_key_values = {
k : np . zeros ( ( 1 , 0 , 32 , 128 ) , dtype = np . float32 ) for k in self . past_names
}
def prepare_input ( self , prompt : str ) :
input_ids , prefix_mask = self . tokenizer . encode ( prompt )
input_ids = np . array ( [ input_ids ] , dtype = np . longlong )
prefix_mask = np . array ( [ prefix_mask ] , dtype = np . longlong )
return input_ids , prefix_mask , self . default_past_key_values
def sample_next_token ( self , logits : np . ndarray , top_k = 50 , top_p = 0.7 , temperature = 1 ) :
# softmax with temperature
exp_logits = np . exp ( logits / temperature )
probs = exp_logits / np . sum ( exp_logits )
# top k
top_k_idx = np . argsort ( - probs ) [ : top_k ]
top_k_probs = probs [ top_k_idx ]
# top p
cumsum_probs = np . cumsum ( top_k_probs )
top_k_probs [ ( cumsum_probs - top_k_probs ) > top_p ] = 0.0
top_k_probs = top_k_probs / np . sum ( top_k_probs )
# sample
next_token = np . random . choice ( top_k_idx , size = 1 , p = top_k_probs )
return next_token [ 0 ] . item ( )
def generate_iterate ( self , prompt : str , max_generated_tokens = 100 , top_k = 50 , top_p = 0.7 , temperature = 1 ) :
input_ids , prefix_mask , past_key_values = self . prepare_input ( prompt )
output_tokens = [ ]
while True :
inputs = {
" input_ids " : input_ids ,
" prefix_mask " : prefix_mask ,
" use_past " : np . array ( len ( output_tokens ) > 0 ) ,
}
inputs . update ( past_key_values )
logits , * past_key_values = self . session . run ( self . output_names , inputs )
past_key_values = { k : v for k , v in zip ( self . past_names , past_key_values ) }
next_token = self . sample_next_token ( logits [ 0 , - 1 ] , top_k = top_k , top_p = top_p , temperature = temperature )
output_tokens + = [ next_token ]
if next_token == self . eop_token_id or len ( output_tokens ) > max_generated_tokens :
break
input_ids = np . array ( [ [ next_token ] ] , dtype = np . longlong )
prefix_mask = np . concatenate ( [ prefix_mask , np . array ( [ [ 0 ] ] , dtype = np . longlong ) ] , axis = 1 )
yield process_response ( self . tokenizer . decode ( output_tokens ) )
return process_response ( self . tokenizer . decode ( output_tokens ) )
class ChatGLMTokenizer :
def __init__ ( self , vocab_file ) :
assert vocab_file is not None
self . vocab_file = vocab_file
self . special_tokens = [ " [MASK] " , " [gMASK] " , " [sMASK] " , " <unused_0> " , " <sop> " , " <eop> " , " <ENC> " , " <dBLOCK> " ]
self . text_tokenizer = SentencePieceProcessor ( str ( vocab_file ) )
def __len__ ( self ) :
return len ( self . text_tokenizer )
def __getitem__ ( self , key : str ) :
return self . text_tokenizer [ key ]
def preprocess ( self , text : str , linebreak = True , whitespaces = True ) :
if linebreak :
text = text . replace ( " \\ n " , " <n> " )
if whitespaces :
text = text . replace ( " \\ t " , " <|tab|> " )
text = re . sub ( r " { 2,80} " , self . replace_spaces_with_blank , text )
return text
def encode (
self , text : str , text_pair : str = None ,
linebreak = True , whitespaces = True ,
add_dummy_prefix = True , special_tokens = True ,
) - > tuple [ list [ int ] , list [ int ] ] :
"""
text: Text to encode. Bidirectional part with a [gMASK] and an <sop> for causal LM.
text_pair: causal LM part.
linebreak: Whether to encode newline ( \n ) in text.
whitespaces: Whether to encode multiple whitespaces or tab in text, useful for source code encoding.
special_tokens: Whether to encode special token ([MASK], [gMASK], etc.) in text.
add_dummy_prefix: Whether to add dummy blank space in the beginning.
"""
text = self . preprocess ( text , linebreak , whitespaces )
if not add_dummy_prefix :
text = " <n> " + text
tokens = self . text_tokenizer . encode ( text )
prefix_mask = [ 1 ] * len ( tokens )
if special_tokens :
tokens + = [ self . text_tokenizer [ " [gMASK] " ] , self . text_tokenizer [ " <sop> " ] ]
prefix_mask + = [ 1 , 0 ]
if text_pair is not None :
text_pair = self . preprocess ( text_pair , linebreak , whitespaces )
pair_tokens = self . text_tokenizer . encode ( text_pair )
tokens + = pair_tokens
prefix_mask + = [ 0 ] * len ( pair_tokens )
if special_tokens :
tokens + = [ self . text_tokenizer [ " <eop> " ] ]
prefix_mask + = [ 0 ]
return ( tokens if add_dummy_prefix else tokens [ 2 : ] ) , prefix_mask
def decode ( self , text_ids : list [ int ] ) - > str :
text = self . text_tokenizer . decode ( text_ids )
text = text . replace ( " <n> " , " \n " )
text = text . replace ( " <|tab|> " , " \t " )
text = re . sub ( r " < \ |blank_( \ d \ d?) \ |> " , self . replace_blank_with_spaces , text )
return text
def replace_spaces_with_blank ( match : re . Match [ str ] ) :
return f " <|blank_ { len ( match . group ( ) ) } |> "
def replace_blank_with_spaces ( match : re . Match [ str ] ) :
return " " * int ( match . group ( 1 ) )
#################################################################################
def chat_template ( history : list [ tuple [ str , str ] ] , current : str ) :
prompt = " "
chat_round = 0
for question , answer in history :
prompt + = f " [Round { chat_round } ] \n 问： { question } \n 答： { answer } \n "
chat_round + = 1
prompt + = f " [Round { chat_round } ] \n 问： { current } \n 答： "
return prompt
def process_response ( response : str ) :
response = response . strip ( )
response = response . replace ( " [[训练时间]] " , " 2023年 " )
punkts = [
[ " , " , " ， " ] ,
[ " ! " , " ！ " ] ,
[ " : " , " ： " ] ,
[ " ; " , " ； " ] ,
[ " \ ? " , " ？ " ] ,
]
for item in punkts :
response = re . sub ( r " ([ \ u4e00- \ u9fff]) %s " % item [ 0 ] , r " \ 1 %s " % item [ 1 ] , response )
response = re . sub ( r " %s ([ \ u4e00- \ u9fff]) " % item [ 0 ] , r " %s \ 1 " % item [ 1 ] , response )
return response
#################################################################################
def predict_no_ui_long_connection ( inputs , llm_kwargs , history = [ ] , sys_prompt = " " , observe_window = [ ] , console_slience = False ) :
"""
多线程方法
函数的说明请见 request_llm/bridge_all.py
"""
if glm_onnx_handle is None :
glm_onnx_handle = GetGLMHandle ( )
if len ( observe_window ) > = 1 : observe_window [ 0 ] = load_message + " \n \n " + glm_onnx_handle . info
if not glm_onnx_handle . success :
error = glm_onnx_handle . info
glm_onnx_handle = None
raise RuntimeError ( error )
# ChatGLM_onnx doesn't have a sys_prompt interface, so add the prompt to history
history_feedin = [ ]
history_feedin . append ( [ " What can I do? " , sys_prompt ] )
for i in range ( len ( history ) / / 2 ) :
history_feedin . append ( [ history [ 2 * i ] , history [ 2 * i + 1 ] ] )
watch_dog_patience = 5 # Watchdog patience, set to 5 seconds
response = " "
for response in glm_onnx_handle . stream_chat ( query = inputs , history = history_feedin ) :
print ( response )
if len ( observe_window ) > = 1 :
observe_window [ 0 ] = response
if len ( observe_window ) > = 2 :
if ( time . time ( ) - observe_window [ 1 ] ) > watch_dog_patience :
raise RuntimeError ( " 程序终止。 " )
return response
def predict ( inputs , llm_kwargs , plugin_kwargs , chatbot , history = [ ] , system_prompt = ' ' , stream = True , additional_fn = None ) :
"""
单线程方法
函数的说明请见 request_llm/bridge_all.py
"""
chatbot . append ( ( inputs , " " ) )
global glm_onnx_handle
if glm_onnx_handle is None :
glm_onnx_handle = GetGLMHandle ( )
chatbot [ - 1 ] = ( inputs , load_message + " \n \n " + glm_onnx_handle . info )
yield from update_ui ( chatbot = chatbot , history = [ ] )
if not glm_onnx_handle . success :
glm_onnx_handle = None
return
if additional_fn is not None :
import core_functional
importlib . reload ( core_functional ) # Hot-reload prompt
core_functional = core_functional . get_core_functions ( )
if " PreProcess " in core_functional [ additional_fn ] :
inputs = core_functional [ additional_fn ] [ " PreProcess " ] ( inputs )
inputs = core_functional [ additional_fn ] [ " Prefix " ] + inputs + core_functional [ additional_fn ] [ " Suffix " ]
history_feedin = [ ]
history_feedin . append ( [ " What can I do? " , system_prompt ] )
for i in range ( len ( history ) / / 2 ) :
history_feedin . append ( [ history [ 2 * i ] , history [ 2 * i + 1 ] ] )
response = " [Local Message]: 等待ChatGLM_onnx响应中 ... "
for response in glm_onnx_handle . stream_chat ( query = inputs , history = history_feedin ) :
chatbot [ - 1 ] = ( inputs , response )
yield from update_ui ( chatbot = chatbot , history = history )
if response == " [Local Message]: 等待ChatGLM_onnx响应中 ... " :
response = " [Local Message]: ChatGLM_onnx响应异常 ... "
history . extend ( [ inputs , response ] )
yield from update_ui ( chatbot = chatbot , history = history )
ValeriaWong