Update app.py

app.py

@@ -1,63 +1,51 @@
 import gradio as gr
-from huggingface_hub import InferenceClient
+from transformers import AutoTokenizer, AutoModelForCausalLM
+import torch
 
-# 修改为你的模型
-client = InferenceClient("cmz1024/minimind-zero")
+# 加载模型和tokenizer
+model_name = "cmz1024/minimind-zero"  # 替换为你的模型路径
+tokenizer = AutoTokenizer.from_pretrained(model_name)
+model = AutoModelForCausalLM.from_pretrained(model_name, torch_dtype=torch.float16)
+model.eval()
 
-def respond(
-    message,
-    history: list[tuple[str, str]],
-    system_message,
-    max_tokens,
-    temperature,
-    top_p,
-):
-    messages = [{"role": "system", "content": system_message}]
+if torch.cuda.is_available():
+    model = model.cuda()
 
-    for val in history:
-        if val[0]:
-            messages.append({"role": "user", "content": val[0]})
-        if val[1]:
-            messages.append({"role": "assistant", "content": val[1]})
-
-    messages.append({"role": "user", "content": message})
-
-    try:
-        response = ""
-        for message in client.chat_completion(
-            messages,
-            max_tokens=max_tokens,
-            stream=True,
+def generate_text(prompt, max_length=512, temperature=0.7, top_p=0.9):
+    # 对输入进行编码
+    inputs = tokenizer(prompt, return_tensors="pt")
+    
+    if torch.cuda.is_available():
+        inputs = {k: v.cuda() for k, v in inputs.items()}
+    
+    # 生成文本
+    with torch.no_grad():
+        outputs = model.generate(
+            inputs["input_ids"],
+            max_new_tokens=max_length,
             temperature=temperature,
             top_p=top_p,
-        ):
-            token = message.choices[0].delta.content
-            response += token
-            yield response
-    except Exception as e:
-        yield f"发生错误: {str(e)}"
+            pad_token_id=tokenizer.pad_token_id,
+            eos_token_id=tokenizer.eos_token_id,
+        )
+    
+    # 解码输出
+    generated_text = tokenizer.decode(outputs[0], skip_special_tokens=True)
+    return generated_text
 
-# 创建聊天界面
-demo = gr.ChatInterface(
-    respond,
-    chatbot=gr.Chatbot(height=400),
-    textbox=gr.Textbox(placeholder="请输入您的问题...", container=False, scale=7),
-    title="MinimindZero 演示",
-    description="这是一个牧之训练的中文语言模型演示，您可以通过对话来体验模型的能力。",
-    additional_inputs=[
-        gr.Textbox(value="你是一个友好的中文助手。", label="系统提示词"),
-        gr.Slider(minimum=1, maximum=2048, value=512, step=1, label="最大生成长度"),
-        gr.Slider(minimum=0.1, maximum=4.0, value=0.7, step=0.1, label="温度"),
-        gr.Slider(
-            minimum=0.1,
-            maximum=1.0,
-            value=0.95,
-            step=0.05,
-            label="Top-p",
-        ),
+# 创建Gradio界面
+demo = gr.Interface(
+    fn=generate_text,
+    inputs=[
+        gr.Textbox(label="输入提示词", lines=3),
+        gr.Slider(minimum=1, maximum=1024, value=512, label="最大生成长度"),
+        gr.Slider(minimum=0.1, maximum=2.0, value=0.7, label="Temperature"),
+        gr.Slider(minimum=0.1, maximum=1.0, value=0.9, label="Top-p"),
     ],
-    theme="soft",
+    outputs=gr.Textbox(label="生成结果", lines=10),
+    title="MiniMind 文本生成",
+    description="一个简单的文本生成demo"
 )
 
 if __name__ == "__main__":
-    demo.launch()
+    demo.launch()