Update app.py

app.py

@@ -118,14 +118,14 @@ with gr.Blocks() as demo:
         prefix = history[-1]["content"]
         history.append({"role": "assistant", "content": ""})
         history[-1]["content"] += "Generating with the given prefix...\n"
-        queue = Queue()
+        queue = Queue(maxsize=10)
         class MyStreamer:
             def put(self, tokens):
                 for token in tokens.flatten():
                     text = pipe.tokenizer.decode(token.item())
                     if text == '<|begin_of_text|>':
                         continue
-                    queue.put(text)
+                    queue.put(text, block=True, timeout=5)
             def end(self):
                 queue.put(None)
         def background_fn():
@@ -140,7 +140,11 @@ with gr.Blocks() as demo:
             print()
         Thread(target=background_fn).start()
         while True:
-            text = queue.get()
+            try:
+                text = queue.get()
+            except queue.Full:
+                print("Queue is full. Exiting.")
+                break
             if text is None:
                 break
             history[-1]["content"] += text