1. 项目缘起:为什么我们需要一个自己的Embedding服务?
最近在折腾一些RAG(检索增强生成)或者语义搜索相关的项目时,我遇到了一个非常典型且恼人的问题:no embedding model is loaded. set rag_embedding_model to a valid sentence transformer model。这个错误提示就像一盆冷水,浇灭了我快速验证想法的热情。它背后反映的是一个更普遍的需求——我们总是需要一个稳定、可控、且能按需定制的文本转向量服务。
市面上的云服务当然方便,OpenAI的text-embedding-ada-002、百度的文心、阿里的通义,调用一个API就能拿到高质量的向量。但问题也随之而来:成本、网络延迟、数据隐私、模型固定无法微调。尤其是在做一些内部工具、离线应用或者对响应速度要求极高的场景时,依赖外部API就成了瓶颈。更别提当你兴致勃勃地拉下一个开源项目,准备跑起来看看效果,却卡在模型下载或环境配置上时的那种挫败感。
所以,我决定动手从零搭建一个属于自己的Embedding服务。这个“从零”并不意味着从零开始训练一个BERT,那不现实。我们的“从零”指的是:从选择一个成熟的开源Embedding模型开始,完成本地部署、服务化封装、性能优化,最终提供一个类似云服务那样可以通过HTTP接口调用的文本向量化服务。目标很明确:摆脱对外部服务的强依赖,掌握核心组件的自主权,并且能根据业务需求灵活调整模型、优化性能。
这次实践,我选择了目前中文社区口碑很好的BGE(BAAI General Embedding)系列模型,特别是BGE-M3和轻量级的BGE embedding 4b作为主要实验对象。整个过程涉及模型选型、环境搭建、服务框架选择、接口设计、性能压测以及一些实战中的“坑”和技巧。下面,我就把这套完整的实现路径和思考过程分享出来。
2. 核心组件选型:模型、框架与部署环境
搭建服务的第一步是选择“砖瓦”。我们需要确定三样东西:用哪个模型把文本变成向量?用什么框架来包装和提供这个能力?以及最终把这个服务放在哪里运行?
2.1 Embedding模型选型:为什么是BGE?
开源Embedding模型的选择很多,像Sentence-BERT、Instructor、E5系列都很优秀。我最终聚焦在BGE上,主要是基于以下几点考虑:
- 中文优化与社区热度:BGE由智源研究院推出,针对中文场景做了大量优化,在MTEB中文榜单上长期名列前茅。这意味着用它来处理中文文本,开箱即用的效果就很有保障。从网络热词
bge embedding、embedding 4b bge的搜索热度也能看出,它已经是中文开发者的事实标准之一。 - 模型规格丰富:BGE提供了从大到小各种规格的模型。例如:
BGE-M3:最新的多功能模型,支持稠密向量、稀疏向量和多重向量检索,能力全面但体积较大(约2.2GB)。BGE-large-zh-v1.5:经典的高质量稠密向量模型,适用于大多数检索和语义相似度任务。BGE embedding 4b:这里需要澄清一个常见的误解。4b并非指4亿参数,而是指模型文件名为bge-small-zh-v1.5的量化版本(可能指4-bit量化?),模型本身很小(约50MB),速度极快,非常适合对精度要求不高、但对延迟和资源极其敏感的场景。这正好解决了我们轻量级、快速部署的需求。
- 易用性:BGE模型完美兼容
sentence-transformers库,而后者是Python生态中处理句子嵌入的标杆库,API设计优雅,社区支持好。
注意:模型选择没有银弹。
BGE-M3功能强但耗资源,bge-small速度快但精度有损。我的建议是,生产环境可以先从BGE-large-zh开始,它在效果和资源消耗上取得了很好的平衡;对于需要快速原型验证或资源受限的环境,bge-small(即常说的embedding 4b)是绝佳的起点。
2.2 服务化框架选型:FastAPI的压倒性优势
把模型封装成HTTP服务,我们有几个选择:Flask、FastAPI、或者直接用gradio快速构建一个Web界面。对于纯API服务,FastAPI几乎是当前的最优解,原因如下:
- 性能卓越:基于
Starlette和Pydantic,异步支持原生且强大,天生适合IO密集型的推理服务(网络请求、模型加载都是IO)。 - 开发效率极高:自动生成交互式API文档(Swagger UI和ReDoc),类型提示(Type Hints)带来极佳的开发体验和代码可靠性。
- 生态契合:与机器学习部署库(如
ray serve,text-generation-inference)的理念很契合,社区活跃。
因此,我们的技术栈就确定为:sentence-transformers+FastAPI+Uvicorn(ASGI服务器)。这是一个轻量、高效、且易于维护的组合。
2.3 部署环境准备:避开第一个坑
环境是这一切的基础。这里最大的坑就是网络问题。sentence-transformers在第一次使用某个模型时,会自动从Hugging Face Hub下载模型。如果你身处国内网络环境,这个过程可能极其缓慢甚至失败,直接导致服务启动报错no embedding model is loaded。
解决方案不是修改代码,而是预先准备好模型。有两种推荐做法:
手动下载与离线加载:
- 通过镜像站(如hf-mirror.com)或能稳定访问的环境,提前下载好模型文件。
- 将模型文件夹(包含
pytorch_model.bin、config.json、tokenizer.json等)放置到服务器本地目录,例如./models/bge-large-zh-v1.5。 - 在代码中,初始化模型时指定本地路径:
from sentence_transformers import SentenceTransformer model = SentenceTransformer('/path/to/your/local/models/bge-large-zh-v1.5')
这是最稳定、最推荐的方式,尤其适合生产环境。
环境变量配置(备选):
- 如果还是希望通过库自动下载,可以设置HF镜像的环境变量:
export HF_ENDPOINT=https://hf-mirror.com - 但这并非百分百可靠,取决于镜像站的同步情况和网络状况。
- 如果还是希望通过库自动下载,可以设置HF镜像的环境变量:
我的实战经验是:对于核心的生产依赖,永远优先选择离线部署。这能避免在关键时刻(如服务重启、扩容)被网络问题“背刺”。准备好模型文件后,我们就可以开始编写服务核心代码了。
3. 服务核心实现:从模型加载到API暴露
有了清晰的选型和准备好的模型,接下来就是编码实现。我们的服务核心要做两件事:1. 高效加载并管理Embedding模型;2. 通过HTTP接口暴露向量化能力。
3.1 模型加载与单例模式
在Web服务中,我们肯定不希望每个请求都去加载一次模型(耗时耗内存)。正确的做法是在服务启动时加载一次,之后所有请求共享这个模型实例。这可以通过FastAPI的lifespan事件或直接在全局作用域初始化来实现。我更喜欢使用lifespan,因为它管理起来更清晰。
from contextlib import asynccontextmanager from fastapi import FastAPI from sentence_transformers import SentenceTransformer import numpy as np # 全局变量存放模型 _model = None @asynccontextmanager async def lifespan(app: FastAPI): # 启动时加载模型 global _model print("Loading Embedding Model...") # 这里替换为你的实际模型路径 _model = SentenceTransformer('./models/bge-large-zh-v1.5') # 可以在这里进行一次预热推理,避免第一次请求过慢 _model.encode("预热文本") print("Model loaded successfully.") yield # 关闭时清理(如果需要) print("Shutting down...") app = FastAPI(lifespan=lifespan)为什么用lifespan?它提供了明确的启动和关闭钩子,比在模块顶层直接写加载代码更优雅,也便于未来扩展(例如连接数据库池)。预热推理那一步很重要,因为PyTorch/TensorFlow在第一次推理时会有图构建等开销,预热能确保第一个真实请求的延迟不会异常高。
3.2 API接口设计:兼顾简单与灵活
Embedding服务的核心API通常很简单:输入文本,输出向量。但设计时需要考虑扩展性。
基础编码接口:
from pydantic import BaseModel from typing import List class EncodeRequest(BaseModel): texts: List[str] # 可扩展参数:是否归一化(归一化后便于计算余弦相似度) normalize_embeddings: bool = True # 可扩展参数:批处理大小(针对大量文本) batch_size: int = 32 @app.post("/encode") async def encode_text(request: EncodeRequest): """将文本列表编码为向量列表""" try: # 调用模型进行编码 embeddings = _model.encode( request.texts, normalize_embeddings=request.normalize_embeddings, batch_size=request.batch_size, show_progress_bar=False # 服务端不需要进度条 ) # 将numpy数组转换为列表 embeddings_list = embeddings.tolist() return {"embeddings": embeddings_list, "model": _model.get_sentence_embedding_dimension()} except Exception as e: return {"error": str(e)}关键点:
- 使用
List[str]支持批量处理,能极大提升吞吐量。 - 提供
normalize_embeddings参数。对于大多数语义相似度或检索任务,将向量归一化为单位向量是标准操作,这样余弦相似度就等于点积,计算更高效。 batch_size参数允许调用方根据自身文本长度和数量微调,以平衡内存和速度。
- 使用
健康检查与元信息接口:
@app.get("/health") async def health_check(): return {"status": "healthy", "model": _model.__class__.__name__} @app.get("/model_info") async def model_info(): return { "model_name": _model.model_name_or_path, "embedding_dimension": _model.get_sentence_embedding_dimension(), "max_seq_length": _model.max_seq_length }这两个接口对于服务运维至关重要。健康检查用于负载均衡器或K8s的存活探针;模型信息接口让客户端能动态获取向量维度,避免硬编码。
3.3 处理长文本:截断与分块策略
所有Embedding模型都有一个max_seq_length(例如BGE-large是512)。当文本超过这个长度时,必须处理。sentence-transformers的默认行为是静默截断,但这可能丢失重要信息。
更优的做法是在服务层提供明确的策略,并在API响应中告知客户端。我们可以修改/encode接口,增加一个truncation_strategy参数,或者更简单一点,在服务内部实现一个智能分块函数(对于远长于512字的文档)。
def smart_truncate(text: str, max_length: int = 500) -> str: """简单智能截断:尽量在句末截断""" if len(text) <= max_length: return text # 找到max_length之前的最后一个句号、问号或感叹号 truncate_at = text.rfind('。', 0, max_length) if truncate_at == -1: truncate_at = text.rfind('?', 0, max_length) if truncate_at == -1: truncate_at = text.rfind('!', 0, max_length) if truncate_at == -1: truncate_at = text.rfind('.', 0, max_length) # 英文句号 # 如果还是没找到,就在max_length处硬截断 truncate_at = truncate_at if truncate_at != -1 else max_length return text[:truncate_at + 1] # 包含截断的标点 # 在encode函数中调用 processed_texts = [smart_truncate(t, _model.max_seq_length) for t in request.texts] embeddings = _model.encode(processed_texts, ...)这是一个基础示例。对于真正的长文档检索,更好的做法是“分块-分别嵌入-再聚合”(例如使用langchain的文本分割器),但这通常在上游应用层完成,而非在基础的Embedding服务内。我们的服务应保持职责单一,主要提供基础的向量化能力。
4. 性能优化与生产级考量
一个能用的服务和一个好用的服务之间,隔着性能优化和稳定性建设。当你的服务开始接收真实流量时,以下几个方面的考量至关重要。
4.1 并发处理与异步优化
FastAPI是异步框架,但sentence_transformers的encode方法是CPU/GPU密集型的同步操作。如果在异步路径中直接调用,会阻塞整个事件循环,导致服务并发能力急剧下降。
解决方案:使用run_in_executor将同步的模型推理任务丢到线程池中执行,避免阻塞异步事件循环。
import asyncio from concurrent.futures import ThreadPoolExecutor # 创建线程池 _executor = ThreadPoolExecutor(max_workers=4) # worker数量根据CPU核心数调整 @app.post("/encode") async def encode_text(request: EncodeRequest): loop = asyncio.get_event_loop() try: # 将同步的model.encode函数放到线程池中运行 embeddings = await loop.run_in_executor( _executor, lambda: _model.encode( request.texts, normalize_embeddings=request.normalize_embeddings, batch_size=request.batch_size, show_progress_bar=False ) ) embeddings_list = embeddings.tolist() return {"embeddings": embeddings_list} except Exception as e: return {"error": str(e)}参数max_workers设置多少合适?这没有固定答案。如果模型推理是CPU瓶颈(例如在CPU机器上运行),设置成CPU核心数或稍多一点。如果是GPU推理,GPU本身是瓶颈,线程数可以略多于GPU流处理器数量,但主要目的是不让CPU成为调度瓶颈。建议通过压测来确定,观察GPU利用率和请求延迟。
4.2 批处理:吞吐量的关键
Embedding模型在GPU上运行时,批处理能极大提升吞吐量,因为GPU擅长并行计算。我们的API设计已经支持批量文本输入。但这里有一个重要的权衡点:批大小(batch_size)。
- 批大小太小:GPU算力无法被充分利用,大量时间浪费在kernel启动和数据传输上。
- 批大小太大:可能导致GPU内存溢出(OOM),特别是文本长度不一、动态padding后显存占用激增。
如何找到最佳批大小?需要实测。写一个脚本,用不同长度的文本和不同的批大小进行测试,监控GPU内存使用量和每秒处理的token数(或句子数)。一个常见的经验是,对于固定长度的输入,可以逐步增加批大小直到接近OOM,然后留出20%的安全余量。对于变长输入,可以设定一个“最大总token数”作为批处理的上限,而不是简单的句子数。
4.3 模型量化与轻量化部署
如果你对延迟和资源消耗极其敏感,或者需要在边缘设备部署,模型量化是必须考虑的步骤。前面提到的BGE embedding 4b很可能就是一个量化版本。
量化分为多种精度:FP16(半精度)、INT8、甚至INT4。sentence-transformers本身对量化支持有限,但我们可以借助其他工具:
- 使用ONNX Runtime:将模型导出为ONNX格式,并使用ONNX Runtime进行推理,它支持多种硬件加速和量化。
# 示例:使用 optimum 库导出 pip install optimum[onnxruntime] optimum-cli export onnx --model BAAI/bge-large-zh-v1.5 ./bge-large-zh-onnx/ - 使用Pytorch原生量化:对于高级用户,可以对模型进行动态量化或静态量化,但过程相对复杂。
- 直接使用社区提供的量化模型:在Hugging Face上搜索模型时,可以关注是否有
-int8、-4bit等后缀的版本。
量化带来的影响:精度会有轻微损失(对于Embedding任务,通常损失在可接受范围内),但模型体积和推理速度会有显著改善。建议在决定量化前,务必在你的业务数据集上评估量化前后向量相似度的保真度。
4.4 监控、日志与容错
一个健壮的服务离不开可观测性。
- 日志:使用标准的
logging模块,记录每个请求的摘要(如文本数量、总字符数、处理耗时)、错误信息。这有助于问题排查和用量分析。 - 监控指标:可以考虑集成
Prometheus客户端,暴露一些指标,如:requests_total:总请求数request_duration_seconds:请求耗时直方图batch_size_distribution:批大小分布text_length_distribution:输入文本长度分布
- 容错与限流:
- 输入验证:使用Pydantic严格校验输入,防止恶意或异常数据导致服务崩溃。
- 超时控制:在FastAPI层面或反向代理(如Nginx)设置请求超时,防止慢请求拖垮服务。
- 限流:对于公开服务,必须实施限流(如使用
slowapi或fastapi-limiter),防止被滥用。 - 优雅降级:如果模型加载失败或GPU不可用,是否有后备方案(如降级到CPU模式或返回特定错误码)?
5. 实战部署与测试:让服务跑起来
代码写好了,我们需要把它部署到一个稳定的环境中,并验证其功能和性能。
5.1 使用Docker容器化部署
Docker是保证环境一致性的最佳实践。编写一个Dockerfile:
# 使用带有CUDA的Python基础镜像(如果使用GPU) FROM pytorch/pytorch:2.1.0-cuda11.8-cudnn8-runtime # 或使用CPU镜像 # FROM python:3.10-slim WORKDIR /app # 复制依赖文件并安装 COPY requirements.txt . RUN pip install --no-cache-dir -r requirements.txt -i https://pypi.tuna.tsinghua.edu.cn/simple # 复制模型文件(假设模型已下载到本地./models目录) COPY ./models /app/models # 复制应用代码 COPY . . # 暴露端口 EXPOSE 8000 # 启动命令 CMD ["uvicorn", "main:app", "--host", "0.0.0.0", "--port", "8000", "--workers", "4"]requirements.txt内容:
fastapi==0.104.1 uvicorn[standard]==0.24.0 sentence-transformers==2.2.2 numpy==1.24.3 pydantic==2.5.0构建与运行:
# 构建镜像 docker build -t embedding-service . # 运行容器(CPU版本) docker run -p 8000:8000 --name embedding-api embedding-service # GPU版本需要加参数 docker run --gpus all -p 8000:8000 --name embedding-api embedding-service5.2 功能与性能测试
服务启动后(访问http://localhost:8000/docs可以看到自动生成的API文档),我们需要进行测试。
基础功能测试:使用
curl或Python的requests库调用接口。curl -X POST "http://localhost:8000/encode" \ -H "Content-Type: application/json" \ -d '{"texts": ["今天天气真好", "人工智能是未来科技的核心"], "normalize_embeddings": true}'检查返回的向量维度是否正确(例如BGE-large是1024维),以及两个语义不太相关的句子的向量余弦相似度是否较低。
压力测试:使用
wrk、locust或apache benchmark进行压测。# 使用wrk示例 wrk -t4 -c100 -d30s --script=post.lua --latency http://localhost:8000/encode在
post.lua文件中定义POST请求体和Header。通过压测,你可以找到服务的QPS(每秒查询数)上限,以及在不同并发下的延迟分布(P50, P95, P99)。这是调整workers数量、线程池大小和批处理参数的核心依据。长文本与边界测试:输入超长文本、空文本列表、特殊字符等,观察服务是否稳定,响应是否符合预期(如截断或返回错误信息)。
5.3 集成到现有项目
最后,如何在你自己的RAG或搜索项目中使用这个服务?非常简单,只需将原来直接调用云API或本地模型库的代码,替换为HTTP调用。
# 以前:直接使用sentence-transformers # from sentence_transformers import SentenceTransformer # model = SentenceTransformer('model_name') # embeddings = model.encode(texts) # 现在:调用自建服务 import requests import numpy as np def encode_with_service(texts, service_url="http://localhost:8000"): resp = requests.post( f"{service_url}/encode", json={"texts": texts, "normalize_embeddings": True} ) resp.raise_for_status() result = resp.json() return np.array(result["embeddings"]) # 使用 vectors = encode_with_service(["查询文本1", "查询文本2"])这样,你的应用就和Embedding模型的具体实现解耦了。未来如果需要切换模型(比如从BGE-large换成BGE-M3),或者进行模型版本升级,只需要重启或替换后端的Embedding服务,而无需修改所有上游应用的代码。
从选择一个合适的开源模型,到用FastAPI将其封装成服务,再到考虑性能、并发和生产部署的方方面面,这个过程让我对“Embedding即服务”有了更深的体会。它不再是一个神秘的黑盒,而是一个可以根据自己需求随意拆解、组装和优化的组件。最重要的是,当再次看到no embedding model is loaded的报错时,你心里有底,知道问题出在哪,并且有能力去解决它。这种掌控感,或许就是自建基础设施最大的乐趣和价值所在。