[update]add main

main.py

@@ -5,6 +5,7 @@ from threading import Thread
 
 import gradio as gr
 from transformers import AutoModelForCausalLM, AutoTokenizer
+from transformers.generation.streamers import TextIteratorStreamer
 import torch
 
 from project_settings import project_path
@@ -109,6 +110,57 @@ def chat_with_llm_non_stream(question: str,
     return result
 
 
+def chat_with_llm_streaming(question: str,
+                            history: List[Tuple[str, str]],
+                            pretrained_model_name_or_path: str,
+                            max_new_tokens: int, top_p: float, temperature: float, repetition_penalty: float,
+                            ):
+    device: str = "cuda" if torch.cuda.is_available() else "cpu"
+
+    model, tokenizer = init_model(pretrained_model_name_or_path)
+
+    text_list = list()
+    for pair in history:
+        text_list.extend(pair)
+    text_list.append(question)
+
+    text_encoded = tokenizer.__call__(text_list, add_special_tokens=False)
+    batch_input_ids = text_encoded["input_ids"]
+
+    input_ids = [tokenizer.bos_token_id]
+    for input_ids_ in batch_input_ids:
+        input_ids.extend(input_ids_)
+        input_ids.append(tokenizer.eos_token_id)
+    input_ids = torch.tensor([input_ids], dtype=torch.long).to(device)
+
+    streamer = TextIteratorStreamer(tokenizer=tokenizer)
+
+    generation_kwargs = dict(
+        inputs=input_ids,
+        max_new_tokens=max_new_tokens,
+        do_sample=True,
+        top_p=top_p,
+        temperature=temperature,
+        repetition_penalty=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()
+
+    answer = ""
+    for output_ in streamer:
+        output_ = output_.replace(question, "")
+        output_ = output_.replace(tokenizer.eos_token, "")
+
+        answer += output_
+
+        result = [(question, answer)]
+
+        yield history + result
+
+
 def main():
     description = """
     chat llm
@@ -153,8 +205,8 @@ def main():
         outputs = [
             chatbot
         ]
-        text_box.submit(chat_with_llm_non_stream, inputs, outputs)
-        submit_button.click(chat_with_llm_non_stream, inputs, outputs)
+        text_box.submit(chat_with_llm_streaming, inputs, outputs)
+        submit_button.click(chat_with_llm_streaming, inputs, outputs)
         clear_button.click(
             fn=lambda: ('', ''),
             outputs=[text_box, chatbot],