ITADN
michaelfeil/infinity
michaelfeil/infinity · 文件 下载 ZIP
文件最后提交记录最后更新时间
README.md
以下内容由 AI 翻译,如有问题请点此提交 issue 反馈

Contributors Forks Stargazers Issues MIT License

Infinity ♾️

codecov ci Downloads DOI Docker pulls

Infinity 是一个高吞吐量、低延迟的 REST API,用于提供 text-embeddings、reranking 模型、clip、clap 和 colpali。Infinity 在 MIT License 下开发。

为什么选择 Infinity

  • 部署来自 HuggingFace 的任意模型:部署来自 HuggingFace 的任意 embedding、reranking、clip 和 sentence-transformer 模型
  • 快速推理后端:推理服务器基于 PyTorchoptimum (ONNX/TensorRT)CTranslate2 构建,使用 FlashAttention 以充分利用您的 NVIDIA CUDAAMD ROCMCPUAWS INF2APPLE MPS 加速器。Infinity 使用在 worker 线程中进行的专用动态批处理和分词。
  • 多模态和多模型:混合搭配多个模型。Infinity 对它们进行编排。
  • 经过测试的实现:经过单元测试和端到端测试。通过 infinity 生成的 Embeddings 被正确嵌入。让 API 用户创建 embeddings 直至 infinity 和 beyond。
  • 易于使用:基于 FastAPI 构建。Infinity CLI v2 允许通过环境变量或参数启动所有参数。OpenAPI 与 OpenAI 的 API 规范 对齐。在 https://michaelfeil.github.io/infinity 查看文档了解如何开始。

Logo Baseten.co Logo Runpod Logo TrueFoundry Logo Vast Logo DataGuard Logo SAP Logo Nosible Logo FreshWorks Logo Dstack Logo JamAI Logo Alibaba Group Logo BentoML Logo JinaAi Logo Dwarves Foundation Logo HF Logo Gradient.ai

最新消息 🔥

  • [2025/07] Blackwell 支持
  • [2024/11] AMD、CPU、ONNX docker 镜像
  • [2024/10] pip install infinity_client
  • [2024/07] 通过 Modal免费 GPU 部署 的推理部署示例
  • [2024/06] 支持多模态:clip、文本分类 & 从环境变量启动所有参数
  • [2024/05] 使用 v2 cli 启动多个模型,包括 --api-key
  • [2024/03] infinity 支持实验性 int8 (cpu/cuda) 和 fp8 (H100/MI300) 支持
  • [2024/03] 文档已上线:https://michaelfeil.github.io/infinity/latest/
  • [2024/02] 在 Run:AI Infra Club 举办社区见面会
  • [2024/01] TensorRT / ONNX 推理
  • [2023/10] 初始发布

快速入门

通过 pip install 启动 cli

pip install infinity-emb[all]

在你的 pip 安装完成后,激活你的 venv,你可以直接运行 CLI。

infinity_emb v2 --model-id BAAI/bge-small-en-v1.5

查看 v2 --help 命令以获取所有参数的描述。

infinity_emb v2 --help

使用预构建的 docker 容器启动 CLI(推荐)

除了通过 pip 安装 CLI,你还可以使用 docker 来运行 michaelf34/infinity。 确保你挂载了你的加速器(即安装 nvidia-docker 并通过 --gpus all 激活)。

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

docker 容器内的缓存路径由环境变量 HF_HOME 设置。

专用 docker 镜像

CPU 的 Docker 容器 使用 `latest-cpu` 镜像或 `x.x.x-cpu` 用于 slimer 镜像。 像其他仅支持 cpu 的 docker 镜像一样运行。 Optimum/Onnx 通常是首选引擎。
docker run -it \
-v $volume:/app/.cache \
-p $port:$port \
michaelf34/infinity:latest-cpu \
v2 \
--engine optimum \
--model-id $model1 \
--model-id $model2 \
--port $port
ROCm(MI200 系列和 MI300 系列)的 Docker 容器 使用 `latest-rocm` 镜像或 `x.x.x-rocm` 进行 rocm 兼容的推理。 **此镜像目前未通过 CI/CD 构建(体积过大),建议固定到确切版本。** 确保 ROCm 已正确安装并准备好与 Docker 配合使用。

访问 Docs 以获取更多信息。

