Update app.py

app.py

@@ -1,22 +1,48 @@
-from transformers import AutoModel, AutoTokenizer
-import gradio as gr
-
-# 加载模型和分词器
-model_id = "BioMistral/BioMistral-7B"
-tokenizer = AutoTokenizer.from_pretrained(model_id)
-model = AutoModel.from_pretrained(model_id)
-
-def generate_text(prompt):
-    # 对输入进行编码
-    inputs = tokenizer(prompt, return_tensors="pt")
-    # 生成文本（这里使用模型的输出作为示例）
-    outputs = model(**inputs)
-    
-    # 处理输出（这里返回最后一个隐藏状态的形状作为示例）
-    return outputs.last_hidden_state.tolist()  # 转换为列表以便Gradio显示
-
-# 创建Gradio界面
-iface = gr.Interface(fn=generate_text, inputs="text", outputs="text")
-
-# 启动Gradio界面并共享链接
-iface.launch(share=True)
+from transformers import AutoModelForCausalLM, AutoTokenizer
+import torch
+
+# 加载指令模型
+model = AutoModelForCausalLM.from_pretrained(
+    "MediaTek-Research/Breeze-7B-Instruct-v1_0",
+    device_map="auto",
+    torch_dtype=torch.bfloat16,
+)
+
+# 加载基础模型（如果需要）
+# model_base = AutoModelForCausalLM.from_pretrained(
+#     "MediaTek-Research/Breeze-7B-Base-v1_0",
+#     device_map="auto",
+#     torch_dtype=torch.bfloat16,
+# )
+
+# 加载分词器
+tokenizer = AutoTokenizer.from_pretrained("MediaTek-Research/Breeze-7B-Instruct-v1_0")
+
+# 定义SYS_PROMPT
+SYS_PROMPT = "You are a helpful AI assistant built by MediaTek Research. The user you are helping speaks Traditional Chinese and comes from Taiwan."
+
+# 定义聊天内容
+chat = [
+    {"role": "user", "content": "你好，請問你可以完成什麼任務？"},
+    {"role": "assistant", "content": "你好，我可以幫助您解決各種問題、提供資訊和協助您完成許多不同的任務。例如：回答技術問題、提供建議、翻譯文字、尋找資料或協助您安排行程等。請告訴我如何能幫助您。"},
+    {"role": "user", "content": "太棒了！"},
+]
+
+# 应用聊天模板
+prompt = tokenizer.apply_chat_template(chat, tokenize=False)
+full_prompt = f"<s>{SYS_PROMPT} [INST] {prompt} [/INST]"
+
+# 生成文本
+inputs = tokenizer(full_prompt, return_tensors="pt").to(model.device)
+outputs = model.generate(
+    inputs["input_ids"],
+    max_new_tokens=128,
+    top_p=0.95,
+    top_k=50,
+    repetition_penalty=1.1,
+    temperature=0.7,
+)
+
+# 解码输出
+generated_text = tokenizer.decode(outputs[0], skip_special_tokens=True)
+print(generated_text)