Update app.py

app.py

@@ -1,30 +1,50 @@
+import torch
+from transformers import AutoTokenizer, AutoModelForCausalLM, GenerationConfig
 import gradio as gr
-from transformers import MarianMTModel, MarianTokenizer
 
-# 加载 MarianMT 模型和分词器
-model_name = "Helsinki-NLP/opus-mt-en-zh"
-model = MarianMTModel.from_pretrained(model_name).to("cpu")  # 强制使用 CPU
-tokenizer = MarianTokenizer.from_pretrained(model_name)
+model_name = "deepseek-ai/deepseek-llm-7b-chat"
 
-# 翻译功能
-def translate_text(input_text):
-    # 使用 MarianMT 模型进行翻译
-    inputs = tokenizer(input_text, return_tensors="pt", padding=True, truncation=True)
-    outputs = model.generate(**inputs)
-    translated_text = tokenizer.decode(outputs[0], skip_special_tokens=True)
-    
-    return translated_text
+# 加载 tokenizer 和模型
+tokenizer = AutoTokenizer.from_pretrained(model_name)
+model = AutoModelForCausalLM.from_pretrained(
+    model_name,
+    torch_dtype=torch.bfloat16 if torch.cuda.is_available() else torch.float32,
+    device_map="auto" if torch.cuda.is_available() else None
+)
+model.generation_config = GenerationConfig.from_pretrained(model_name)
+model.generation_config.pad_token_id = model.generation_config.eos_token_id
 
-# 创建 Gradio 界面
-with gr.Blocks() as demo:
-    gr.Markdown("# 英文文本翻译器")
+# 推理函数
+def chat(prompt, max_new_tokens=256, temperature=0.7, top_p=0.9):
+    messages = [{"role": "user", "content": prompt}]
+    input_tensor = tokenizer.apply_chat_template(
+        messages, add_generation_prompt=True, return_tensors="pt"
+    ).to(model.device)
+    with torch.no_grad():
+        outputs = model.generate(
+            input_tensor,
+            max_new_tokens=max_new_tokens,
+            temperature=temperature,
+            top_p=top_p,
+            do_sample=True,
+            eos_token_id=tokenizer.eos_token_id
+        )
+    result = tokenizer.decode(outputs[0][input_tensor.shape[1]:], skip_special_tokens=True)
+    return result
 
-    with gr.Row():
-        text_input = gr.Textbox(label="输入英文文本", lines=5)
-    
-    translate_button = gr.Button("开始翻译")
-    output_text = gr.Textbox(label="翻译后的中文文本", lines=5)
+# Gradio 界面
+iface = gr.Interface(
+    fn=chat,
+    inputs=[
+        gr.Textbox(lines=4, label="请输入您的问题"),
+        gr.Slider(32, 1024, step=16, value=256, label="最大生成长度"),
+        gr.Slider(0.1, 1.5, step=0.1, value=0.7, label="Temperature"),
+        gr.Slider(0.1, 1.0, step=0.05, value=0.9, label="Top-p")
+    ],
+    outputs=gr.Textbox(label="模型回复"),
+    title="DeepSeek LLM 7B Chat 演示",
+    description="基于 Hugging Face Spaces 的部署示例"
+)
 
-    translate_button.click(fn=translate_text, inputs=text_input, outputs=output_text)
-
-demo.launch()
+if __name__ == "__main__":
+    iface.launch()