跳到內容

池化模型

vLLM 還支援池化模型,包括嵌入模型、重排序模型和獎勵模型。

在 vLLM 中,池化模型實現了 VllmModelForPooling 介面。這些模型使用 [Pooler][vllm.model_executor.layers.Pooler] 來提取輸入最終的隱藏狀態,然後將其返回。

注意

我們目前支援池化模型主要是為了方便。如相容性矩陣所示,大多數 vLLM 功能不適用於池化模型,因為它們只在生成或解碼階段工作,因此效能提升可能不會很大。

如果模型沒有實現此介面,您可以設定 --task 來告訴 vLLM 將模型轉換為池化模型。

--task 模型型別 支援的池化任務
embed 嵌入模型 encode, embed
classify 分類模型 encode, classify, score
reward 獎勵模型 encode

池化任務

在 vLLM 中,我們定義了以下池化任務和相應的 API

任務 API
encode encode
embed embed, score*
classify classify
score score

*如果模型不支援 score 任務,score API 將回退到 embed 任務。

vLLM 中的每個池化模型都根據 [Pooler.get_supported_tasks][vllm.model_executor.layers.Pooler.get_supported_tasks] 支援一個或多個這些任務。

預設情況下,分配給每個任務的池化器具有以下屬性

任務 池化型別 歸一化 Softmax
encode ALL
embed LAST ✅︎
classify LAST ✅︎

這些預設值可能會被 vLLM 中模型的實現所覆蓋。

載入 Sentence Transformers 模型時,我們嘗試根據其 Sentence Transformers 配置檔案 (modules.json) 覆蓋預設值,該檔案優先於模型的預設值。

您可以透過 --override-pooler-config 選項進一步自定義此設定,該選項優先於模型和 Sentence Transformers 的預設值。

注意

如果 vLLM 中模型的實現定義了不基於 PoolerConfig 的自有池化器,則上述配置可能會被忽略。

離線推理

LLM 類提供了多種用於離線推理的方法。有關初始化模型時的選項列表,請參閱配置

LLM.encode

vLLM 中所有的池化模型都可以使用 encode 方法。它直接返回提取的隱藏狀態,這對於獎勵模型很有用。

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 方法為每個提示輸出一個嵌入向量。它主要為嵌入模型設計。

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 方法為每個提示輸出一個機率向量。它主要為分類模型設計。

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 對應的端點

  • 池化 APILLM.encode 類似,適用於所有型別的池化模型。
  • 嵌入 APILLM.embed 類似,接受文字和多模態輸入用於嵌入模型。
  • 分類 APILLM.classify 類似,適用於序列分類模型。
  • 分數 APILLM.score 類似,適用於交叉編碼器模型。

Matryoshka 嵌入

Matryoshka 嵌入Matryoshka 表示學習 (MRL) 是一種用於訓練嵌入模型的技術。它允許使用者在效能和成本之間進行權衡。

警告

並非所有嵌入模型都使用 Matryoshka 表示學習進行訓練。為避免濫用 dimensions 引數,對於嘗試更改不支援 Matryoshka 嵌入的模型輸出維度的請求,vLLM 會返回錯誤。

例如,在使用 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": [<allowed output dimensions>]}(離線)或 --hf_overrides '{"is_matryoshka": true}', --hf_overrides '{"matryoshka_dimensions": [<allowed output 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",
          task="embed",
          trust_remote_code=True)
outputs = llm.embed(["Follow the white rabbit."],
                    pooling_params=PoolingParams(dimensions=32))
print(outputs[0].outputs)

程式碼示例可在此處找到: examples/offline_inference/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/online_serving/openai_embedding_matryoshka_fy.py