Update app.py

app.py

@@ -1,7 +1,8 @@
 import gradio as gr
-from transformers import pipeline, AutoTokenizer
+from transformers import pipeline, AutoTokenizer, TextIteratorStreamer
 import torch
 import spaces
+from threading import Thread
 import os
 
 @spaces.GPU
@@ -23,8 +24,15 @@ def generate(
         prompt = user_input
     else:
         prompt = f"<|im_start|>system\nYou are an AI assistant that follows instruction extremely well. Help as much as you can.<|im_end|>\n<|im_start|>user\n{user_input}<|im_end|>\n<|im_start|>assistant\n"
-    outputs = pipe(prompt, max_new_tokens=max_new_tokens, do_sample=True,
-                   temperature=temperature, top_k=top_k, top_p=top_p, repetition_penalty=1.1)
+    streamer = TextIteratorStreamer(tokenizer, timeout=240.0, skip_prompt=True, skip_special_tokens=True)
+    generation_kwargs = dict(text_inputs=prompt, streamer=streamer, max_new_tokens=max_new_tokens, do_sample=True, top_p=top_p, top_k=top_k, 
+                              temperature=temperature, num_beams=1, repetition_penalty=repetition_penalty, eos_token_id=pipe.tokenizer.eos_token_id)
+    t = Thread(target=pipe.__call__, kwargs=generation_kwargs)
+    t.start()
+    outputs = []
+    for chunk in streamer:
+        outputs.append(chunk)
+        yield "".join(outputs)
     return outputs[0]["generated_text"]
 
 model_choices = ["M4-ai/Hercules-Mini-1.8B", "Locutusque/Hyperion-3.0-Mistral-7B-DPO", "Locutusque/OpenCerebrum-1.5-Mistral-11B-Evolved-beta", "M4-ai/tau-1.8B", "Locutusque/OpenCerebrum-1.5-Mistral-7b-v0.2-alpha", "Locutusque/Hercules-4.0-Mistral-v0.2-7B", "Locutusque/Hercules-3.1-Mistral-7B"]