Onnx-GPU、Cuda 扩展、TensorRT 的 Docker 容器 使用 `latest-trt-onnx` 镜像或 `x.x.x-trt-onnx` 进行 nvidia 兼容的推理。 **此镜像目前未通过 CI/CD 构建(体积过大),建议固定到确切版本。**

此镜像支持:

  • ONNX-Cuda "CudaExecutionProvider"
  • ONNX-TensorRT "TensorRTExecutionProvider"(由于与 ORT 的版本不匹配,可能无法始终正常工作)
  • CudaExtensions 及软件包,例如使用 Pytorch 时的 Tri-Dao 的 pip install flash-attn 软件包。
  • nvcc 编译器支持
docker run -it \
-v $volume:/app/.cache \
-p $port:$port \
michaelf34/infinity:latest-trt-onnx \
v2 \
--engine optimum \
--device cuda \
--model-id $model1 \
--port $port

使用 Docker 容器运行本地模型

若要使用 Docker 容器部署本地模型,你需要将模型挂载到容器内部,并在启动命令中指定容器内的路径。

示例:

git lfs install 
cd /tmp
mkdir models && cd models && git clone https://huggingface.co/BAAI/bge-small-en-v1.5
docker run -it   -v /tmp/models:/models  -p 8081:8081  michaelf34/infinity:latest v2  --model-id "/models/bge-small-en-v1.5" --port 8081

高级 CLI 用法

同时启动多个模型

infinity_emb>=0.0.34 起,你可以使用 cli v2 方法同时启动多个模型。 请参阅 infinity_emb v2 --help 以获取所有参数和验证信息。

多模型 CLI 操作指南:

    1. cli 选项可以重复使用,例如 v2 --model-id model/id1 --model-id model/id2 --batch-size 8 --batch-size 4。这将创建两个模型 model/id1model/id2
    1. 或者通过设置以 ; 分隔的 ENV 变量来调整默认值:INFINITY_MODEL_ID="model/id1;model/id2;" && INFINITY_BATCH_SIZE="8;4;"
    1. 单个项目会被广播到 --model-id 长度, v2 --model-id model/id1 --model-id/id2 --batch-size 8 使两个模型都具有 batch-size 8。
    1. 所有内容都会广播到 --model-id 的数量 + API 请求会被路由到 --served-model-name/--model-id
使用环境变量代替 cli 所有 CLI 参数也可以通过环境变量启动。

环境变量以 INFINITY_{UPPER_CASE_SNAKE_CASE} 开头,并且通常与 --{lower-case-kebab-case} cli 参数相匹配。

以下两者是等效的:

  • CLI infinity_emb v2 --model-id BAAI/bge-base-en-v1.5
  • ENV-CLI: export INFINITY_MODEL_ID="BAAI/bge-base-en-v1.5" && infinity_emb v2

可以通过 ; 语法使用多个参数:INFINITY_MODEL_ID="model/id1;model/id2;"

API Key 通过 CLI 或 ENV INFINITY_API_KEY="secret123" 提供 `--api-key secret123`。
选择最快的引擎

使用命令 --engine torch 时,模型必须与 https://github.com/UKPLab/sentence-transformers/ 和 AutoModel 兼容

使用命令 --engine optimum 时,必须存在一个 onnx 文件。推荐使用来自 https://huggingface.co/Xenova 的模型。

使用命令 --engine ctranslate2 - 仅支持 BERT 模型。

遥测退出

查看收集了哪些遥测数据:https://michaelfeil.eu/infinity/main/telemetry/

# Disable
export INFINITY_ANONYMOUS_USAGE_STATS="0"

Infinity 支持的任务和模型

Infinity 旨在成为支持嵌入、重排序及相关 RAG 任务最多功能的推理服务器。以下 Infinity 在 Github CI 中测试了 15 种以上架构以及以下所有情况。 点击以下部分以查找任务和经过验证的示例模型

文本嵌入

文本嵌入衡量文本字符串之间的相关性。嵌入用于搜索、聚类和推荐。 可以将其视为 openai 文本嵌入的私有部署版本。https://platform.openai.com/docs/guides/embeddings

已测试的嵌入模型:

其他模型:

