跳到内容

池化模型

vLLM 也支持池化模型,例如嵌入、分类和奖励模型。

在 vLLM 中,池化模型实现了 VllmModelForPooling 接口。这些模型使用 Pooler 在返回输入之前提取最终的隐藏状态。

注意

我们目前主要为了方便支持池化模型。这不保证会比直接使用 Hugging Face Transformers 或 Sentence Transformers 提供任何性能改进。

我们计划在 vLLM 中优化池化模型。请在 Issue #21796 提出您的建议!

配置

模型运行器

通过选项 --runner pooling 以池化模式运行模型。

提示

在绝大多数情况下,无需设置此选项,因为 vLLM 可以通过 --runner auto 自动检测合适的模型运行器。

模型转换

vLLM 可以通过选项 --convert <type> 来调整模型以适应各种池化任务。

如果已设置(手动或自动)--runner pooling 但模型未实现 VllmModelForPooling 接口,vLLM 将尝试根据下表中显示的架构名称自动转换模型。

架构 --convert 支持的池化任务
*ForTextEncoding, *EmbeddingModel, *Model embed token_embed, embed
*ForRewardModeling, *RewardModel embed token_embed, embed
*For*Classification, *ClassificationModel classify token_classify, classify, score

提示

您可以明确设置 --convert <type> 来指定如何转换模型。

池化任务

vLLM 中的每个池化模型都支持一个或多个这些任务,具体取决于 Pooler.get_supported_tasks,从而启用相应的 API。

任务 API
embed LLM.embed(...), LLM.score(...)*, LLM.encode(..., pooling_task="embed")
classify LLM.classify(...), LLM.encode(..., pooling_task="classify")
score LLM.score(...)
token_classify LLM.reward(...), LLM.encode(..., pooling_task="token_classify")
token_embed LLM.encode(..., pooling_task="token_embed")
plugin LLM.encode(..., pooling_task="plugin")

* 如果模型不支持 score 任务,LLM.score(...) API 将回退到 embed 任务。

池化器配置

预定义模型

如果模型定义的 Pooler 接受 pooler_config,则可以通过 --pooler-config 选项覆盖其某些属性。

已转换模型

如果模型已通过 --convert(如上所示)进行转换,则为每个任务分配的池化器具有以下默认属性:

任务 池化类型 标准化 Softmax
embed LAST ✅︎
classify LAST ✅︎

加载 Sentence Transformers 模型时,其 Sentence Transformers 配置文件(modules.json)的优先级高于模型的默认设置。

您可以通过 --pooler-config 选项进一步自定义此设置,该选项的优先级高于模型和 Sentence Transformers 的默认设置。

离线推理

LLM 类提供各种用于离线推理的方法。有关初始化模型时的选项列表,请参阅 配置

LLM.embed

embed 方法为每个提示输出一个嵌入向量。它主要设计用于嵌入模型。

from vllm import LLM

llm = LLM(model="intfloat/e5-small", runner="pooling")
(output,) = llm.embed("Hello, my name is")

embeds = output.outputs.embedding
print(f"Embeddings: {embeds!r} (size={len(embeds)})")

可以在此处找到代码示例: examples/offline_inference/basic/embed.py

LLM.classify

classify 方法为每个提示输出一个概率向量。它主要设计用于分类模型。

from vllm import LLM

llm = LLM(model="jason9693/Qwen2.5-1.5B-apeach", runner="pooling")
(output,) = llm.classify("Hello, my name is")

probs = output.outputs.probs
print(f"Class Probabilities: {probs!r} (size={len(probs)})")

可以在此处找到代码示例: examples/offline_inference/basic/classify.py

LLM.score

score 方法输出句子对之间的相似度得分。它专为嵌入模型和交叉编码器模型设计。嵌入模型使用余弦相似度,而 交叉编码器模型 在 RAG 系统中用作候选查询-文档对之间的重排序器。

注意

