Merge branch 'master' of https://github.com/ValeriaWong/chatgpt_academic into ValeriaWong-master

config.py

@@ -70,7 +70,7 @@ MAX_RETRY = 2
 
 # 模型选择是 (注意: LLM_MODEL是默认选中的模型, 它*必须*被包含在AVAIL_LLM_MODELS列表中 )
 LLM_MODEL = "gpt-3.5-turbo" # 可选 ↓↓↓
-AVAIL_LLM_MODELS = ["gpt-3.5-turbo-16k", "gpt-3.5-turbo", "azure-gpt-3.5", "api2d-gpt-3.5-turbo", "gpt-4", "api2d-gpt-4", "chatglm", "moss", "newbing", "stack-claude"]
+AVAIL_LLM_MODELS = ["gpt-3.5-turbo-16k", "gpt-3.5-turbo", "azure-gpt-3.5", "api2d-gpt-3.5-turbo", "gpt-4", "api2d-gpt-4", "chatglm","chatglm_onnx","moss", "newbing", "stack-claude"]
 # P.S. 其他可用的模型还包括 ["gpt-3.5-turbo-0613", "gpt-3.5-turbo-16k-0613", "claude-1-100k", "claude-2", "internlm", "jittorllms_rwkv", "jittorllms_pangualpha", "jittorllms_llama"]
 
 

request_llm/bridge_all.py

@@ -19,6 +19,8 @@ from .bridge_chatgpt import predict as chatgpt_ui
 from .bridge_chatglm import predict_no_ui_long_connection as chatglm_noui
 from .bridge_chatglm import predict as chatglm_ui
 
+from .bridge_chatglm_onnx import predict_no_ui_long_connection as chatglm_onnx_noui
+from .bridge_chatglm_onnx import predict as chatglm_onnx_ui
 # from .bridge_tgui import predict_no_ui_long_connection as tgui_noui
 # from .bridge_tgui import predict as tgui_ui
 
@@ -164,7 +166,14 @@ model_info = {
         "tokenizer": tokenizer_gpt35,
         "token_cnt": get_token_num_gpt35,
     },
-
+    "chatglm_onnx": {
+        "fn_with_ui": chatglm_onnx_ui,
+        "fn_without_ui": chatglm_onnx_noui,
+        "endpoint": None,
+        "max_token": 1024,
+        "tokenizer": tokenizer_gpt35,
+        "token_cnt": get_token_num_gpt35,
+    },
 }
 
 

request_llm/bridge_chatglm_onnx.py

@@ -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)
+
+
+
+

request_llm/requirements_chatglm_onnx.txt

@@ -0,0 +1,11 @@
+protobuf
+transformers==4.27.1
+cpm_kernels
+torch>=1.10
+mdtex2html
+sentencepiece
+numpy
+onnxruntime
+sentencepiece
+streamlit
+streamlit-chat