| 1 | # SPDX-License-Identifier: Apache-2.0 |
| 2 | # SPDX-FileCopyrightText: Copyright contributors to the vLLM project |
| 3 | from collections.abc import Iterable |
| 4 | from typing import cast |
| 5 | |
| 6 | import torch |
| 7 | |
| 8 | from vllm import PromptType, TextPrompt |
| 9 | from vllm.config import ModelConfig |
| 10 | from vllm.entrypoints.chat_utils import ( |
| 11 | BaseMultiModalItemTracker, |
| 12 | ChatCompletionContentPartParam, |
| 13 | ChatCompletionContentPartTextParam, |
| 14 | ConversationMessage, |
| 15 | MultiModalItemTracker, |
| 16 | _parse_chat_message_content_parts, |
| 17 | ) |
| 18 | from vllm.exceptions import VLLMValidationError |
| 19 | from vllm.inputs import MultiModalDataDict, MultiModalUUIDDict |
| 20 | |
| 21 | from .typing import ( |
| 22 | ScoreContentPartParam, |
| 23 | ScoreData, |
| 24 | ScoreInput, |
| 25 | ScoringData, |
| 26 | ) |
| 27 | |
| 28 | |
| 29 | def get_num_special_tokens_for_pair(tokenizer) -> int: |
| 30 | """Get number of special tokens added for a text pair encoding.""" |
| 31 | method = getattr(tokenizer, "num_special_tokens_to_add", None) |
| 32 | if method is not None: |
| 33 | try: |
| 34 | return method(pair=True) |
| 35 | except TypeError: |
| 36 | pass |
| 37 | # Fallback: compute by tokenizing empty strings |
| 38 | empty_encoding = tokenizer("", text_pair="", add_special_tokens=True) |
| 39 | return len(empty_encoding["input_ids"]) |
| 40 | |
| 41 | |
| 42 | def truncate_text_to_tokens( |
| 43 | text: str, |
| 44 | tokenizer, |
| 45 | max_tokens: int, |
| 46 | ) -> str: |
| 47 | """Truncate text to a maximum number of content tokens. |
| 48 | |
| 49 | Uses offset_mapping to slice the original text at the exact character |
| 50 | boundary, avoiding lossy encode→decode round-trips that can shift |
| 51 | the token count due to BPE merge boundary changes. |
| 52 | """ |
| 53 | if max_tokens <= 0: |
| 54 | return "" |
| 55 | |
| 56 | while text: |
| 57 | encoding = tokenizer( |
| 58 | text, |
| 59 | add_special_tokens=False, |
| 60 | return_offsets_mapping=True, |
| 61 | ) |
| 62 | offsets = encoding["offset_mapping"] |
| 63 | if len(encoding["input_ids"]) <= max_tokens: |
| 64 | break |
| 65 | |
| 66 | char_end = min(offsets[max_tokens - 1][1], offsets[max_tokens][0]) |
| 67 | if char_end >= len(text): |
| 68 | char_end = len(text) - 1 |
| 69 | text = text[:char_end] |
| 70 | |
| 71 | return text |
| 72 | |
| 73 | |
| 74 | def compute_maxsim_score(q_emb: torch.Tensor, d_emb: torch.Tensor) -> torch.Tensor: |
| 75 | """Compute ColBERT MaxSim score. |
| 76 | |
| 77 | Args: |
| 78 | q_emb: Query token embeddings [query_len, dim] |
| 79 | d_emb: Document token embeddings [doc_len, dim] |
| 80 | |
| 81 | Returns: |
| 82 | MaxSim score (sum over query tokens of max similarity to any doc token) |
| 83 | |
| 84 | """ |
| 85 | # compute in float32 for numerical stability |
| 86 | # [query_len, doc_len] |
| 87 | token_scores = torch.matmul(q_emb.float(), d_emb.float().T) |
| 88 | # Max over document tokens, sum over query tokens |
| 89 | return token_scores.amax(dim=-1).sum() |
| 90 | |
| 91 | |
| 92 | def _validate_mm_score_input( |
| 93 | data: list[ScoreInput], |
| 94 | is_multimodal_model: bool, |
| 95 | architecture: str, |
| 96 | ) -> list[ScoreData]: |
| 97 | out: list[ScoreData] = [] |
| 98 | for d in data: |
| 99 | if isinstance(d, str): |
| 100 | out.append(d) |
| 101 | else: |
| 102 | if not is_multimodal_model: |
| 103 | raise VLLMValidationError( |
| 104 | f"MultiModalParam is not supported for {architecture}" |
| 105 | ) |
| 106 | content = cast(list[ScoreContentPartParam], d.get("content", [])) |
| 107 | out.append(content) |
| 108 | return out |
| 109 | |
| 110 | |
| 111 | def _validate_score_input_lens( |
| 112 | data_1: list[ScoreData], |
| 113 | data_2: list[ScoreData], |
| 114 | ): |
| 115 | len_1 = len(data_1) |
| 116 | len_2 = len(data_2) |
| 117 | |
| 118 | if len_1 > 1 and len_1 != len_2: |
| 119 | raise VLLMValidationError("Input lengths must be either 1:1, 1:N or N:N") |
| 120 | if len_1 == 0: |
| 121 | raise VLLMValidationError("At least one text element must be given") |
| 122 | if len_2 == 0: |
| 123 | raise VLLMValidationError("At least one text_pair element must be given") |
| 124 | |
| 125 | |
| 126 | def validate_score_input( |
| 127 | data_1: ScoreInput | list[ScoreInput], |
| 128 | data_2: ScoreInput | list[ScoreInput], |
| 129 | is_multimodal_model: bool, |
| 130 | architecture: str, |
| 131 | ) -> ScoringData: |
| 132 | if not isinstance(data_1, list): |
| 133 | data_1 = [data_1] |
| 134 | |
| 135 | if not isinstance(data_2, list): |
| 136 | data_2 = [data_2] |
| 137 | |
| 138 | score_input_1 = _validate_mm_score_input(data_1, is_multimodal_model, architecture) |
| 139 | score_input_2 = _validate_mm_score_input(data_2, is_multimodal_model, architecture) |
| 140 | _validate_score_input_lens(score_input_1, score_input_2) |
| 141 | return ScoringData(data_1=score_input_1, data_2=score_input_2) |
| 142 | |
| 143 | |
| 144 | def score_data_to_prompts( |
| 145 | data_list: list[ScoreData], |
| 146 | role: str, |
| 147 | model_config: ModelConfig, |
| 148 | ) -> list[PromptType]: |
| 149 | """Convert a list of ScoreData into PromptType objects. |
| 150 | |
| 151 | For plain text inputs, returns the string directly. |
| 152 | For multimodal inputs (list of content parts), parses them into |
| 153 | a :class:`TextPrompt` with attached ``multi_modal_data`` / |
| 154 | ``multi_modal_uuids``. |
| 155 | |
| 156 | This is used by late-interaction scoring where each query/document |
| 157 | is encoded independently. |
| 158 | """ |
| 159 | prompts: list[PromptType] = [] |
| 160 | for data in data_list: |
| 161 | if isinstance(data, str): |
| 162 | prompts.append(data) |
| 163 | else: |
| 164 | text, mm_data, mm_uuids = parse_score_data_single(data, role, model_config) |
| 165 | prompt: TextPrompt = TextPrompt(prompt=text) |
| 166 | if mm_data is not None: |
| 167 | prompt["multi_modal_data"] = mm_data |
| 168 | if mm_uuids is not None: |
| 169 | prompt["multi_modal_uuids"] = mm_uuids |
| 170 | prompts.append(prompt) |
| 171 | return prompts |
| 172 | |
| 173 | |
| 174 | def _ensure_str(content: list[ConversationMessage]) -> str: |
| 175 | """Extract a single string prompt from parsed conversation content.""" |
| 176 | assert len(content) == 1 |
| 177 | prompt = content[0]["content"] |
| 178 | if prompt is not None and isinstance(prompt, str): |
| 179 | return cast(str, prompt) |
| 180 | raise ValueError(f"Only string content is supported, but got {content}.") |
| 181 | |
| 182 | |
| 183 | def _parse_score_content( |
| 184 | role: str, |
| 185 | data: ScoreData, |
| 186 | mm_tracker: BaseMultiModalItemTracker, |
| 187 | ) -> list[ConversationMessage]: |
| 188 | parts: Iterable[ChatCompletionContentPartParam] |
| 189 | if isinstance(data, str): |
| 190 | parts = [ChatCompletionContentPartTextParam(type="text", text=data)] |
| 191 | else: |
| 192 | parts = cast(Iterable[ChatCompletionContentPartParam], data) |
| 193 | |
| 194 | mm_parser = mm_tracker.create_parser() |
| 195 | |
| 196 | parse_res = _parse_chat_message_content_parts( |
| 197 | role=role, |
| 198 | parts=parts, |
| 199 | mm_tracker=mm_tracker, |
| 200 | wrap_dicts=False, |
| 201 | interleave_strings=False, |
| 202 | multimodal_content_part_separator="", |
| 203 | ) |
| 204 | |
| 205 | if parse_res: |
| 206 | return parse_res |
| 207 | |
| 208 | mm_placeholder_storage = mm_parser.mm_placeholder_storage() |
| 209 | |
| 210 | if ( |
| 211 | len(mm_placeholder_storage) != 1 |
| 212 | or len(next(iter(mm_placeholder_storage.values()))) != 1 |
| 213 | ): |
| 214 | raise ValueError("Only one multi-modal item is supported") |
| 215 | |
| 216 | return next(iter(mm_placeholder_storage.values()))[0] |
| 217 | |
| 218 | |
| 219 | def parse_score_data_single( |
| 220 | data: ScoreData, |
| 221 | role: str, |
| 222 | model_config: ModelConfig, |
| 223 | ) -> tuple[str, MultiModalDataDict | None, MultiModalUUIDDict | None]: |
| 224 | """Parse **one** ScoreData into a text prompt and its own multi-modal |
| 225 | data. |
| 226 | |
| 227 | Unlike :func:`parse_score_data`, each call creates an **independent** |
| 228 | :class:`MultiModalItemTracker` so multi-modal items are kept separate. |
| 229 | This is the correct behaviour for late-interaction scoring, where |
| 230 | query and document are encoded independently. |
| 231 | """ |
| 232 | mm_tracker = MultiModalItemTracker(model_config) |
| 233 | content = _parse_score_content(role, data, mm_tracker) |
| 234 | |
| 235 | prompt = _ensure_str(content) |
| 236 | mm_items, mm_uuids = mm_tracker.resolve_items() |
| 237 | return prompt, mm_items, mm_uuids |
| 238 | |
| 239 | |
| 240 | def parse_score_data( |
| 241 | data_1: ScoreData, |
| 242 | data_2: ScoreData, |
| 243 | model_config: ModelConfig, |
| 244 | ) -> tuple[str, str, MultiModalDataDict | None, MultiModalUUIDDict | None]: |
| 245 | """Parse a query-document pair into text prompts and shared multi-modal |
| 246 | data. |
| 247 | |
| 248 | Uses a **single** :class:`MultiModalItemTracker` so that multi-modal |
| 249 | items from both inputs are merged into one ``mm_data`` dict. This is |
| 250 | the correct behaviour for cross-encoder scoring, where query and |
| 251 | document are concatenated into a single model prompt. |
| 252 | """ |
| 253 | mm_tracker = MultiModalItemTracker(model_config) |
| 254 | |
| 255 | content_1 = _parse_score_content("query", data_1, mm_tracker) |
| 256 | content_2 = _parse_score_content("document", data_2, mm_tracker) |
| 257 | |
| 258 | prompt_1 = _ensure_str(content_1) |
| 259 | prompt_2 = _ensure_str(content_2) |
| 260 | mm_items, mm_uuids = mm_tracker.resolve_items() |
| 261 | |
| 262 | return prompt_1, prompt_2, mm_items, mm_uuids |
| 263 | |
| 264 | |
| 265 | def compress_token_type_ids(token_type_ids: list[int]) -> int: |
| 266 | """Return position of the first 1 or the length of the list |
| 267 | if not found. |
| 268 | """ |
| 269 | first_one = len(token_type_ids) |
| 270 | err_msg = ( |
| 271 | "Token type ids are expected to be a sequence" |
| 272 | " of zeros followed by a sequence of ones" |
| 273 | ) |
| 274 | for i, type_id in enumerate(token_type_ids): |
| 275 | if type_id == 0 and first_one < i: |
| 276 | raise ValueError(err_msg) |
| 277 | elif type_id == 1 and first_one > i: |
| 278 | first_one = i |
| 279 | elif type_id > 1: |
| 280 | raise ValueError(err_msg) |
| 281 | |
| 282 | return first_one |