Print out traceback to Output when something gets wrong.

app.py

@@ -1,6 +1,7 @@
 #!/usr/bin/env python
 # or gradio app.py
 
+import traceback
 import gradio as gr
 import chat_client
 
@@ -11,8 +12,13 @@ def generate(prompt, model, endseq, max_length,
         do_sample, top_k, top_p, temperature,
         add_stoptoken, copy_output):
 
-    client = chat_client.ModelClient(CHAT_URL)
-    client.open_session(f"bigscience/{model}-petals", max_length)
+    try:
+        client = chat_client.ModelClient(CHAT_URL)
+        client.open_session(f"bigscience/{model}-petals", max_length)
+    except Exception:
+        print(traceback.format_exc())
+        yield [prompt, "Error: " + traceback.format_exc()]
+        return
 
     if add_stoptoken:
         prompt += "</s>" if "bloomz" in model else "\n\n"
@@ -44,20 +50,25 @@ def generate(prompt, model, endseq, max_length,
     # don't wait to generator to return first result
     yield [prompt2, output]
 
-    for out in client.generate(prompt,
+    try:
+        for out in client.generate(prompt,
                     max_new_tokens=1,
                     do_sample=do_sample,
                     temperature=temperature,
                     top_k=top_k,
                     top_p=top_p,
                     extra_stop_sequences=seq
-        ):
+            ):
 
-        output += out
-        if copy_output:
-            prompt2 += out
+            output += out
+            if copy_output:
+                prompt2 += out
 
-        yield [prompt2, output]
+            yield [prompt2, output]
+    except Exception:
+        print(traceback.format_exc())
+        yield [prompt, "Error: " + traceback.format_exc()]
+        return
 
 with gr.Blocks() as iface:
     gr.Markdown("""# Petals playground