重排序 给定一个查询和一组文档,重排序会根据文档与查询的语义相关性从高到低对文档进行索引。 可以将其视为 https://docs.cohere.com/reference/rerank 的本地部署版本

已测试的重排序模型:

其他重排序模型:

多模态与跨模态 - 图像和音频嵌入 允许进行图像<->文本或图像<->音频搜索的专用嵌入模型。 通常,这些模型支持文本<->文本、文本<->其他以及其他<->其他搜索,但在进行跨模态搜索时存在准确性权衡。

图像<->文本模型可用于例如照片库搜索,用户可输入关键词查找照片,或使用照片查找相关图像。 音频<->文本模型不太流行,例如可用于根据文本描述查找音乐歌曲或相关音乐歌曲。

已测试的图像<->文本模型:

已测试的音频<->文本模型:

不支持:

  • 纯视觉模型,例如 nomic-ai/nomic-embed-vision-v1.5
ColBert 风格的晚期交互 Embeddings ColBert Embeddings 不执行任何特殊的 Pooling 方法,而是返回原始的 **token embeddings**。 随后,**token embeddings** 将在 VectorDB (Qdrant / Vespa) 中使用 MaxSim Metric 进行评分。

若通过 RestAPI 使用,晚期交互 embeddings 最好通过 base64 编码进行传输。 示例 notebook: https://colab.research.google.com/drive/14FqLc0N_z92_VgL_zygWV5pJZkaskyk7?usp=sharing

已测试的 colbert 模型:

ColPali 风格的晚期交互图像<->文本嵌入 用法与 ColBert 类似,但扫描的是图像<->文本,而不仅仅是文本。

若通过 RestAPI 使用,晚期交互嵌入最好通过 base64 编码进行传输。 示例笔记本:https://colab.research.google.com/drive/14FqLc0N_z92_VgL_zygWV5pJZkaskyk7?usp=sharing

已测试的 ColPali/ColQwen 模型:

文本分类 一种 bert 风格的多标签文本分类。将其分类到不同的类别中。

已测试的模型:

通过 Python API 使用 Infinity

代替 cli 和 RestAPI,通过 Python API 使用 infinity 的接口。 这提供了最大的灵活性。Python API 基于 asyncio 及其 await/async 特性,以允许并发处理请求。CLI 的参数也可以通过 Python 使用。

嵌入

import asyncio
from infinity_emb import AsyncEngineArray, EngineArgs, AsyncEmbeddingEngine

sentences = ["Embed this is sentence via Infinity.", "Paris is in France."]
array = AsyncEngineArray.from_args([
  EngineArgs(model_name_or_path = "BAAI/bge-small-en-v1.5", engine="torch", embedding_dtype="float32", dtype="auto")
])

async def embed_text(engine: AsyncEmbeddingEngine): 
    async with engine: 
        embeddings, usage = await engine.embed(sentences=sentences)
    # or handle the async start / stop yourself.
    await engine.astart()
    embeddings, usage = await engine.embed(sentences=sentences)
    await engine.astop()
asyncio.run(embed_text(array[0]))

重排序

Reranking 为您提供查询与多个文档之间相似度的评分。 请将其与 VectorDB+Embeddings 结合使用,或针对少量文档作为独立组件使用。 请从 huggingface 中选择一个兼容 AutoModelForSequenceClassification 的单分类模型。

import asyncio
from infinity_emb import AsyncEngineArray, EngineArgs, AsyncEmbeddingEngine
query = "What is the python package infinity_emb?"
docs = ["This is a document not related to the python package infinity_emb, hence...", 
    "Paris is in France!",
    "infinity_emb is a package for sentence embeddings and rerankings using transformer models in Python!"]
array = AsyncEmbeddingEngine.from_args(
  [EngineArgs(model_name_or_path = "mixedbread-ai/mxbai-rerank-xsmall-v1", engine="torch")]
)

async def rerank(engine: AsyncEmbeddingEngine): 
    async with engine:
        ranking, usage = await engine.rerank(query=query, docs=docs)
        print(list(zip(ranking, docs)))
    # or handle the async start / stop yourself.
    await engine.astart()
    ranking, usage = await engine.rerank(query=query, docs=docs)
    await engine.astop()

asyncio.run(rerank(array[0]))

使用 CLI 时,使用此命令启动 rerankers:

