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

Infinity 提供一个强大的、灵活的嵌入和重排序平台，适用于各种机器学习需求。如果您正在寻找一个开源的解决方案来处理文本和多模态数据的嵌入，这可能正是您需要的工具。