Update app.py

app.py

@@ -1,4 +1,5 @@
 import os
+
 os.environ['CUDA_LAUNCH_BLOCKING'] = '1'
 
 import torch
@@ -7,28 +8,48 @@ from transformers import AutoModelForCausalLM, AutoTokenizer
 
 
 def init_model():
-    model = AutoModelForCausalLM.from_pretrained("Linly-AI/Chinese-LLaMA-2-7B-hf", device_map="cuda:0", torch_dtype=torch.float16, trust_remote_code=True)
+    model = AutoModelForCausalLM.from_pretrained("Linly-AI/Chinese-LLaMA-2-7B-hf", device_map="cuda:0",
+                                                 torch_dtype=torch.bfloat16, trust_remote_code=True)
     tokenizer = AutoTokenizer.from_pretrained("Linly-AI/Chinese-LLaMA-2-7B-hf", use_fast=False, trust_remote_code=True)
     return model, tokenizer
-    
 
 
-def chat(prompt, top_k, temperature):
-    prompt = f"### Instruction:{prompt.strip()}  ### Response:"
-    inputs = tokenizer(prompt, return_tensors="pt").to("cuda:0")
-    generate_ids = model.generate(inputs.input_ids, do_sample=True, max_new_tokens=2048, top_k=int(top_k), top_p=0.84, temperature=float(temperature), repetition_penalty=1.15, eos_token_id=2, bos_token_id=1, pad_token_id=0)
-    response = tokenizer.batch_decode(generate_ids, skip_special_tokens=True, clean_up_tokenization_spaces=False)[0]
-    response = response.lstrip(prompt)
-    print('-log: ',prompt, response)
-    return response
+def process(message, history):
+    input_prompt = ""
+    for interaction in history:
+        input_prompt = f"{input_prompt} User: {str(interaction[0]).strip(' ')} Bot: {str(interaction[1]).strip(' ')}"
+    input_prompt = f"{input_prompt} ### Instruction:{message.strip()}  ### Response:"
+    inputs = tokenizer(input_prompt, return_tensors="pt").to("cuda:0")
+    try:
+        generate_ids = model.generate(inputs.input_ids, max_new_tokens=2048, do_sample=True, top_k=30, top_p=0.84,
+                                      temperature=1.0, repetition_penalty=1.15, eos_token_id=2, bos_token_id=1,
+                                      pad_token_id=0)
+        response = tokenizer.batch_decode(generate_ids, skip_special_tokens=True, clean_up_tokenization_spaces=False)[0]
+        print('log:', response)
+        response = response.split("### Response:")[-1]
+        return response
+    except:
+        return "Error: 会话超长，请重试！"
 
 
 if __name__ == '__main__':
+    examples = ["Python和JavaScript编程语言的主要区别是什么？", "影响消费者行为的主要因素是什么？", "请用pytorch实现一个带ReLU激活函数的全连接层的代码",
+                "请用C++编程语言实现“给你两个字符串haystack和needle，在haystack字符串中找出needle字符串的第一个匹配项的下标（下标从 0 开始）。如果needle不是haystack的一部分，则返回-1。",
+                "为什么有些人选择使用纸质地图或寻求方向，而不是依赖GPS设备或智能手机应用程序？",
+                "应对压力最有效的方法是什么？"]
     model, tokenizer = init_model()
-    demo = gr.Interface(
-        fn=chat,
-        inputs=["text", gr.Slider(1, 60, value=10, step=1), gr.Slider(0.1, 2.0, value=1.0, step=0.1)],
-        outputs="text",
+    demo = gr.ChatInterface(
+        process,
+        chatbot=gr.Chatbot(height=600),
+        textbox=gr.Textbox(placeholder="Input", container=False, scale=7),
+        title="Linly ChatFlow",
+        description="",
+        theme="soft",
+        examples=examples,
+        cache_examples=True,
+        retry_btn="Retry",
+        undo_btn="Delete Previous",
+        clear_btn="Clear",
     )
-    demo.launch()
-
+    demo.queue(concurrency_count=75).launch(share=True, server_name="0.0.0.0", server_port=7862, show_error=True,
+                                            debug=True)