【FastAPI】利用FastAPI构建大模型接口服务

最新推荐文章于 2025-06-05 09:02:43 发布

原创最新推荐文章于 2025-06-05 09:02:43 发布 · 7.6k 阅读

46 ·

CC 4.0 BY-SA版权

文章标签：

#fastapi

该文介绍了如何利用FastAPI和uvicorn搭建一个本地化的接口服务，将大语言模型（如transformers库中的模型）与应用解耦。服务端通过AutoModelForCausalLM和AutoTokenizer处理输入文本，生成响应，客户端则通过HTTP请求与服务端交互进行对话。

摘要生成于 C知道，由 DeepSeek-R1 满血版支持，前往体验 >

0. 摘要：

利用FastAPI和uvicorn构建本地化部署大语言模型的接口服务，实现大模型与后续应用（如langchain等）的解耦。整个流程分为服务端与客户端。

1. 服务端

import uvicorn
from fastapi import FastAPI
from pydantic import BaseModel
import uvicorn, json, datetime
import torch
import os
from transformers import AutoModelForCausalLM, AutoTokenizer
from transformers.generation.utils import GenerationConfig
os.environ['CUDA_VISIBLE_DEVICES'] = "5,6,7"

app = FastAPI()

class Query(BaseModel):
    text: str

path = "/workdir/model/baichuan13b_chat/"
tokenizer = AutoTokenizer.from_pretrained(path, use_fast=False, trust_remote_code=True)
model = AutoModelForCausalLM.from_pretrained(path, device_map="auto", torch_dtype=torch.float16, trust_remote_code=True)
model.generation_config = GenerationConfig.from_pretrained(path)


@app.post("/chat/")
async def chat(query: Query):
    input_ids = tokenizer([query.text]).input_ids
    output_ids = model.generate(
        torch.as_tensor(input_ids).cuda(),
        do_sample=False,
        temperature=0.1,
        repetition_penalty=1,
        max_new_tokens=1024)
    output_ids = output_ids[0][len(input_ids[0]):]
    outputs = tokenizer.decode(output_ids, skip_special_tokens=True, spaces_between_special_tokens=False)
    return {"result": outputs}


if __name__ == "__main__":
    uvicorn.run(app, host="0.0.0.0", port=6667)

2. 客户端

import requests

url = "http://0.0.0.0:6667/chat/"
query = {"text": "你好，请做一段自我介绍。"}

response = requests.post(url, json=query)

if response.status_code == 200:
    result = response.json()
    print("BOT:", result["result"])
else:
    print("Error:", response.status_code, response.text)

在这里插入图片描述