Using Gradio's exception handling instead of custom error textarea.

chat.py

@@ -4,6 +4,7 @@
 import traceback
 import gradio as gr
 import chat_client
+import json
 
 CHAT_URL='ws://chat.petals.ml/api/v2/generate'
 #CHAT_URL='ws://localhost:8000/api/v2/generate'
@@ -20,9 +21,8 @@ def generate(state, prompt, model, context, output, *args):
     state['generate'] = True
 
     try:
-        for x in _generate(state, prompt, model, context, output, *args):
-            yield x
-    except BrokenPipeError:
+        yield from _generate(state, prompt, model, context, output, *args)
+    except (json.decoder.JSONDecodeError, BrokenPipeError):
         # Broken session, try to renew
         # TODO This is a bit fragile because of recursive call...
         print("Retrying session...")
@@ -47,9 +47,7 @@ def _generate(state, prompt, model, context, output, endseq, max_length,
             state['model'] = model
         except Exception:
             print(traceback.format_exc())
-            yield state, state['history'], prompt, output, \
-                gr.update(visible=True, value=traceback.format_exc())
-            return
+            raise gr.Error(traceback.format_exc())
     else:
         context = ''
 
@@ -99,7 +97,7 @@ def _generate(state, prompt, model, context, output, endseq, max_length,
     output += prompt2
 
     # Update widgets even before we get the first response
-    yield state, state['history'] + [[prompt, '']], None, output, gr.update(visible=False)
+    yield state, state['history'] + [[prompt, '']], None, output
 
     orig_history = state['history']
     new_line = ''
@@ -115,7 +113,7 @@ def _generate(state, prompt, model, context, output, endseq, max_length,
 
             if not state['generate']:
                 client.close_session()
-                yield state, [], None, '', ''
+                yield state, [], None, ''
                 # Stopping generation
                 return
 
@@ -129,16 +127,16 @@ def _generate(state, prompt, model, context, output, endseq, max_length,
                 if len(spl) > 1:
                     state['history'] = orig_history + [[prompt, new_line]]
                     output += new_line
-                    yield state, state['history'], None, output, ''
+                    yield state, state['history'], None, output
                     # Stopping generation
                     return
     
             # Keep original history untouched as we're adding just
             # a chunks at one moment.
             state['history'] = orig_history + [[prompt, new_line]]
-            yield state, state['history'], None, output, ''
+            yield state, state['history'], None, output
 
-    except BrokenPipeError:
+    except (json.decoder.JSONDecodeError, BrokenPipeError):
         # Session was interrupted
         # Handled in upstream func
         client.close_session()
@@ -153,15 +151,12 @@ def _generate(state, prompt, model, context, output, endseq, max_length,
         state['model'] = None
 
         print(traceback.format_exc())
-        # TODO Store errors outside output log
-        yield state, state['history'], prompt, output, \
-            gr.update(visible=True, value=traceback.format_exc())
-        return
+        raise gr.Error(traceback.format_exc())
 
 def reset(state):
     """Resets the session and clears the chat window."""
     state.update(EMPTY_STATE)
-    return state, [], '', gr.update(visible=False, value='')
+    return state, [], ''
 
 with gr.Blocks() as iface_chat:
     gr.Markdown("""**Let's talk to Bloom in a chat!**""")
@@ -198,7 +193,6 @@ with gr.Blocks() as iface_chat:
     chat = gr.Chatbot(label='Chat window')
     prompt = gr.Textbox(show_label=False, label='Prompt',
             placeholder="Prompt Here and press Enter...").style(container=False)
-    error = gr.Textbox(label="Error log", visible=False, elem_id="error")
 
     with gr.Row():
         button_generate = gr.Button("Generate")
@@ -212,11 +206,11 @@ with gr.Blocks() as iface_chat:
 
     inputs = [state, prompt, model, context, output, endseq,
         max_length, do_sample, top_k, top_p, temperature]
-    outputs=[state, chat, prompt, output, error]
+    outputs=[state, chat, prompt, output]
 
     prompt.submit(generate, inputs=inputs, outputs=outputs)
     button_generate.click(generate, inputs=inputs, outputs=outputs)
-    button_reset.click(reset, inputs=[state], outputs=[state, chat, output, error])
+    button_reset.click(reset, inputs=[state], outputs=[state, chat, output])
 
     examples = gr.Examples(inputs=[context, prompt, model, do_sample, top_k, top_p, temperature],
         examples=[