Pooling 模型

Pooling 模型#

vLLM 也支持 pooling 模型，包括嵌入、重排序和奖励模型。

在 vLLM 中，pooling 模型实现了 VllmModelForPooling 接口。这些模型使用 Pooler 来提取输入的最终隐藏状态，然后再返回它们。

注意

我们目前主要为了方便起见而支持 pooling 模型。如 兼容性矩阵 所示，大多数 vLLM 功能不适用于 pooling 模型，因为它们仅在生成或解码阶段工作，因此性能可能不会有太大提升。

对于 pooling 模型，我们支持以下 --task 选项。选择的选项设置用于提取最终隐藏状态的默认 pooler

任务

Pooling 类型

归一化

Softmax

嵌入 (embed)

LAST

✅︎

❌

分类 (classify)

LAST

❌

✅︎

句子对评分 (score)

*

*

*

奖励建模 (reward)

ALL

❌

❌

*默认 pooler 始终由模型定义。

注意

如果模型在 vLLM 中的实现定义了自己的 pooler，则默认 pooler 将设置为该 pooler，而不是此表中指定的 pooler。

加载 Sentence Transformers 模型时，我们会尝试根据其 Sentence Transformers 配置文件 (modules.json) 覆盖默认 pooler。

提示

您可以通过 --override-pooler-config 选项自定义模型的 pooling 方法，该选项的优先级高于模型和 Sentence Transformers 的默认设置。

离线推理#

LLM 类提供了各种离线推理方法。有关初始化模型时的选项列表，请参阅 引擎参数。

LLM.encode#

encode 方法适用于 vLLM 中的所有 pooling 模型。它直接返回提取的隐藏状态，这对于奖励模型非常有用。

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 方法为每个提示输出一个嵌入向量。它主要为嵌入模型而设计。

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 方法为每个提示输出一个概率向量。它主要为分类模型而设计。

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。

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 相对应的端点