池化模型
vLLM 还支持池化模型,包括嵌入、重排序和奖励模型。
在 vLLM 中,池化模型实现了 VllmModelForPooling 接口。这些模型使用一个 [Pooler][vllm.model_executor.layers.Pooler] 在返回之前提取输入的最终隐藏状态。
注意
目前我们主要为了方便而支持池化模型。如兼容性矩阵所示,大多数 vLLM 功能不适用于池化模型,因为它们只在生成或解码阶段工作,因此性能提升可能不大。
对于池化模型,我们支持以下 --task
选项。选定的选项设置了用于提取最终隐藏状态的默认池化器
任务 | 池化类型 | 归一化 | Softmax |
---|---|---|---|
嵌入 (embed ) |
LAST |
✅︎ | ❌ |
分类 (classify ) |
LAST |
❌ | ✅︎ |
句子对评分 (score ) |
* | * | * |
*默认池化器始终由模型定义。
注意
如果模型在 vLLM 中的实现定义了其自己的池化器,则默认池化器将设置为该池化器,而不是此表中指定的池化器。
加载 Sentence Transformers 模型时,我们尝试根据其 Sentence Transformers 配置文件 (modules.json
) 覆盖默认池化器。
提示
您可以通过 --override-pooler-config
选项自定义模型的池化方法,该选项优先于模型和 Sentence Transformers 的默认设置。
离线推理¶
LLM 类为离线推理提供了各种方法。有关初始化模型时的选项列表,请参阅配置。
LLM.encode
¶
encode 方法适用于 vLLM 中的所有池化模型。它直接返回提取的隐藏状态,这对于奖励模型很有用。
from vllm import LLM
llm = LLM(model="Qwen/Qwen2.5-Math-RM-72B", task="reward")
(output,) = llm.encode("Hello, my name is")
data = output.outputs.data
print(f"Data: {data!r}")
LLM.embed
¶
embed 方法为每个 prompt 输出一个嵌入向量。它主要为嵌入模型设计。
from vllm import LLM
llm = LLM(model="intfloat/e5-mistral-7b-instruct", task="embed")
(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 方法为每个 prompt 输出一个概率向量。它主要为分类模型设计。
from vllm import LLM
llm = LLM(model="jason9693/Qwen2.5-1.5B-apeach", task="classify")
(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", task="score")
(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
在线服务¶
我们的兼容 OpenAI 的服务器提供了对应于离线 API 的端点
- Pooling API 类似于
LLM.encode
,适用于所有类型的池化模型。 - Embeddings API 类似于
LLM.embed
,接受文本和多模态输入用于嵌入模型。 - Classification API 类似于
LLM.classify
,适用于序列分类模型。 - Score API 类似于
LLM.score
,用于交叉编码器模型。
Matryoshka 嵌入¶
Matryoshka 嵌入或Matryoshka 表示学习 (MRL) 是一种用于训练嵌入模型的技术。它允许用户在性能和成本之间进行权衡。
警告
并非所有嵌入模型都使用 Matryoshka 表示学习进行训练。为了避免滥用 dimensions
参数,vLLM 对于尝试更改不支持 Matryoshka 嵌入的模型的输出维度的请求会返回错误。
例如,在使用 BAAI/bge-m3
模型时设置 dimensions
参数将导致以下错误。
手动启用 Matryoshka 嵌入¶
目前没有官方接口用于指定对 Matryoshka 嵌入的支持。在 vLLM 中,如果 config.json
中的 is_matryoshka
为 True
,则允许将输出更改为任意维度。使用 matryoshka_dimensions
可以控制允许的输出维度。
对于支持 Matryoshka 嵌入但未被 vLLM 识别的模型,请使用 hf_overrides={"is_matryoshka": True}
, hf_overrides={"matryoshka_dimensions": [<允许的输出维度>]}
(离线)或 --hf_overrides '{"is_matryoshka": true}'
, --hf_overrides '{"matryoshka_dimensions": [<允许的输出维度>]}'
(在线)手动覆盖配置。
这是一个启用 Matryoshka 嵌入来服务模型的示例。
离线推理¶
您可以通过使用 PoolingParams 中的 dimensions 参数更改支持 Matryoshka 嵌入的嵌入模型的输出维度。
from vllm import LLM, PoolingParams
model = LLM(model="jinaai/jina-embeddings-v3",
task="embed",
trust_remote_code=True)
outputs = model.embed(["Follow the white rabbit."],
pooling_params=PoolingParams(dimensions=32))
print(outputs[0].outputs)
代码示例可在以下位置找到: examples/offline_inference/embed_matryoshka_fy.py
在线推理¶
使用以下命令启动 vllm 服务器。
您可以通过使用 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/online_serving/openai_embedding_matryoshka_fy.py