Update app.py

app.py

@@ -1,31 +1,34 @@
-import gradio as gr
 from transformers import AutoModelForCausalLM, AutoTokenizer
+import torch
 
 # 加载模型和分词器
-model_name = "nvidia/NVLM-D-72B"
+model_name = "umingpeng/Meta-Llama-3.1-8B-Instruct"  # 选择合适的模型
 tokenizer = AutoTokenizer.from_pretrained(model_name)
 model = AutoModelForCausalLM.from_pretrained(model_name)
 
-def generate_response(prompt):
-    # 编码输入
-    inputs = tokenizer(prompt, return_tensors="pt")
+# 设置设备
+device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
+model.to(device)
+
+def generate_response(user_input):
+    # 构建输入格式
+    messages = [
+        {"role": "user", "content": user_input}
+    ]
+    
+    # 使用聊天模板处理输入
+    tokenized_input = tokenizer.apply_chat_template(messages, tokenize=True, add_generation_prompt=True, return_tensors="pt").to(device)
     
-    # 生成输出
-    outputs = model.generate(**inputs)
+    # 生成响应
+    with torch.no_grad():
+        output = model.generate(**tokenized_input, max_new_tokens=256)
     
-    # 解码输出
-    response = tokenizer.decode(outputs[0], skip_special_tokens=True)
+    # 解码生成的文本
+    response = tokenizer.decode(output[0], skip_special_tokens=True)
     return response
 
-# 创建 Gradio 接口
-iface = gr.Interface(
-    fn=generate_response,
-    inputs=gr.inputs.Textbox(lines=2, placeholder="输入你的问题..."),
-    outputs="text",
-    title="NVLM-D-72B 交互式问答",
-    description="使用 NVIDIA 的 NVLM-D-72B 模型进行问答。"
-)
-
-# 启动应用
+# 主程序
 if __name__ == "__main__":
-    iface.launch()
+    user_input = input("你想问什么？")
+    response = generate_response(user_input)
+    print("助手:", response)