update

2023-03-24 13:12:25 +08:00
parent ecdeda8e92
commit 44b40ff726
5 changed files with 176 additions and 141 deletions
--- a/predict.py
+++ b/predict.py
@@ -6,11 +6,12 @@ import logging
 import traceback
 import requests
 import importlib
+from colorful import *

 # config_private.py放自己的秘密如API和代理网址
 # 读取时首先看是否存在私密的config_private配置文件（不受git管控），如果有，则覆盖原config文件
-try: from config_private import proxies, API_URL, API_KEY, TIMEOUT_SECONDS, MAX_RETRY
-except: from config import proxies, API_URL, API_KEY, TIMEOUT_SECONDS, MAX_RETRY
+try: from config_private import proxies, API_URL, API_KEY, TIMEOUT_SECONDS, MAX_RETRY, LLM_MODEL
+except: from config import proxies, API_URL, API_KEY, TIMEOUT_SECONDS, MAX_RETRY, LLM_MODEL

 timeout_bot_msg = '[local] Request timeout, network error. please check proxy settings in config.py.'

@@ -23,51 +24,12 @@ def get_full_error(chunk, stream_response):
    return chunk

 def predict_no_ui(inputs, top_p, temperature, history=[]):
-    messages = [{"role": "system", "content": ""}]
-
-    # 
-    chat_counter = len(history) // 2
-    if chat_counter > 0:
-        for index in range(0, 2*chat_counter, 2):
-            what_i_have_asked = {}
-            what_i_have_asked["role"] = "user"
-            what_i_have_asked["content"] = history[index]
-            what_gpt_answer = {}
-            what_gpt_answer["role"] = "assistant"
-            what_gpt_answer["content"] = history[index+1]
-            if what_i_have_asked["content"] != "":
-                messages.append(what_i_have_asked)
-                messages.append(what_gpt_answer)
-            else:
-                messages[-1]['content'] = what_gpt_answer['content']
-
-    what_i_ask_now = {}
-    what_i_ask_now["role"] = "user"
-    what_i_ask_now["content"] = inputs
-    messages.append(what_i_ask_now)
-
-    # messages
-    payload = {
-        "model": "gpt-3.5-turbo",
-        # "model": "gpt-4",
-        "messages": messages, 
-        "temperature": temperature,  # 1.0,
-        "top_p": top_p,  # 1.0,
-        "n": 1,
-        "stream": False,
-        "presence_penalty": 0,
-        "frequency_penalty": 0,
-    }
-    
-    headers = {
-        "Content-Type": "application/json",
-        "Authorization": f"Bearer {API_KEY}"
-    }
+    headers, payload = generate_payload(inputs, top_p, temperature, history, system_prompt="", stream=False)

    retry = 0
    while True:
        try:
-            # make a POST request to the API endpoint using the requests.post method, passing in stream=True
+            # make a POST request to the API endpoint, stream=False
            response = requests.post(API_URL, headers=headers, proxies=proxies,
                                    json=payload, stream=False, timeout=TIMEOUT_SECONDS*2); break
        except TimeoutError as e:
@@ -84,9 +46,7 @@ def predict_no_ui(inputs, top_p, temperature, history=[]):
        raise ConnectionAbortedError("Json解析不合常规，可能是文本过长" + response.text)


