Upload 2 files (#1)

- Upload 2 files (85ea1dda7c5244f8926fdf8c902d37d1462db21c)

app.py

@@ -1,10 +1,80 @@
-import gradio as gr
-
-with gr.Blocks(fill_height=True) as demo:
-    with gr.Sidebar():
-        gr.Markdown("# Inference Provider")
-        gr.Markdown("This Space showcases the WiroAI/WiroAI-Finance-Qwen-1.5B model, served by the hf-inference API. Sign in with your Hugging Face account to use this API.")
-        button = gr.LoginButton("Sign in")
-    gr.load("models/WiroAI/WiroAI-Finance-Qwen-1.5B", accept_token=button, provider="hf-inference")
-    
-demo.launch()
+from fastapi import FastAPI, HTTPException
+from fastapi.middleware.cors import CORSMiddleware
+from pydantic import BaseModel
+import transformers
+import torch
+
+# ---- 初始化模型 ----
+model_id = "WiroAI/WiroAI-Finance-Qwen-1.5B"
+
+# 缓存 pipeline 避免重复加载
+_pipeline = None
+
+
+def get_pipeline():
+    global _pipeline
+    if _pipeline is None:
+        _pipeline = transformers.pipeline(
+            "text-generation",
+            model=model_id,
+            model_kwargs={"torch_dtype": torch.bfloat16},
+            device_map="auto" if torch.cuda.is_available() else None,
+        )
+        _pipeline.model.eval()
+    return _pipeline
+
+
+# ---- FastAPI 配置 ----
+app = FastAPI(title="WiroAI Finance Chat API")
+
+# 允许跨域（重要！）
+app.add_middleware(
+    CORSMiddleware,
+    allow_origins=["*"],
+    allow_methods=["*"],
+    allow_headers=["*"],
+)
+
+
+# ---- 请求/响应模型 ----
+class ChatRequest(BaseModel):
+    message: str
+    max_new_tokens: int = 512
+    temperature: float = 0.9
+
+
+class ChatResponse(BaseModel):
+    content: str
+
+
+# ---- API 端点 ----
+@app.post("/chat", response_model=ChatResponse)
+async def generate_response(request: ChatRequest):
+    try:
+        pipeline = get_pipeline()
+
+        # 构建对话历史
+        messages = [
+            {"role": "system", "content": "You are a finance chatbot developed by Wiro AI"},
+            {"role": "user", "content": request.message}
+        ]
+
+        # 终止符配置
+        terminators = [
+            pipeline.tokenizer.eos_token_id,
+            pipeline.tokenizer.convert_tokens_to_ids("")
+        ]
+
+        # 生成回复
+        outputs = pipeline(
+            messages,
+            max_new_tokens=request.max_new_tokens,
+            eos_token_id=terminators,
+            do_sample=True,
+            temperature=request.temperature,
+        )
+
+        return {"content": outputs[0]["generated_text"][-1]['content']}
+
+    except Exception as e:
+        raise HTTPException(status_code=500, detail=str(e))

requirements

@@ -0,0 +1,6 @@
+fastapi>=0.68.0
+uvicorn>=0.15.0
+transformers>=4.40.0
+torch>=2.3.0
+accelerate>=0.30.0
+pydantic>=2.0