@tduncan / Nvidia-Optimization-Stack-VLLM
Code Commits Issues Pulls Cloud Agent CI/CD
main
main
md 140 lines 6.11 KB
Rendered Raw
1 # Token Embedding Usages
2
3 ## Summary
4
5 - Model Usage: Token classification models
6 - Pooling Tasks: `token_embed`
7 - Offline APIs:
8 - `LLM.encode(..., pooling_task="token_embed")`
9 - Online APIs:
10 - Pooling API (`/pooling`)
11
12 The difference between the (sequence) embedding task and the token embedding task is that (sequence) embedding outputs one embedding for each sequence, while token embedding outputs an embedding for each token.
13
14 Many embedding models support both (sequence) embedding and token embedding. For further details on (sequence) embedding, please refer to [this page](embed.md).
15
16 !!! note
17
18 Pooling multitask support has been removed since v0.21. When the default pooling task (embed) is not
19 what you want, you need to manually specify it via `PoolerConfig(task="token_embed")` offline or
20 `--pooler-config.task token_embed` online.
21
22 ## Typical Use Cases
23
24 ### Multi-Vector Retrieval
25
26 For implementation examples, see:
27
28 Offline: [examples/pooling/token_embed/multi_vector_retrieval_offline.py](../../../examples/pooling/token_embed/multi_vector_retrieval_offline.py)
29
30 Online: [examples/pooling/token_embed/multi_vector_retrieval_online.py](../../../examples/pooling/token_embed/multi_vector_retrieval_online.py)
31
32 ### Late interaction
33
34 Similarity scores can be computed using late interaction between two input prompts via the score API. For more information, see [Score API](scoring.md).
35
36 ### Extract last hidden states
37
38 Models of any architecture can be converted into embedding models using `--convert embed`. Token embedding can then be used to extract the last hidden states from these models.
39
40 ## Supported Models
41
42 --8<-- [start:supported-token-embed-models]
43
44 ### Text-only Models
45
46 | Architecture | Models | Example HF Models | [LoRA](../../features/lora.md) | [PP](../../serving/parallelism_scaling.md) |
47 | ------------ | ------ | ----------------- | -------------------- | ------------------------- |
48 | `ColBERTLfm2Model` | LFM2 | `LiquidAI/LFM2-ColBERT-350M` | | |
49 | `ColBERTModernBertModel` | ModernBERT | `lightonai/GTE-ModernColBERT-v1` | | |
50 | `ColBERTJinaRobertaModel` | Jina XLM-RoBERTa | `jinaai/jina-colbert-v2` | | |
51 | `HF_ColBERT` | BERT | `answerdotai/answerai-colbert-small-v1`, `colbert-ir/colbertv2.0` | | |
52 | `*Model`<sup>C</sup>, `*ForCausalLM`<sup>C</sup>, etc. | Generative models | N/A | \* | \* |
53
54 ### Multimodal Models
55
56 !!! note
57 For more information about multimodal models inputs, see [this page](../supported_models.md#list-of-multimodal-language-models).
58
59 | Architecture | Models | Inputs | Example HF Models | [LoRA](../../features/lora.md) | [PP](../../serving/parallelism_scaling.md) |
60 | ------------ | ------ | ----- | ----------------- | ------------------------------ | ------------------------------------------ |
61 | `ColModernVBertForRetrieval` | ColModernVBERT | T / I | `ModernVBERT/colmodernvbert-merged` | | |
62 | `ColPaliForRetrieval` | ColPali | T / I | `vidore/colpali-v1.3-hf` | | |
63 | `ColQwen3` | Qwen3-VL | T / I | `TomoroAI/tomoro-colqwen3-embed-4b`, `TomoroAI/tomoro-colqwen3-embed-8b` | | |
64 | `ColQwen3_5` | ColQwen3.5 | T + I + V | `athrael-soju/colqwen3.5-4.5B-v3`, `vultr/VultronRetrieverPrime-Qwen3.5-8B` | | |
65 | `OpsColQwen3Model` | Qwen3-VL | T / I | `OpenSearch-AI/Ops-Colqwen3-4B`, `OpenSearch-AI/Ops-Colqwen3-8B` | | |
66 | `Qwen3VLNemotronEmbedModel` | Qwen3-VL | T / I | `nvidia/nemotron-colembed-vl-4b-v2`, `nvidia/nemotron-colembed-vl-8b-v2` | ✅︎ | ✅︎ |
67 | `*ForConditionalGeneration`<sup>C</sup>, `*ForCausalLM`<sup>C</sup>, etc. | Generative models | \* | N/A | \* | \* |
68
69 <sup>C</sup> Automatically converted into an embedding model via `--convert embed`. ([details](./README.md#model-conversion))
70 \* Feature support is the same as that of the original model.
71
72 If your model is not in the above list, we will try to automatically convert the model using [as_embedding_model][vllm.model_executor.models.adapters.as_embedding_model].
73
74 ### Special models
75
76 | Architecture | Models | Example HF Models | [LoRA](../../features/lora.md) | [PP](../../serving/parallelism_scaling.md) |
77 | ------------ | ------ | ----------------- | -------------------- | ------------------------- |
78 | `JinaForRanking` | Qwen3-based | `jinaai/jina-reranker-v3` | | |
79
80 jina-reranker-v3 is a listwise document reranker model with a novel `last but not late interaction` architecture. More information can be found at: [examples/pooling/token_embed/jina_reranker_v3_offline.py](../../../examples/pooling/token_embed/jina_reranker_v3_offline.py)
81
82 --8<-- [end:supported-token-embed-models]
83
84 ## Offline Inference
85
86 ### Pooling Parameters
87
88 The following [pooling parameters][vllm.PoolingParams] are supported.
89
90 ```python
91 --8<-- "vllm/pooling_params.py:common-pooling-params"
92 --8<-- "vllm/pooling_params.py:embed-pooling-params"
93 ```
94
95 ### `LLM.encode`
96
97 The [encode][vllm.entrypoints.pooling.offline.PoolingOfflineMixin.encode] method is available to all pooling models in vLLM.
98
99 Set `pooling_task="token_embed"` when using `LLM.encode` for token embedding Models:
100
101 ```python
102 from vllm import LLM
103
104 llm = LLM(model="answerdotai/answerai-colbert-small-v1", runner="pooling")
105 (output,) = llm.encode("Hello, my name is", pooling_task="token_embed")
106
107 data = output.outputs.data
108 print(f"Data: {data!r}")
109 ```
110
111 ### `LLM.score`
112
113 The [score][vllm.entrypoints.pooling.offline.PoolingOfflineMixin.score] method outputs similarity scores between sentence pairs.
114
115 All models that support token embedding task also support using the score API to compute similarity scores by calculating the late interaction of two input prompts.
116
117 ```python
118 from vllm import LLM
119
120 llm = LLM(model="answerdotai/answerai-colbert-small-v1", runner="pooling")
121 (output,) = llm.score(
122 "What is the capital of France?",
123 "The capital of Brazil is Brasilia.",
124 )
125
126 score = output.outputs.score
127 print(f"Score: {score}")
128 ```
129
130 ## Online Serving
131
132 Please refer to the [Pooling API](README.md#pooling-api) and use `"task":"token_embed"`.
133
134 ## More examples
135
136 More examples can be found here: [examples/pooling/token_embed](../../../examples/pooling/token_embed)
137
138 ## Supported Features
139
140 Token embedding features should be consistent with (sequence) embedding. For more information, see [this page](embed.md#supported-features).