feasible

app.py

@@ -1,43 +1,50 @@
 from peft import PeftModel, PeftConfig
-from transformers import AutoModelForCausalLM, AutoTokenizer
+from transformers import AutoModelForCausalLM, AutoTokenizer, TextIteratorStreamer
 from transformers_stream_generator.main import NewGenerationMixin, StreamGenerationConfig
+from threading import Thread
 import gradio as gr
 import torch
 
 config = PeftConfig.from_pretrained("Junity/Genshin-World-Model", trust_remote_code=True)
-model = AutoModelForCausalLM.from_pretrained("../Baichuan/models--baichuan-inc--Baichuan-13B-Base\snapshots\Baichuan-13B-Base", trust_remote_code=True)
-model = PeftModel.from_pretrained(model, r"../Baichuan/r64alpha32dropout0.5loss0.007/checkpoint-5000", trust_remote_code=True)
-tokenizer = AutoTokenizer.from_pretrained("Junity/Genshin-World-Model", trust_remote_code=True)
-
+model = AutoModelForCausalLM.from_pretrained("baichuan-inc/Baichuan-13B-Base", torch_dtype=torch.float32, device_map="auto", trust_remote_code=True)
+model = PeftModel.from_pretrained(model, r"Junity/Genshin-World-Model", torch_dtype=torch.float32, device_map="auto", trust_remote_code=True)
+tokenizer = AutoTokenizer.from_pretrained("baichuan-inc/Baichuan-13B-Base", trust_remote_code=True)
 history = []
-device = "cpu"
-
+device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
+if device == "cuda":
+    model.cuda()
+    model = model.half()
 
-def respond(role_name, msg, chatbot, character):
-    global history
-    if role_name is not None:
-        history.append(role_name + "：" + msg)
+def respond(role_name, msg, textbox):
+    if textbox != '':
+        textbox = textbox + "\n" + role_name + "：" + msg + ('。' if msg[-1] not in ['。', '！', '？'] else '') + '\n'
+        yield ["", textbox]
     else:
-        history.append(msg)
-    total_input = []
-    for i, message in enumerate(history[::-1]):
-        content_tokens = tokenizer.encode(message + '\n')
-        total_input = content_tokens + total_input
-        if content_tokens + total_input > 4096:
-            break
-    total_input = total_input[-4096:]
-    input_ids = torch.LongTensor([total_input]).to(device)
+        textbox = textbox + role_name + "：" + msg + ('。' if msg[-1] not in ['。', '！', '？'] else '') + '\n'
+        yield ["", textbox]
+    input_ids = tokenizer.encode(textbox)[-4096:]
+    input_ids = torch.LongTensor([input_ids]).to(device)
     generation_config = model.generation_config
     stream_config = StreamGenerationConfig(**generation_config.to_dict(), do_stream=True)
+    gen_kwargs = {}
+    gen_kwargs.update(dict(
+        input_ids=input_ids,
+        temperature=1.0,
+        top_p=0.75,
+        repetition_penalty=1.2,
+        max_new_tokens=256
+    ))
+    outputs = []
+    print(input_ids)
+    streamer = TextIteratorStreamer(tokenizer, timeout=60.0, skip_prompt=True, skip_special_tokens=True)
+    gen_kwargs["streamer"] = streamer
 
-    def stream_generator():
-        outputs = []
-        for token in model.generate(input_ids, generation_config=stream_config):
-            outputs.append(token.item())
-            yield None, tokenizer.decode(outputs, skip_special_tokens=True)
-
-    return stream_generator()
+    thread = Thread(target=model.generate, kwargs=gen_kwargs)
+    thread.start()
 
+    for new_text in streamer:
+        textbox += new_text
+        yield ["", textbox]
 
 with gr.Blocks() as demo:
     gr.Markdown(
@@ -53,7 +60,7 @@ with gr.Blocks() as demo:
     with gr.Row():
         clear = gr.Button("Clear")
         sub = gr.Button("Submit")
-    chatbot = gr.Chatbot()
-    sub.click(fn=respond, inputs=[role_name, msg, chatbot], outputs=[msg, chatbot])
-    clear.click(lambda: None, None, chatbot, queue=False)
-    demo.queue().launch()
+    textbox = gr.Textbox(interactive=False)
+    sub.click(fn=respond, inputs=[role_name, msg, textbox], outputs=[msg, textbox])
+    clear.click(lambda: None, None, textbox, queue=False)
+    demo.queue().launch(server_port=6006)