| 1 | # SPDX-License-Identifier: Apache-2.0 |
| 2 | # SPDX-FileCopyrightText: Copyright contributors to the vLLM project |
| 3 | |
| 4 | from io import BytesIO |
| 5 | |
| 6 | import pybase64 as base64 |
| 7 | import torch |
| 8 | import torch.nn.functional as F |
| 9 | from PIL import Image |
| 10 | from safetensors.torch import load_file |
| 11 | from transformers import AutoModel, AutoTokenizer |
| 12 | |
| 13 | from tests.conftest import HfRunner |
| 14 | from vllm.entrypoints.chat_utils import ( |
| 15 | ChatCompletionContentPartImageParam, |
| 16 | ChatCompletionContentPartTextParam, |
| 17 | ) |
| 18 | from vllm.entrypoints.pooling.scoring.typing import ScoreMultiModalParam |
| 19 | from vllm.entrypoints.pooling.scoring.utils import compute_maxsim_score |
| 20 | from vllm.transformers_utils.repo_utils import hf_api |
| 21 | |
| 22 | |
| 23 | class ColBERTScoringHfRunner(torch.nn.Module): |
| 24 | def __init__(self, model_name, linear_weights_key): |
| 25 | super().__init__() |
| 26 | |
| 27 | self.tokenizer = AutoTokenizer.from_pretrained(model_name) |
| 28 | |
| 29 | self.device = torch.device("cuda" if torch.cuda.is_available() else "cpu") |
| 30 | |
| 31 | extra = {} |
| 32 | if self.device.type == "cpu": |
| 33 | extra["attn_implementation"] = "eager" |
| 34 | |
| 35 | self.model = AutoModel.from_pretrained( |
| 36 | model_name, |
| 37 | **extra, |
| 38 | ).to(self.device) |
| 39 | self.model.eval() |
| 40 | |
| 41 | path = hf_api().hf_hub_download(model_name, filename="model.safetensors") |
| 42 | weights = load_file(path) |
| 43 | |
| 44 | self.linear_weight = weights[linear_weights_key].to(self.device).float() |
| 45 | |
| 46 | @torch.inference_mode() |
| 47 | def forward(self, texts): |
| 48 | embeddings = [] |
| 49 | for text in texts: |
| 50 | inputs = self.tokenizer(text, return_tensors="pt").to(self.device) |
| 51 | hidden = self.model(**inputs).last_hidden_state.float() |
| 52 | projected = F.linear(hidden, self.linear_weight.float()) |
| 53 | normalised = F.normalize(projected, p=2, dim=-1) |
| 54 | embeddings.append(normalised.squeeze(0).cpu()) |
| 55 | return embeddings |
| 56 | |
| 57 | @torch.inference_mode() |
| 58 | def predict(self, prompts: list[list[str]], *args, **kwargs): |
| 59 | hf_embeddings = [self(prompt) for prompt in prompts] |
| 60 | hf_outputs = [ |
| 61 | compute_maxsim_score(*map(torch.tensor, pair)).item() |
| 62 | for pair in hf_embeddings |
| 63 | ] |
| 64 | return torch.as_tensor(hf_outputs) |
| 65 | |
| 66 | |
| 67 | class EncoderScoringHfRunner(HfRunner): |
| 68 | def __init__(self, *args, **kwargs): |
| 69 | super().__init__(*args, **kwargs, is_sentence_transformer=True) |
| 70 | |
| 71 | @torch.inference_mode() |
| 72 | def predict(self, prompts: list[list[str]], *args, **kwargs): |
| 73 | hf_embeddings = [self.encode(prompt) for prompt in prompts] |
| 74 | hf_outputs = [ |
| 75 | F.cosine_similarity(*map(torch.tensor, pair), dim=0) |
| 76 | for pair in hf_embeddings |
| 77 | ] |
| 78 | return torch.as_tensor(hf_outputs) |
| 79 | |
| 80 | |
| 81 | def make_base64_image( |
| 82 | width: int = 64, height: int = 64, color: tuple[int, int, int] = (255, 0, 0) |
| 83 | ) -> str: |
| 84 | """Create a small solid-color PNG image and return its base64 data URI.""" |
| 85 | img = Image.new("RGB", (width, height), color) |
| 86 | buf = BytesIO() |
| 87 | img.save(buf, format="PNG") |
| 88 | b64 = base64.b64encode(buf.getvalue()).decode() |
| 89 | return f"data:image/png;base64,{b64}" |
| 90 | |
| 91 | |
| 92 | def make_image_mm_param( |
| 93 | image_uri: str, |
| 94 | text: str | None = None, |
| 95 | ) -> ScoreMultiModalParam: |
| 96 | """Build a ScoreMultiModalParam containing an image (and optional text).""" |
| 97 | content: list = [ |
| 98 | ChatCompletionContentPartImageParam( |
| 99 | type="image_url", |
| 100 | image_url={"url": image_uri}, |
| 101 | ), |
| 102 | ] |
| 103 | if text is not None: |
| 104 | content.append( |
| 105 | ChatCompletionContentPartTextParam(type="text", text=text), |
| 106 | ) |
| 107 | return ScoreMultiModalParam(content=content) |