Upload app.py

app.py

@@ -0,0 +1,83 @@
+from threading import Thread
+
+import torch
+import gradio as gr
+from transformers import AutoTokenizer, AutoModelForCausalLM, TextIteratorStreamer
+
+model_id = "fireballoon/baichuan-vicuna-chinese-7b"
+torch_device = "cuda" if torch.cuda.is_available() else "cpu"
+print("Running on device:", torch_device)
+print("CPU threads:", torch.get_num_threads())
+
+
+if torch_device == "cuda":
+    model = AutoModelForCausalLM.from_pretrained(model_id, torch_dtype=torch.float16).cuda()
+else:
+    model = AutoModelForCausalLM.from_pretrained(model_id)
+tokenizer = AutoTokenizer.from_pretrained(model_id, use_fast=False)
+
+
+def run_generation(history, *args, **kwargs):
+    # Get the model and tokenizer, and tokenize the user text.
+
+    instruction = "A chat between a curious user and an artificial intelligence assistant. " \
+                  "The assistant gives helpful, detailed, and polite answers to the user's questions."
+    context = ''.join([f" USER: {turn[0].strip()} ASSISTANT: {turn[1].strip()} </s>" for turn in history[:-1]])
+    prompt = instruction + context + f" USER: {history[-1][0].strip()} ASSISTANT:"
+
+    input_ids = tokenizer(prompt, return_tensors="pt").input_ids.cuda()
+
+    print()
+    print(prompt)
+    print('##', input_ids.size())
+
+    # Start generation on a separate thread, so that we don't block the UI. The text is pulled from the streamer
+    # in the main thread. Adds timeout to the streamer to handle exceptions in the generation thread.
+    streamer = TextIteratorStreamer(tokenizer, timeout=10., skip_prompt=True, skip_special_tokens=True)
+    generate_kwargs = dict(
+        input_ids=input_ids,
+        streamer=streamer,
+        max_new_tokens=2048,
+        do_sample=True,
+        temperature=0.7,
+        repetition_penalty=1.1,
+        top_p=0.85
+    )
+    t = Thread(target=model.generate, kwargs=generate_kwargs)
+    t.start()
+
+    # Pull the generated text from the streamer, and update the model output.
+    history[-1][1] = ""
+
+    print("")
+    for new_text in streamer:
+        history[-1][1] += new_text
+        print(new_text, end="", flush=True)
+        yield history
+    print('</s>')
+    return history
+
+
+def reset_textbox():
+    return gr.update(value='')
+
+
+with gr.Blocks() as demo:
+    gr.Markdown(
+        "# Baichuan Vicuna Chinese\n"
+        f"[{model_id}](https://huggingface.co/{model_id})：使用中英双语sharegpt数据全参数微调的对话模型，基于baichuan-7b"
+    )
+    chatbot = gr.Chatbot().style(height=600)
+    msg = gr.Textbox()
+    clear = gr.ClearButton([msg, chatbot])
+
+    def user(user_message, history):
+        return gr.update(value="", interactive=False), history + [[user_message, None]]
+
+    response = msg.submit(user, [msg, chatbot], [msg, chatbot], queue=False).then(
+        run_generation, chatbot, chatbot
+    )
+    response.then(lambda: gr.update(interactive=True), None, [msg], queue=False)
+
+    demo.queue()
+    demo.launch(server_name='0.0.0.0')