Update app.py

app.py

@@ -156,22 +156,30 @@ with gr.Blocks() as demo:
             }
             
             start_time = time.time()
+            keep_streaming = True
             with requests.post(API_URL, headers=HEADERS, data=json.dumps(data), stream=True, timeout=30) as r:
-                for response in r.iter_lines():
-                    if len(response) > 0:
-                        text = response.decode()
-                        if text != "data: [DONE]":
-                            if text.startswith("data: "):
-                                text = text[5:]
-                            delta = json.loads(text)["choices"][0]["text"]
-    
-                            if history[-1][1] is None:
-                                history[-1][1] = delta
-                            else:
-                                history[-1][1] += delta
+                for line in r.iter_lines():
+                    if line and keep_streaming:
+                        assert r.status_code == 200
+                        json_response = json.loads(line)
+        
+                        try:
+                            if "fragment" not in json_response["result"]:
+                                keep_streaming = False
+                                break
+        
+                            delta = json_response["result"]["fragment"]["data"]["text"]
+                        except Exception as e:
+                            raise e
+
+                        if history[-1][1] is None:
+                            history[-1][1] = delta
+                        else:
+                            history[-1][1] += delta
+                        yield history
 
-                            yield history
                     if time.time() - start_time > MAX_SEC:
+                        keep_streaming = False
                         break
                     
             if history[-1][1].endswith('</s>'):