Update app.py

app.py

@@ -1,58 +1,98 @@
 import gradio as gr
 from fastapi import FastAPI, Request
 import uvicorn
-from transformers import AutoTokenizer, AutoModel
-import torch
 import os
+from huggingface_hub import InferenceClient
+import logging
 
-# 初始化接口
-app = FastAPI()
-
-# --- 模型配置 ---
-# 如果同学本地有模型文件，可以改成文件夹路径
-MODEL_PATH = "jiang1002/chatglm-6b-adgen" 
+# 设置日志
+logging.basicConfig(level=logging.INFO)
+logger = logging.getLogger(__name__)
 
-print("🚀 正在加载模型，请稍候...")
+# 初始化 FastAPI
+app = FastAPI()
 
-tokenizer = AutoTokenizer.from_pretrained(MODEL_PATH, trust_remote_code=True)
+# 从环境变量获取 Hugging Face Token
+HF_TOKEN = os.getenv("HF_TOKEN")
+if not HF_TOKEN:
+    logger.warning("⚠️ 未设置 HF_TOKEN 环境变量，API 调用可能失败")
 
-# 自动检测设备：有显卡用显卡，没显卡用 CPU
-if torch.cuda.is_available():
-    print("✨ 检测到 GPU，正在使用显卡加速...")
-    model = AutoModel.from_pretrained(MODEL_PATH, trust_remote_code=True).half().cuda()
-else:
-    print("☁️ 未检测到 GPU，正在使用 CPU 模式（速度较慢）...")
-    model = AutoModel.from_pretrained(MODEL_PATH, trust_remote_code=True).float()
+# 初始化 InferenceClient
+# 这里用 provider="auto" 让系统自动选择可用提供商
+client = InferenceClient(provider="auto", token=HF_TOKEN)
 
-model.eval()
-print("✅ 模型加载成功！")
+# 你的模型名称
+MODEL_ID = "jiang1002/chatglm-6b-adgen"  # 或者换成其他公开模型
 
-# --- 1. 给别人用的 API 接口 ---
+# --- 1. API 接口 ---
 @app.post("/generate")
 async def generate(request: Request):
     try:
         data = await request.json()
         prompt = data.get("text", "")
-        # 模型推理
-        response, _ = model.chat(tokenizer, prompt, history=[])
-        return {"success": True, "result": response}
+        messages = data.get("messages", [])
+        
+        # 如果提供了完整的 messages 格式，就用它
+        if messages:
+            response = client.chat.completions.create(
+                model=MODEL_ID,
+                messages=messages
+            )
+            result = response.choices[0].message.content
+        else:
+            # 否则用简单的 prompt 格式
+            response = client.text_generation(
+                prompt,
+                model=MODEL_ID,
+                max_new_tokens=512,
+                temperature=0.7
+            )
+            result = response
+            
+        return {"success": True, "result": result}
     except Exception as e:
+        logger.error(f"API 调用失败: {str(e)}")
         return {"success": False, "error": str(e)}
 
-# --- 2. 给自己用的网页界面 ---
-def chat_func(msg, hist):
-    res, _ = model.chat(tokenizer, msg, history=hist)
-    return res
+# --- 2. Gradio 聊天界面 ---
+def chat_func(message, history):
+    """Gradio 聊天函数"""
+    try:
+        # 将历史记录转换为 messages 格式
+        messages = []
+        for human, assistant in history:
+            messages.append({"role": "user", "content": human})
+            messages.append({"role": "assistant", "content": assistant})
+        messages.append({"role": "user", "content": message})
+        
+        # 调用 Inference API
+        response = client.chat.completions.create(
+            model=MODEL_ID,
+            messages=messages,
+            max_tokens=512,
+            temperature=0.7
+        )
+        
+        return response.choices[0].message.content
+    except Exception as e:
+        logger.error(f"聊天失败: {str(e)}")
+        return f"调用失败: {str(e)}"
 
+# 创建 Gradio 界面
 demo = gr.ChatInterface(
-    fn=chat_func, 
-    title="ChatGLM 广告生成助手",
-    description="本程序已同时开启 API 接口（路径：/generate）"
+    fn=chat_func,
+    title="ChatGLM 广告生成助手 (使用 Inference Providers)",
+    description="后台使用 Hugging Face Inference Providers，无需本地 GPU"
 )
 
-# 挂载 Gradio 到 FastAPI
+# 挂载 Gradio
 app = gr.mount_gradio_app(app, demo, path="/")
 
+# 添加健康检查端点
+@app.get("/health")
+async def health():
+    return {"status": "ok", "model": MODEL_ID}
+
 if __name__ == "__main__":
-    # 启动服务器，默认端口 7860
-    uvicorn.run(app, host="0.0.0.0", port=7860)
+    port = int(os.getenv("PORT", 7860))
+    uvicorn.run(app, host="0.0.0.0", port=port)