vLLM 只能执行 RAG 的模型推理组件(例如,嵌入、重排序)。要处理更高级别的 RAG,您应该使用 LangChain 等集成框架。

from vllm import LLM

llm = LLM(model="BAAI/bge-reranker-v2-m3", runner="pooling")
(output,) = llm.score(
    "What is the capital of France?",
    "The capital of Brazil is Brasilia.",
)

score = output.outputs.score
print(f"Score: {score}")

可以在此处找到代码示例: examples/offline_inference/basic/score.py

LLM.reward

reward 方法可用于 vLLM 中的所有奖励模型。

from vllm import LLM

llm = LLM(model="internlm/internlm2-1_8b-reward", runner="pooling", trust_remote_code=True)
(output,) = llm.reward("Hello, my name is")

data = output.outputs.data
print(f"Data: {data!r}")

可以在此处找到代码示例: examples/offline_inference/basic/reward.py

LLM.encode

encode 方法可用于 vLLM 中的所有池化模型。

注意

请使用更具体的某个方法,或在使用 LLM.encode 时直接设置任务。

  • 对于嵌入,请使用 LLM.embed(...)pooling_task="embed"
  • 对于分类 logits,请使用 LLM.classify(...)pooling_task="classify"
  • 对于相似度得分,请使用 LLM.score(...)
  • 对于奖励,请使用 LLM.reward(...)pooling_task="token_classify"
  • 对于 token 分类,请使用 pooling_task="token_classify"
  • 对于多向量检索,请使用 pooling_task="token_embed"
  • 对于 IO 处理器插件,请使用 pooling_task="plugin"
from vllm import LLM

llm = LLM(model="intfloat/e5-small", runner="pooling")
(output,) = llm.encode("Hello, my name is", pooling_task="embed")

data = output.outputs.data
print(f"Data: {data!r}")

在线服务

我们的 OpenAI 兼容服务器 提供了与离线 API 相对应的端点。

  • 嵌入 API 类似于 LLM.embed,它接受文本和 多模态输入 用于嵌入模型。
  • 分类 API 类似于 LLM.classify,适用于序列分类模型。
  • 得分 API 类似于用于交叉编码器模型的 LLM.score
  • 池化 API 类似于 LLM.encode,适用于所有类型的池化模型。

注意

在使用 池化 API 时,请使用更具体的端点之一或直接设置任务。

  • 对于嵌入,请使用 嵌入 API"task":"embed"
  • 对于分类 logits,请使用 分类 API"task":"classify"
  • 对于相似度得分,请使用 得分 API
  • 对于奖励,请使用 "task":"token_classify"
  • 对于 token 分类,请使用 "task":"token_classify"
  • 对于多向量检索,请使用 "task":"token_embed"
  • 对于 IO 处理器插件,请使用 "task":"plugin"
# start a supported embeddings model server with `vllm serve`, e.g.
# vllm serve intfloat/e5-small
import requests

host = "localhost"
port = "8000"
model_name = "intfloat/e5-small"

api_url = f"http://{host}:{port}/pooling"

prompts = [
    "Hello, my name is",
    "The president of the United States is",
    "The capital of France is",
    "The future of AI is",
]
prompt = {"model": model_name, "input": prompts, "task": "embed"}

response = requests.post(api_url, json=prompt)

for output in response.json()["data"]:
    data = output["data"]
    print(f"Data: {data!r} (size={len(data)})")

Matryoshka 嵌入

Matryoshka 嵌入Matryoshka 表示学习 (MRL) 是一种用于训练嵌入模型的技术。它允许用户在性能和成本之间进行权衡。

警告

并非所有嵌入模型都使用 Matryoshka 表示学习进行训练。为避免滥用 dimensions 参数,vLLM 会对试图更改不支持 Matryoshka 嵌入的模型输出维度的请求返回错误。

例如,在使用 BAAI/bge-m3 模型时设置 dimensions 参数将导致以下错误。

{"object":"error","message":"Model \"BAAI/bge-m3\" does not support matryoshka representation, changing output dimensions will lead to poor results.","type":"BadRequestError","param":null,"code":400}

