infinity

Infinity 项目介绍

Infinity 是一个高吞吐量、低延迟的 REST API，用于提供文本嵌入、重排名模型、CLIP、CLAP 和 COLPALI。该项目在 MIT 许可证下开发。

为什么选择 Infinity

支持 HuggingFace 的所有模型部署: 支持从 HuggingFace 部署任何文本嵌入、重排名、CLIP 和句子转换器模型。
快速推理后端: 推理服务器基于 torch、optimum (ONNX/TensorRT) 和 CTranslate2 构建，使用 FlashAttention 充分利用硬件加速器，如 NVIDIA CUDA、AMD ROCM、CPU、AWS INF2 或 APPLE MPS。支持动态批处理和工作线程中专用的标记化。
多模态、多模型: 支持各种模型的混合搭配与协调。
经过测试的实现: 单元测试和端到端测试都已完成，让 API 用户可以在无穷的创意中创建嵌入。
易于使用: 基于 FastAPI 构建，Infinity CLI v2 允许通过环境变量或参数启动所有参数。OpenAPI 与 OpenAI 的 API 规格对齐。有关如何开始使用的文档请查看此处。

开始使用

通过 pip 安装启动 CLI

pip install infinity-emb[all]

安装完成后，你可以直接在激活的虚拟环境中运行 CLI。

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

使用预构建的 Docker 容器启动 CLI

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

使用 Python API 启动

import asyncio
from infinity_emb import AsyncEngineArray, EngineArgs, AsyncEmbeddingEngine

sentences = ["Embed this 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)
    await engine.astart()
    embeddings, usage = await engine.embed(sentences=sentences)
    await engine.astop()

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

重排序功能

重排名为查询和多个文档间的相似性提供得分。

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)))
    await engine.astart()
    ranking, usage = await engine.rerank(query=query, docs=docs)
    await engine.astop()

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

文本分类

使用 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(): 
    async with engine:
        predictions, usage = await engine.classify(sentences=sentences)
    await engine.astart()
    predictions, usage = await engine.classify(sentences=sentences)
    await engine.astop()
asyncio.run(classifier(array["SamLowe/roberta-base-go_emotions"]))