Infinity: 高性能文本嵌入和重排序API

infinity

Infinity: 开源高性能文本嵌入和重排序API

Infinity是一个开源的高性能文本嵌入和重排序API,旨在为开发者提供快速、灵活的自然语言处理能力。它具有以下主要特点和优势:

强大的模型支持

Infinity支持部署任何来自HuggingFace的嵌入、重排序、CLIP和句子转换器模型。这意味着开发者可以轻松使用最新的NLP模型,无需自己训练或微调。支持的模型类型包括:

文本嵌入模型
重排序模型
CLIP多模态模型
文本分类模型

这种广泛的模型支持使Infinity能够应用于各种NLP任务,如语义搜索、文档排序、图像-文本匹配等。

高性能推理后端

Infinity采用了多种技术来实现高性能推理:

基于PyTorch、ONNX/TensorRT和CTranslate2构建
使用FlashAttention加速
支持NVIDIA CUDA、AMD ROCm、CPU、AWS INF2和Apple MPS等多种硬件加速器
动态批处理和专用线程中的分词

这些优化使Infinity能够实现高吞吐量和低延迟,满足生产环境的需求。

多模态和多模型支持

Infinity允许混合使用多个不同的模型,并能够协调它们的工作。这为开发复杂的NLP应用提供了极大的灵活性。例如,可以同时使用文本嵌入模型和图像-文本匹配模型来构建多模态搜索系统。

易用性

Infinity提供了多种使用方式:

命令行界面(CLI):可以通过pip安装后直接使用CLI启动服务。
Docker容器:提供了预构建的Docker镜像,方便部署。
Python API:可以在Python代码中直接调用Infinity的功能。

此外,Infinity还提供了符合OpenAI API规范的REST API,便于集成到现有系统中。

广泛的应用

Infinity已经在多个知名公司和项目中得到应用,包括:

RunPod
TrueFoundry
Vast.ai
SAP
Freshworks
等

这些案例展示了Infinity在实际生产环境中的可靠性和性能。

Infinity应用案例

快速开始

要开始使用Infinity,有几种方法:

1. 通过pip安装CLI

首先安装Infinity:

pip install infinity-emb[all]

然后可以直接运行CLI:

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

使用v2 --help命令可以查看所有参数的说明。

2. 使用预构建的Docker容器(推荐)

使用Docker运行Infinity更加方便,特别是在需要GPU加速的情况下:

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

3. 通过Python API使用

Infinity还提供了Python API,可以直接在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]))

主要功能

1. 文本嵌入

Infinity支持各种流行的文本嵌入模型,如:

mixedbread-ai/mxbai-embed-large-v1
WhereIsAI/UAE-Large-V1
BAAI/bge-base-en-v1.5
Alibaba-NLP/gte-large-en-v1.5
jinaai/jina-embeddings-v2-base-code
intfloat/multilingual-e5-large-instruct

这些模型可以将文本转换为高维向量,用于语义搜索、文本聚类等任务。

2. 重排序

重排序功能可以为查询和多个文档之间的相似度打分,常用于搜索结果的优化:

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)))

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

支持的重排序模型包括:

mixedbread-ai/mxbai-rerank-xsmall-v1
BAAI/bge-reranker-base
jinaai/jina-reranker-v1-turbo-en

3. 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"]))

支持的CLIP模型包括:

wkcn/TinyCLIP-ViT-8M-16-Text-3M-YFCC15M
jinaai/jina-clip-v1 (需要额外安装timm库)

4. 文本分类

Infinity还支持文本分类任务,如情感分析、情绪检测等:

import asyncio
from infinity_emb import AsyncEngineArray, EngineArgs, AsyncEmbeddingEngine

sentences = ["This is awesome.", "I am bored."]
engine_args = EngineArgs(
    model_name_or_path = "distilbert-base-uncased-finetuned-sst-2-english", 
    engine="torch"
)
array = AsyncEngineArray.from_args([engine_args])

async def classify(engine: AsyncEmbeddingEngine): 
    await engine.astart()
    classifications, usage = await engine.classify(sentences=sentences)
    await engine.astop()

asyncio.run(classify(array["distilbert-base-uncased-finetuned-sst-2-english"]))