Update app.py

app.py

@@ -72,15 +72,19 @@ def generate_tokens(model, generator):
     global stop_generation
     app.logger.info('generate_tokens started')
     #with lock:
-    for token in generator:            
-        if token == model.token_eos() or stop_generation:
-            stop_generation = False
-            app.logger.info('Abort generating')
-            yield b''  # End of chunk
-            break
-            
-        token_str = model.detokenize([token])#.decode("utf-8", errors="ignore")
-        yield token_str 
+    try:
+        for token in generator:            
+            if token == model.token_eos() or stop_generation:
+                stop_generation = False
+                app.logger.info('Abort generating')
+                yield b''  # End of chunk
+                break
+                
+            token_str = model.detokenize([token])#.decode("utf-8", errors="ignore")
+            yield token_str 
+    except Exception as e:
+        app.logger.info('generator exception')
+        yield b''  # End of chunk
 
 @app.route('/stop_generation', methods=['GET'])
 def handler_stop_generation():
@@ -233,4 +237,4 @@ def generate_response():
     return Response(generate_tokens(model, generator), content_type='text/plain', status=200, direct_passthrough=True)
 
 if __name__ == "__main__":
-    app.run(host="0.0.0.0", port=7860, debug=False, threaded=False)
+    app.run(host="0.0.0.0", port=7860, debug=False, threaded=True)