Update app.py

app.py

@@ -31,29 +31,27 @@ def get_input_ids(inst, history):
     return input_ids
 
 
-@spaces.GPU
-def chat(inst, history, temperature, top_p, repetition_penalty):
+def generate(generation_kwargs):
     with torch.no_grad():
-        streamer = TextIteratorStreamer(tokenizer, skip_prompt=True, skip_special_tokens=True)
-        input_ids = get_input_ids(inst, history)
-        if len(input_ids) > MAX_INPUT_LIMIT:
-            yield "The input is too long, please clear the history."
-            return
-        generate_config = dict(
-            max_new_tokens=MAX_NEW_TOKENS,
-            temperature=temperature,
-            top_p=top_p,
-            repetition_penalty=repetition_penalty
-        )
-        print(generate_config)
-        generation_kwargs = dict(input_ids=torch.tensor([input_ids]).to(model.device), do_sample=True,
-                                 streamer=streamer, **generate_config)
         Thread(target=model.generate, kwargs=generation_kwargs).start()
 
-        outputs = ""
-        for new_text in streamer:
-            outputs += new_text
-            yield outputs
+
+@spaces.GPU
+def chat(inst, history, temperature, top_p, repetition_penalty):
+    streamer = TextIteratorStreamer(tokenizer, skip_prompt=True, skip_special_tokens=True)
+    input_ids = get_input_ids(inst, history)
+    if len(input_ids) > MAX_INPUT_LIMIT:
+        yield "The input is too long, please clear the history."
+        return
+    generation_kwargs = dict(input_ids=torch.tensor([input_ids]).to(model.device),
+                             streamer=streamer, do_sample=True, max_new_tokens=MAX_NEW_TOKENS,
+                             temperature=temperature, top_p=top_p, repetition_penalty=repetition_penalty)
+    generate(generation_kwargs)
+
+    outputs = ""
+    for new_text in streamer:
+        outputs += new_text
+        yield outputs
 
 
 additional_inputs = [
@@ -93,7 +91,8 @@ gr.ChatInterface(chat,
                  description='Hello, I am Blossom, an open source conversational large language model.🌠'
                              '<a href="https://github.com/Azure99/BlossomLM">GitHub</a>',
                  theme="soft",
-                 examples=[["Hello"], ["What is MBTI"], ["用Python实现二分查找"], ["为switch写一篇小红书种草文案，带上emoji"]],
+                 examples=[["Hello"], ["What is MBTI"], ["用Python实现二分查找"],
+                           ["为switch写一篇小红书种草文案，带上emoji"]],
                  additional_inputs=additional_inputs,
                  additional_inputs_accordion=gr.Accordion(label="Config", open=True),
                  clear_btn="🗑️Clear",