Improved switching between models.

chat.py

@@ -5,6 +5,7 @@ import traceback
 import gradio as gr
 import chat_client
 import json
+import re
 
 CHAT_URL='ws://chat.petals.ml/api/v2/generate'
 #CHAT_URL='ws://localhost:8000/api/v2/generate'
@@ -38,6 +39,12 @@ def _generate(state, prompt, model, context, output, endseq, max_length,
     print('prompt', prompt)
     eos = "</s>\n" if "bloomz" in model else "\n\n"
 
+    if state['model'] != model and output:
+        # If the connection is resumed, output is truncated in generate().
+        # So this happen when user change model.
+        context = output
+        output = ''
+
     if state['model'] != model or \
         state['client'] == None or state['client'].is_session() == False:
 
@@ -48,6 +55,7 @@ def _generate(state, prompt, model, context, output, endseq, max_length,
         except Exception:
             print(traceback.format_exc())
             raise gr.Error(traceback.format_exc())
+
     else:
         context = ''
 
@@ -63,6 +71,7 @@ def _generate(state, prompt, model, context, output, endseq, max_length,
         prompt2 = prompt.replace("\n\n", eos) + "</s>\n"
     else:
         context = context.replace("</s>", eos)
+        context = re.sub(r"\n\n+", "\n\n", context)
         prompt2 = prompt.replace("</s>", eos) + "\n\n"
     
     prompt2 = f"{context}Human: {prompt2}AI:"