手动启用 Matryoshka 嵌入

目前没有指定 Matryoshka 嵌入支持的官方接口。在 vLLM 中,如果 config.json 中的 is_matryoshkaTrue,则可以将输出维度更改为任意值。使用 matryoshka_dimensions 来控制允许的输出维度。

对于支持 Matryoshka 嵌入但未被 vLLM 识别的模型,可以通过 hf_overrides={"is_matryoshka": True}hf_overrides={"matryoshka_dimensions": [<允许的输出维度>]}(离线),或 --hf-overrides '{"is_matryoshka": true}'--hf-overrides '{"matryoshka_dimensions": [<允许的输出维度>]}'(在线)手动覆盖配置。

以下是启用 Matryoshka 嵌入的模型服务的示例。

vllm serve Snowflake/snowflake-arctic-embed-m-v1.5 --hf-overrides '{"matryoshka_dimensions":[256]}'

离线推理

您可以通过在 PoolingParams 中使用 dimensions 参数来更改支持 Matryoshka 嵌入的嵌入模型的输出维度。

from vllm import LLM, PoolingParams

llm = LLM(
    model="jinaai/jina-embeddings-v3",
    runner="pooling",
    trust_remote_code=True,
)
outputs = llm.embed(
    ["Follow the white rabbit."],
    pooling_params=PoolingParams(dimensions=32),
)
print(outputs[0].outputs)

可以在此处找到代码示例: examples/pooling/embed/embed_matryoshka_fy.py

在线推理

使用以下命令启动 vLLM 服务器。

vllm serve jinaai/jina-embeddings-v3 --trust-remote-code

您可以通过 dimensions 参数来更改支持 Matryoshka 嵌入的嵌入模型的输出维度。

curl http://127.0.0.1:8000/v1/embeddings \
  -H 'accept: application/json' \
  -H 'Content-Type: application/json' \
  -d '{
    "input": "Follow the white rabbit.",
    "model": "jinaai/jina-embeddings-v3",
    "encoding_format": "float",
    "dimensions": 32
  }'

预期输出

{"id":"embd-5c21fc9a5c9d4384a1b021daccaf9f64","object":"list","created":1745476417,"model":"jinaai/jina-embeddings-v3","data":[{"index":0,"object":"embedding","embedding":[-0.3828125,-0.1357421875,0.03759765625,0.125,0.21875,0.09521484375,-0.003662109375,0.1591796875,-0.130859375,-0.0869140625,-0.1982421875,0.1689453125,-0.220703125,0.1728515625,-0.2275390625,-0.0712890625,-0.162109375,-0.283203125,-0.055419921875,-0.0693359375,0.031982421875,-0.04052734375,-0.2734375,0.1826171875,-0.091796875,0.220703125,0.37890625,-0.0888671875,-0.12890625,-0.021484375,-0.0091552734375,0.23046875]}],"usage":{"prompt_tokens":8,"total_tokens":8,"completion_tokens":0,"prompt_tokens_details":null}}

可以在此处找到 OpenAI 客户端示例: examples/pooling/embed/openai_embedding_matryoshka_fy.py

已弃用功能

Encode 任务

我们将 encode 任务拆分为两个更具体的 token 级任务:token_embedtoken_classify

  • token_embedembed 相同,使用标准化作为激活函数。
  • token_classifyclassify 相同,默认使用 softmax 作为激活函数。

从 PoolingParams 中移除 softmax

我们将在 v0.15 中移除 PoolingParams 中的 softmaxactivation。改为使用 use_activation,因为我们允许 classifytoken_classify 使用任何激活函数。

as_reward_model

警告

我们将在 v0.15 中移除 --convert reward,而是使用 --convert embed

池化模型现在默认支持所有池化,您无需任何设置即可使用。

  • 提取隐藏状态优先使用 token_embed 任务。
  • 奖励模型优先使用 token_classify 任务。