infinity_emb v2 --model-id mixedbread-ai/mxbai-rerank-xsmall-v1

Image-Embeddings: CLIP 模型

CLIP 模型能够同时编码图像和文本。

import asyncio
from infinity_emb import AsyncEngineArray, EngineArgs, AsyncEmbeddingEngine

sentences = ["This is awesome.", "I am bored."]
images = ["http://images.cocodataset.org/val2017/000000039769.jpg"]
engine_args = EngineArgs(
    model_name_or_path = "wkcn/TinyCLIP-ViT-8M-16-Text-3M-YFCC15M", 
    engine="torch"
)
array = AsyncEngineArray.from_args([engine_args])

async def embed(engine: AsyncEmbeddingEngine): 
    await engine.astart()
    embeddings, usage = await engine.embed(sentences=sentences)
    embeddings_image, _ = await engine.image_embed(images=images)
    await engine.astop()

asyncio.run(embed(array["wkcn/TinyCLIP-ViT-8M-16-Text-3M-YFCC15M"]))

音频嵌入:CLAP 模型

CLAP 模型能够同时编码音频和文本。

import asyncio
from infinity_emb import AsyncEngineArray, EngineArgs, AsyncEmbeddingEngine
import requests
import soundfile as sf
import io

sentences = ["This is awesome.", "I am bored."]

url = "https://bigsoundbank.com/UPLOAD/wav/2380.wav"
raw_bytes = requests.get(url, stream=True).content

audios = [raw_bytes]
engine_args = EngineArgs(
    model_name_or_path = "laion/clap-htsat-unfused",
    dtype="float32", 
    engine="torch"

)
array = AsyncEngineArray.from_args([engine_args])

async def embed(engine: AsyncEmbeddingEngine): 
    await engine.astart()
    embeddings, usage = await engine.embed(sentences=sentences)
    embedding_audios = await engine.audio_embed(audios=audios)
    await engine.astop()

asyncio.run(embed(array["laion/clap-htsat-unfused"]))

文本分类

使用 Infinity 的 classify 功能进行文本分类,该功能支持情感分析、情绪检测以及更多分类任务。

import asyncio
from infinity_emb import AsyncEngineArray, EngineArgs, AsyncEmbeddingEngine

sentences = ["This is awesome.", "I am bored."]
engine_args = EngineArgs(
    model_name_or_path = "SamLowe/roberta-base-go_emotions", 
    engine="torch", model_warmup=True)
array = AsyncEngineArray.from_args([engine_args])

async def classifier(engine: AsyncEmbeddingEngine): 
    async with engine:
        predictions, usage = await engine.classify(sentences=sentences)
    # or handle the async start / stop yourself.
    await engine.astart()
    predictions, usage = await engine.classify(sentences=sentences)
    await engine.astop()
asyncio.run(classifier(array["SamLowe/roberta-base-go_emotions"]))

通过 Python 客户端使用 Infinity

Infinity 提供了用于 RestAPI 客户端使用的生成客户端代码。

如果您想通过 RestAPI 调用远程 infinity 实例,请在本地安装以下软件包:

pip install infinity_client

有关更多信息,请查看 Client Readme https://github.com/michaelfeil/infinity/tree/main/libs/client_infinity/infinity_client

集成:

文档

https:///michaelfeil.github.io/infinity 查看文档,了解如何入门。 启动后,Swagger Ui 将在 {url}:{port}/docs 下可用,在此情况下为 http://localhost:7997/docs。您还可以在此处找到交互式预览:https://infinity.modal.michaelfeil.eu/docs(以及 https://michaelfeil-infinity.hf.space/docs)

贡献与开发

在 Ubuntu 22.04 上通过 Poetry 1.8.1 和 Python3.11 安装

cd libs/infinity_emb
poetry install --extras all --with lint,test

要通过 CI:

cd libs/infinity_emb
make precommit

所有贡献必须以与本仓库 MIT 许可证兼容的方式进行。

引用

@software{feil_2023_11630143,
  author       = {Feil, Michael},
  title        = {Infinity - To Embeddings and Beyond},
  month        = oct,
  year         = 2023,
  publisher    = {Zenodo},
  doi          = {10.5281/zenodo.11630143},
  url          = {https://doi.org/10.5281/zenodo.11630143}
}

💚 当前贡献者