[update]edit main

main.py

@@ -2,11 +2,12 @@
 # -*- coding: utf-8 -*-
 import argparse
 import os
+from threading import Thread
 
 import gradio as gr
 from transformers import AutoModel, AutoTokenizer
 from transformers.models.auto import AutoModelForCausalLM, AutoTokenizer
-# from transformers.utils.quantization_config import BitsAndBytesConfig
+from transformers.generation.streamers import TextIteratorStreamer
 import torch
 
 from project_settings import project_path
@@ -72,7 +73,7 @@ def main():
     )
     model = model.bfloat16().eval()
 
-    def fn(text: str):
+    def fn_non_stream(text: str):
         input_ids = tokenizer(
             text,
             return_tensors="pt",
@@ -84,8 +85,12 @@ def main():
 
         with torch.no_grad():
             outputs = model.generate(
-                input_ids=input_ids, max_new_tokens=args.max_new_tokens, do_sample=True,
-                top_p=args.top_p, temperature=args.temperature, repetition_penalty=args.repetition_penalty,
+                input_ids=input_ids,
+                max_new_tokens=args.max_new_tokens,
+                do_sample=True,
+                top_p=args.top_p,
+                temperature=args.temperature,
+                repetition_penalty=args.repetition_penalty,
                 eos_token_id=tokenizer.eos_token_id
             )
             outputs = outputs.tolist()[0][len(input_ids[0]):]
@@ -94,6 +99,44 @@ def main():
 
         return [(text, response)]
 
+    def fn_stream(text: str):
+        input_ids = tokenizer(
+            text,
+            return_tensors="pt",
+            add_special_tokens=False,
+        ).input_ids.to(args.device)
+        bos_token_id = torch.tensor([[tokenizer.bos_token_id]], dtype=torch.long).to(args.device)
+        eos_token_id = torch.tensor([[tokenizer.eos_token_id]], dtype=torch.long).to(args.device)
+        input_ids = torch.concat([bos_token_id, input_ids, eos_token_id], dim=1)
+
+        streamer = TextIteratorStreamer(tokenizer=tokenizer)
+
+        generation_kwargs = dict(
+            inputs=input_ids,
+            max_new_tokens=args.max_new_tokens,
+            do_sample=True,
+            top_p=args.top_p,
+            temperature=args.temperature,
+            repetition_penalty=args.repetition_penalty,
+            eos_token_id=tokenizer.eos_token_id,
+            pad_token_id=tokenizer.pad_token_id,
+            streamer=streamer,
+        )
+        thread = Thread(target=model.generate, kwargs=generation_kwargs)
+        thread.start()
+
+        output = ""
+        for output_ in streamer:
+            output_ = output_.tolist()[0][len(input_ids[0]):]
+            output_ = tokenizer.decode(output_)
+            output_ = output_.strip().replace(tokenizer.eos_token, "").strip()
+
+            output += output_
+
+            result = [(text, output)]
+            chatbot.value = result
+            yield result
+
     with gr.Blocks() as blocks:
         gr.Markdown(value=description)
 
@@ -108,8 +151,8 @@ def main():
 
         gr.Examples(examples, text_box)
 
-        text_box.submit(fn, [text_box], [chatbot])
-        submit_button.click(fn, [text_box], [chatbot])
+        text_box.submit(fn_stream, [text_box], [chatbot])
+        submit_button.click(fn_stream, [text_box], [chatbot])
         clear_button.click(
             fn=lambda: ("", ""),
             outputs=[text_box, chatbot],