Infinity 向量嵌入服务

镜像概述和主要用途

Infinity 是一个高吞吐量、低延迟的 REST API，用于提供向量嵌入服务，支持所有 sentence-transformer 模型和框架。该项目基于 MIT 许可证开发，为 Gradient.ai 等平台提供推理支持。

Infinity 的核心功能是将文本转换为低维密集向量，这些向量可用于检索、分类、聚类或语义搜索等任务，也可用于为大型语言模型(LLM)提供向量数据库支持。

核心功能和特性

主要特性

• 支持多种模型部署：可部署 MTEB 中的任何嵌入模型，兼容 SentenceTransformers 生态系统
• 多推理后端支持：基于 torch、optimum(onnx/tensorrt) 和 CTranslate2 构建，利用 FlashAttention 优化性能
• 多硬件支持：支持 NVIDIA CUDA、AMD ROCM、CPU、AWS INF2 和 APPLE MPS 加速器
• 动态批处理：GPU 忙碌时将新嵌入请求排队，设备就绪后立即处理新请求
• 经过验证的实现：单元测试和端到端测试确保嵌入结果准确性
• 易用 API：基于 FastAPI 构建，集成 Swagger 文档，API 与 OpenAI 的嵌入规范对齐

最新功能

• 2024/06：支持多模态(clip)、文本分类及环境变量配置
• 2024/05：支持多模型部署(v2 CLI)和 API 密钥认证
• 2024/03：实验性支持 int8(CPU/CUDA)和 fp8(H100/MI300)精度
• 2024/01：支持 TensorRT/ONNX 推理
• 动态批处理优化和多模型并发部署

使用场景和适用范围

适用场景

• 语义搜索：将查询和文档转换为向量，实现高效相似性搜索
• 文本分类：利用嵌入向量进行文本分类任务
• 聚类分析：将相似文本聚合成组
• 推荐系统：基于内容相似性提供推荐
• 问答系统：提高问答系统的上下文理解能力
• 向量数据库集成：为 Milvus、Pinecone 等向量数据库提供嵌入支持
• LLM 应用增强：为大型语言模型提供上下文检索能力

支持的模型类型

• 嵌入模型：所有 sentence-transformers 模型，如 BAAI/bge-small-en-v1.5、mixedbread-ai/mxbai-embed-large-v1 等
• 重排序模型：如 mixedbread-ai/mxbai-rerank-xsmall-v1、BAAI/bge-reranker-base 等
• 多模态模型：CLIP 模型，支持图像和文本同时编码
• 文本分类模型：如 ProsusAI/finbert、SamLowe/roberta-base-go_emotions 等

详细的使用方法和配置说明

Docker 部署（推荐）

使用预构建的 Docker 容器是部署 Infinity 的推荐方式。确保已安装 nvidia-docker 以支持 GPU 加速。

基本部署命令

port=7997
model="BAAI/bge-small-en-v1.5"
volume=$PWD/data

docker run -it --gpus all \
 -v $volume:/app/.cache \
 -p $port:$port \
 michaelf34/infinity:latest \
 v2 \
 --model-id $model \
 --port $port

部署多个模型

port=7997
model1="michaelfeil/bge-small-en-v1.5"
model2="mixedbread-ai/mxbai-rerank-xsmall-v1"
volume=$PWD/data

docker run -it --gpus all \
 -v $volume:/app/.cache \
 -p $port:$port \
 michaelf34/infinity:latest \
 v2 \
 --model-id $model1 \
 --model-id $model2 \
 --port $port

环境变量配置

• HF_HOME: 容器内缓存路径，默认为 /app/.cache
• INFINITY_MODEL_ID: 模型ID，多个模型用分号分隔
• INFINITY_PORT: 服务端口号

CLI 参数说明

使用 --help 命令查看所有可用参数：

docker run michaelf34/infinity:latest v2 --help

主要参数：

• --model-id: 模型ID或路径，可指定多个
• --port: 服务端口号
• --engine: 推理引擎，可选值: torch, optimum, ctranslate2
• --embedding-dtype: 嵌入输出数据类型
• --dtype: 模型计算数据类型
• --api-key: API密钥认证
• --host: 绑定主机地址

API 使用示例

嵌入生成

curl -X POST "http://localhost:7997/v1/embeddings" \
  -H "Content-Type: application/json" \
  -d '{
    "input": ["这是一个测试句子", "巴黎在法国"],
    "model": "BAAI/bge-small-en-v1.5"
  }'

重排序

curl -X POST "http://localhost:7997/v1/rerank" \
  -H "Content-Type: application/json" \
  -d '{
    "query": "什么是Python包infinity_emb?",
    "documents": [
      "这是一个与Python包infinity_emb无关的文档",
      "巴黎在法国!",
      "infinity_emb是一个用于句子嵌入和重排序的Python包"
    ],
    "model": "mixedbread-ai/mxbai-rerank-xsmall-v1"
  }'

Python API 使用

除了 REST API，还可以通过 Python API 直接使用 Infinity：

import asyncio
from infinity_emb import AsyncEngineArray, EngineArgs

sentences = ["通过Infinity嵌入这句话", "巴黎在法国。"]
array = AsyncEngineArray.from_args([
  EngineArgs(model_name_or_path="BAAI/bge-small-en-v1.5", engine="torch")
])

async def embed_text(): 
    async with array[0]: 
        embeddings, usage = await array[0].embed(sentences=sentences)
        print("嵌入结果:", embeddings)
        print("使用情况:", usage)

asyncio.run(embed_text())

多模型部署

从版本 0.0.34 开始，Infinity 支持同时部署多个模型：

docker run -it --gpus all \
 -p 7997:7997 \
 michaelf34/infinity:latest \
 v2 \
 --model-id BAAI/bge-small-en-v1.5 \
 --model-id mixedbread-ai/mxbai-rerank-xsmall-v1 \
 --port 7997

集成方案

Infinity 可与多种工具和框架集成：

• Serverless 部署：Runpod 上的无服务器部署
• Langchain 集成：作为 Langchain 的嵌入组件
• 向量数据库：与 Milvus、Qdrant 等向量数据库配合使用
• 云部署：通过 dstack 在云平台上部署
• RAG 应用：为检索增强生成应用提供支持

常见问题

什么是嵌入模型？

嵌入模型可以将任何文本映射到低维密集向量，用于检索、分类、聚类或语义搜索等任务，也可用于大型语言模型的向量数据库。最知名的架构是仅编码器的Transformer（如BERT），最流行的实现包括SentenceTransformers。

支持哪些模型？

所有 sentence-transformers 组织的模型都受支持。使用不同引擎时要求不同：

• --engine torch: 模型必须与 SentenceTransformers 和 AutoModel 兼容
• --engine optimum: 必须有 ONNX 文件，推荐使用 Xenova 提供的模型
• --engine ctranslate2: 仅支持 BERT 模型

最新的模型趋势可查看 MTEB 排行榜：[***]

如何与 Langchain 集成？

Infinity 已正式集成到 Langchain（需要 langchain>=0.342）：

from langchain.embeddings.infinity import InfinityEmbeddings
from langchain.docstore.document import Document

documents = [Document(page_content="Hello world!", metadata={"source": "unknown"})]

emb_model = InfinityEmbeddings(model="BAAI/bge-small", infinity_api_url="http://localhost:7997/v1")
print(emb_model.embed_documents([doc.page_content for doc in documents]))

文档和资源

• 完整文档：[***]
• 启动后 Swagger UI：http://localhost:7997/docs
• GitHub 仓库：[***]