【限时免费】从本地脚本到生产级API：用FastAPI将BGE-M3打造成高可用文本检索服务...-优快云博客

从本地脚本到生产级API：用FastAPI将BGE-M3打造成高可用文本检索服务

【免费下载链接】bge-m3 BGE-M3，一款全能型多语言嵌入模型，具备三大检索功能：稠密检索、稀疏检索和多元向量检索，覆盖超百种语言，可处理不同粒度输入，从短句到长达8192个token的文档。通用预训练支持，统一微调示例，适用于多场景文本相似度计算，性能卓越，潜力无限。项目地址: https://gitcode.com/hf_mirrors/BAAI/bge-m3

引言

你是否已经能在本地用BGE-M3生成高质量的文本嵌入，并渴望将其强大的检索能力分享给你的网站或App用户？当一个强大的语言模型躺在你的硬盘里时，它的价值是有限的。只有当它变成一个稳定、可调用的API服务时，才能真正赋能万千应用。本文将手把手教你如何实现这一转变，将BGE-M3从本地脚本升级为生产级的API服务。

技术栈选型与环境准备

推荐框架：FastAPI

FastAPI是一个轻量级、高性能的Python Web框架，特别适合构建API服务。它的优势包括：

异步支持：天然支持异步请求处理，适合高并发场景。
自动文档生成：内置Swagger UI和OpenAPI支持，方便调试和文档管理。
类型安全：基于Pydantic的数据验证，减少运行时错误。

环境准备

创建一个干净的Python环境，并安装以下依赖库：

# requirements.txt
fastapi==0.103.1
uvicorn==0.23.2
FlagEmbedding==1.0.0
torch==2.0.1

运行以下命令安装依赖：

pip install -r requirements.txt

核心逻辑封装：适配BGE-M3的推理函数

模型加载与推理函数

我们将BGE-M3的推理逻辑封装为两个函数：load_model和run_inference。

from FlagEmbedding import BGEM3FlagModel

def load_model(model_name: str = "BAAI/bge-m3", use_fp16: bool = True):
    """
    加载BGE-M3模型。
    :param model_name: 模型名称或路径，默认为"BAAI/bge-m3"。
    :param use_fp16: 是否使用FP16加速推理，默认为True。
    :return: 加载后的模型实例。
    """
    model = BGEM3FlagModel(model_name, use_fp16=use_fp16)
    return model

def run_inference(model, sentences: list, batch_size: int = 12, max_length: int = 8192):
    """
    运行BGE-M3推理，生成密集嵌入。
    :param model: 已加载的BGE-M3模型。
    :param sentences: 输入的文本列表。
    :param batch_size: 批处理大小，默认为12。
    :param max_length: 最大序列长度，默认为8192。
    :return: 密集嵌入向量。
    """
    embeddings = model.encode(sentences, batch_size=batch_size, max_length=max_length)['dense_vecs']
    return embeddings

代码说明

load_model函数负责加载模型，支持FP16加速。
run_inference函数接收文本列表，返回对应的密集嵌入向量。输入为字符串列表，输出为NumPy数组。

API接口设计：优雅地处理输入与输出

设计API端点

我们将创建一个FastAPI应用，提供以下端点：

/embed：接收文本列表，返回密集嵌入向量。

from fastapi import FastAPI
from pydantic import BaseModel
import numpy as np

app = FastAPI()

class TextRequest(BaseModel):
    texts: list[str]

model = load_model()

@app.post("/embed")
async def embed_texts(request: TextRequest):
    """
    接收文本列表，返回密集嵌入向量。
    :param request: 包含文本列表的请求体。
    :return: 密集嵌入向量。
    """
    embeddings = run_inference(model, request.texts)
    return {"embeddings": embeddings.tolist()}

数据返回策略

我们选择将嵌入向量转换为列表返回，而不是直接返回NumPy数组。原因如下：

兼容性：列表是JSON原生支持的数据类型，方便客户端解析。
可读性：列表格式更直观，便于调试和测试。

实战测试：验证你的API服务

启动服务

运行以下命令启动FastAPI服务：

uvicorn main:app --reload

测试API

使用curl测试API：

curl -X POST "http://127.0.0.1:8000/embed" -H "Content-Type: application/json" -d '{"texts": ["What is BGE M3?", "Definition of BM25"]}'

或者使用Python requests库：

import requests

response = requests.post(
    "http://127.0.0.1:8000/embed",
    json={"texts": ["What is BGE M3?", "Definition of BM25"]}
)
print(response.json())

生产化部署与优化考量

部署方案

在生产环境中，建议使用以下组合：

Gunicorn + Uvicorn Worker：提高并发处理能力。
Docker：容器化部署，确保环境一致性。

优化建议

批处理优化：根据实际负载调整batch_size，平衡显存占用和推理速度。
模型量化：使用FP16或INT8量化进一步减少显存占用和提升推理速度。

结语

创作声明：本文部分内容由AI辅助生成（AIGC），仅供参考

【限时免费】 从本地脚本到生产级API：用FastAPI将BGE-M3打造成高可用文本检索服务...