Update app.py

app.py

@@ -148,6 +148,40 @@ with gr.Blocks() as demo:
         return "", history + [[user_message, None]]
 
 
+    def connect_server(data):
+        for _ in range(3):
+            s = requests.Session()
+            r = s.post(API_URL, headers=HEADERS, json=data, stream=True, timeout=30)
+            time.sleep(1)
+            if r.status_code == 200:
+                return r
+        return None
+
+
+    def stream_response_from_server(r):
+        # start_time = time.time()
+        keep_streaming = True
+        for line in r.iter_lines():
+            time.sleep(0.001)
+            # if time.time() - start_time > MAX_SEC:
+            #     keep_streaming = False
+            #     break
+            
+            if line and keep_streaming:
+                if r.status_code != 200:
+                    continue
+                json_response = json.loads(line)
+
+                if "fragment" not in json_response["result"]:
+                    keep_streaming = False
+                    break
+
+                delta = json_response["result"]["fragment"]["data"]["text"]
+                yield delta
+
+            # start_time = time.time()
+
+    
     def bot(history, max_new_tokens, temperature, top_p, system_prompt):
         chat_data = []
         system_prompt = system_prompt.strip()
@@ -163,7 +197,7 @@ with gr.Blocks() as demo:
         if len(message) > MAX_INPUT_LENGTH:
             raise Exception()
 
-        response = ''
+        response = '[ERROR]'
         if refusal_condition(history[-1][0]):
             history = [['[安全拒答啟動]', '[安全拒答啟動] 請清除再開啟對話']]
             response = '[REFUSAL]'
@@ -180,41 +214,23 @@ with gr.Blocks() as demo:
                 }
             }
             
-            start_time = time.time()
-            keep_streaming = True
-            s = requests.Session()
-            with s.post(API_URL, headers=HEADERS, json=data, stream=True, timeout=30) as r:
-                for line in r.iter_lines():
-                    time.sleep(0.005)
-                    if time.time() - start_time > MAX_SEC:
-                        keep_streaming = False
-                        break
-                    
-                    if line and keep_streaming:
-                        if r.status_code != 200:
-                            continue
-                        json_response = json.loads(line)
-        
-                        if "fragment" not in json_response["result"]:
-                            keep_streaming = False
-                            break
+            r = connect_server(data)
+            if r is not None:
+                for delta in stream_response_from_server(r):
+                    if history[-1][1] is None:
+                        history[-1][1] = ''
+                    history[-1][1] += delta
+                    yield history
+
+                if history[-1][1].endswith('</s>'):
+                    history[-1][1] = history[-1][1][:-4]
+                    yield history
+    
+                response = history[-1][1]
     
-                        delta = json_response["result"]["fragment"]["data"]["text"]
-                       
-                        if history[-1][1] is None:
-                            history[-1][1] = ''
-                        history[-1][1] += delta
-                        yield history
-
-            if history[-1][1].endswith('</s>'):
-                history[-1][1] = history[-1][1][:-4]
-                yield history
-
-            response = history[-1][1]
-
-            if refusal_condition(history[-1][1]):
-                history[-1][1] = history[-1][1] + '\n\n**[免責聲明: Breeze-7B-Instruct 和 Breeze-7B-Instruct-64k 並未針對問答進行安全保護，因此語言模型的任何回應不代表 MediaTek Research 立場。]**'
-                yield history
+                if refusal_condition(history[-1][1]):
+                    history[-1][1] = history[-1][1] + '\n\n**[免責聲明: Breeze-7B-Instruct 和 Breeze-7B-Instruct-64k 並未針對問答進行安全保護，因此語言模型的任何回應不代表 MediaTek Research 立場。]**'
+                    yield history
 
         print('== Record ==\nQuery: {query}\nResponse: {response}'.format(query=repr(message), response=repr(history[-1][1])))
         insert_to_db(message, response, float(temperature), float(top_p))