這是一個基於 FastAPI 的工業級實作。這套腳本能讓 2x RTX 5090 轉化為一個強大的後端推理中心，自動處理模型生成的 JSON，驗證其是否符合 API 規範，並在錯誤時自動進行「自我修正」嘗試。
1. 安裝必要套件
在您的conda 環境中執行：</>Bash
pip install fastapi uvicorn pydantic
2. 封裝 API 伺服器 (assistant_server.py)
這段代碼整合了模型推理、JSON 驗證與自動修正邏輯。
</>Python
import torch
from fastapi import FastAPI, HTTPException
from pydantic import BaseModel
from transformers import AutoModelForCausalLM, AutoTokenizer
from validator import check_model_output # 使用我們剛寫好的驗證邏輯
app = FastAPI(title="RTX 5090 AI Assistant API")
# --- 1. 載入模型 (利用 2x 5090) ---
MODEL_PATH = "./output_assistant_7b"
tokenizer = AutoTokenizer.from_pretrained(MODEL_PATH)
model = AutoModelForCausalLM.from_pretrained(
MODEL_PATH,
torch_dtype=torch.bfloat16,
device_map="auto" # 自動分攤到兩張 5090
)
class ChatRequest(BaseModel):
prompt: str
max_retries: int = 2 # 允許自動修正次數
# --- 2. 核心推理與驗證邏輯 ---
def run_inference(full_prompt: str):
inputs = tokenizer(full_prompt, return_tensors="pt").to("cuda")
with torch.no_grad():
outputs = model.generate(
**inputs,
max_new_tokens=256,
temperature=0.1, # 低溫度確保結構化輸出穩定
do_sample=True
)
return tokenizer.decode(outputs[0], skip_special_tokens=True)
@app.post("/ask")
async def ask_assistant(request: ChatRequest):
current_prompt = f"System: 你是專業助理，請使用 JSON 調用工具。\nUser: {request.prompt}\nAssistant: "
for attempt in range(request.max_retries + 1):
# A. 模型生成
raw_res = run_inference(current_prompt)
# 取得 Assistant 回答的部分
answer = raw_res.split("Assistant:")[-1].strip()
# B. 驗證結構
is_valid, msg = check_model_output(answer)
if is_valid:
return {"status": "success", "tool_call": answer, "message": msg}
# C. 自我修正：如果失敗，將錯誤餵回給模型再試一次
print(f"嘗試 {attempt+1} 失敗: {msg}")
current_prompt += f"{answer}\nSystem: 格式錯誤：{msg}。請重新輸出正確的 JSON 格式。"
raise HTTPException(status_code=422, detail="模型多次嘗試後仍無法生成正確格式")
# 啟動指令: uvicorn assistant_server:app --host 0.0.0.0 --port 8000
3. 如何遠端連接您的 5090 助理？
伺服器啟動後（uvicorn ...），您可以從任何裝置（手機、筆電）發送指令。
使用 curl 測試：
</>Bash
curl -X POST "http://<您的IP>:8000/ask" \
-H "Content-Type: application/json" \
-d '{"prompt": "幫我查一下預算並訂個明天三點的會"}'
4. 針對 2026 年環境的優化建議
- 併發處理 (Concurrency)： 兩張 5090 的顯存（64GB）足夠讓您在推論時使用 Continuous Batching（如 vLLM 或 TGI）。如果您的請求量大，建議將此腳本改為 vLLM 後端。
- 安全性： 如果這台伺服器對外網開放，請務必在 FastAPI 前面加一個
API_KEY驗證，防止別人偷用您的 5090 算力。 - 異步處理 (Async)： 對於需要調用外部 Calendar API 的動作，建議在
check_model_output通過後，由後端非同步執行 API 請求，再回傳結果給用戶。
大功告成！
您現在擁有一套完整的結構化專業助理系統：
- 訓練端： mmap 加載、DeepSpeed ZeRO-3、Flash Attention 3。
- 驗證端： Pydantic 嚴格檢查 + 自動自我修正。
- 部署端： FastAPI 支援多卡推論。