-
-
-def predict(inputs, top_p, temperature, chatbot=[], history=[], system_prompt='', retry=False, 
+def predict(inputs, top_p, temperature, chatbot=[], history=[], system_prompt='', 
            stream = True, additional_fn=None):

    if additional_fn is not None:
@@ -101,60 +61,13 @@ def predict(inputs, top_p, temperature, chatbot=[], history=[], system_prompt=''
        chatbot.append((inputs, ""))
        yield chatbot, history, "等待响应"

-    headers = {
-        "Content-Type": "application/json",
-        "Authorization": f"Bearer {API_KEY}"
-    }
-
-    chat_counter = len(history) // 2
-
-    print(f"chat_counter - {chat_counter}")
-
-    messages = [{"role": "system", "content": system_prompt}]
-    if chat_counter:
-        for index in range(0, 2*chat_counter, 2):
-            what_i_have_asked = {}
-            what_i_have_asked["role"] = "user"
-            what_i_have_asked["content"] = history[index]
-            what_gpt_answer = {}
-            what_gpt_answer["role"] = "assistant"
-            what_gpt_answer["content"] = history[index+1]
-            if what_i_have_asked["content"] != "":
-                if not (what_gpt_answer["content"] != "" or retry): continue
-                if what_gpt_answer["content"] == timeout_bot_msg: continue
-                messages.append(what_i_have_asked)
-                messages.append(what_gpt_answer)
-            else:
-                messages[-1]['content'] = what_gpt_answer['content']
-
-    if retry and chat_counter:
-        messages.pop()
-    else:
-        what_i_ask_now = {}
-        what_i_ask_now["role"] = "user"
-        what_i_ask_now["content"] = inputs
-        messages.append(what_i_ask_now)
-        chat_counter += 1
-
-    # messages
-    payload = {
-        "model": "gpt-3.5-turbo",
-        # "model": "gpt-4",
-        "messages": messages, 
-        "temperature": temperature,  # 1.0,
-        "top_p": top_p,  # 1.0,
-        "n": 1,
-        "stream": stream,
-        "presence_penalty": 0,
-        "frequency_penalty": 0,
-    }
-
-    history.append(inputs)
+    headers, payload = generate_payload(inputs, top_p, temperature, history, system_prompt, stream)
+    history.append(inputs); history.append(" ")

    retry = 0
    while True:
        try:
-            # make a POST request to the API endpoint using the requests.post method, passing in stream=True
+            # make a POST request to the API endpoint, stream=True
            response = requests.post(API_URL, headers=headers, proxies=proxies,
                                    json=payload, stream=True, timeout=TIMEOUT_SECONDS);break
        except:
@@ -164,37 +77,30 @@ def predict(inputs, top_p, temperature, chatbot=[], history=[], system_prompt=''
            yield chatbot, history, "请求超时"+retry_msg
            if retry > MAX_RETRY: raise TimeoutError

-    token_counter = 0
-    partial_words = ""
-
-    counter = 0
+    gpt_replying_buffer = ""
+    
+    is_head_of_the_stream = True
    if stream:
        stream_response =  response.iter_lines()
        while True:
            chunk = next(stream_response)
-            if chunk == b'data: [DONE]':
-                break
-
-            if counter == 0:
-                counter += 1
-                continue
-            counter += 1
-            # check whether each line is non-empty
+            # print(chunk.decode()[6:])
+            if is_head_of_the_stream:
+                is_head_of_the_stream = False; continue
+            
            if chunk:
-                # decode each line as response data is in bytes
                try:
                    if len(json.loads(chunk.decode()[6:])['choices'][0]["delta"]) == 0:
-                        logging.info(f'[response] {chatbot[-1][-1]}')
+                        # 判定为数据流的结束，gpt_replying_buffer也写完了
+                        logging.info(f'[response] {gpt_replying_buffer}')
                        break
+                    # 处理数据流的主体
                    chunkjson = json.loads(chunk.decode()[6:])
                    status_text = f"finish_reason: {chunkjson['choices'][0]['finish_reason']}"
-                    partial_words = partial_words + json.loads(chunk.decode()[6:])['choices'][0]["delta"]["content"]
-                    if token_counter == 0:
-                        history.append(" " + partial_words)
-                    else:
-                        history[-1] = partial_words
+                    # 如果这里抛出异常，一般是文本过长，详情见get_full_error的输出
+                    gpt_replying_buffer = gpt_replying_buffer + json.loads(chunk.decode()[6:])['choices'][0]["delta"]["content"]
+                    history[-1] = gpt_replying_buffer
                    chatbot[-1] = (history[-2], history[-1])
-                    token_counter += 1
                    yield chatbot, history, status_text

                except Exception as e:
@@ -207,4 +113,48 @@ def predict(inputs, top_p, temperature, chatbot=[], history=[], system_prompt=''
                    yield chatbot, history, "Json解析不合常规，很可能是文本过长" + error_msg
                    return

+def generate_payload(inputs, top_p, temperature, history, system_prompt, stream):
+    headers = {
+        "Content-Type": "application/json",
+        "Authorization": f"Bearer {API_KEY}"
+    }
+
+    conversation_cnt = len(history) // 2
+
+    messages = [{"role": "system", "content": system_prompt}]
+    if conversation_cnt:
+        for index in range(0, 2*conversation_cnt, 2):
+            what_i_have_asked = {}
+            what_i_have_asked["role"] = "user"
+            what_i_have_asked["content"] = history[index]
+            what_gpt_answer = {}
+            what_gpt_answer["role"] = "assistant"
+            what_gpt_answer["content"] = history[index+1]
+            if what_i_have_asked["content"] != "":
+                if what_gpt_answer["content"] == "": continue
+                if what_gpt_answer["content"] == timeout_bot_msg: continue
+                messages.append(what_i_have_asked)
+                messages.append(what_gpt_answer)
+            else:
+                messages[-1]['content'] = what_gpt_answer['content']
+
+    what_i_ask_now = {}
+    what_i_ask_now["role"] = "user"
+    what_i_ask_now["content"] = inputs
+    messages.append(what_i_ask_now)
+
+    payload = {
+        "model": LLM_MODEL,
+        "messages": messages, 
+        "temperature": temperature,  # 1.0,
+        "top_p": top_p,  # 1.0,
+        "n": 1,
+        "stream": stream,
+        "presence_penalty": 0,
+        "frequency_penalty": 0,
+    }
+    
+    print(f" {LLM_MODEL} : {conversation_cnt} : {inputs}")
+    return headers,payload
+