[update]edit main

main.py

@@ -36,10 +36,10 @@ def get_args():
 description = """
 ## Qwen-7B
 
-基于 [Qwen-7B](https://huggingface.co/qgyd2021/Qwen-7B) 模型, 在 [chinese_modern_poetry](https://huggingface.co/datasets/Iess/chinese_modern_poetry) 的 prompt 数据集上训练. 
+基于 [Qwen-7B](https://huggingface.co/qgyd2021/Qwen-7B) 模型, 在 [chinese_modern_poetry](https://huggingface.co/datasets/Iess/chinese_modern_poetry) 的 prompt 数据集上训练了 2 个 epoch. 
 
 可用于生成现代诗. 如下: 
-使用下列意象写一首现代诗：智慧，刀刃
+使用下列意象写一首现代诗：智慧，刀刃. 
 """
 
 
@@ -64,7 +64,7 @@ def main():
         args.pretrained_model_name_or_path,
         trust_remote_code=True,
         low_cpu_mem_usage=True,
-        torch_dtype=torch.float16,
+        torch_dtype=torch.bfloat16,
         device_map="auto",
         offload_folder="./offload",
         offload_state_dict=True,
@@ -88,30 +88,31 @@ def main():
             outputs = outputs.tolist()[0][len(input_ids[0]):]
             response = tokenizer.decode(outputs)
             response = response.strip().replace(tokenizer.eos_token, "").strip()
-        return response
 
-    def fn_stream(text: str,
-                  max_new_tokens: int = 200,
-                  top_p: float = 0.85,
-                  temperature: float = 0.35,
-                  repetition_penalty: float = 1.2
-                  ):
-        return
+        return [(text, response)]
 
     with gr.Blocks() as blocks:
         gr.Markdown(value=description)
 
-        chatbot = gr.Chatbot([], elem_id="chatbot").style(height=400)
+        chatbot = gr.Chatbot([], elem_id="chatbot", height=400)
         with gr.Row():
             with gr.Column(scale=4):
-                input_text_box = gr.Textbox(show_label=False, placeholder="Enter text and press enter").style(container=False)
+                text_box = gr.Textbox(show_label=False, placeholder="Enter text and press enter", container=False)
             with gr.Column(scale=1):
-                button = gr.Button("Generate")
-
-        gr.Examples(examples, input_text_box)
-
-        input_text_box.submit(fn, [input_text_box], [chatbot])
-        button.click(fn, [input_text_box], [chatbot])
+                submit_button = gr.Button("💬Submit")
+            with gr.Column(scale=1):
+                clear_button = gr.Button("🗑️Clear", variant="secondary")
+
+        gr.Examples(examples, text_box)
+
+        text_box.submit(fn, [text_box], [chatbot])
+        submit_button.click(fn, [text_box], [chatbot])
+        clear_button.click(
+            fn=lambda: ("", ""),
+            outputs=[text_box, chatbot],
+            queue=False,
+            api_name=False,
+        )
 
     blocks.queue().launch()