| |
| |
|
|
| from io import BytesIO |
|
|
| import pybase64 as base64 |
| import torch |
| import torch.nn.functional as F |
| from huggingface_hub import hf_hub_download |
| from PIL import Image |
| from safetensors.torch import load_file |
| from transformers import AutoModel, AutoTokenizer |
|
|
| from tests.conftest import HfRunner |
| from vllm.entrypoints.chat_utils import ( |
| ChatCompletionContentPartImageParam, |
| ChatCompletionContentPartTextParam, |
| ) |
| from vllm.entrypoints.pooling.scoring.typing import ScoreMultiModalParam |
| from vllm.entrypoints.pooling.scoring.utils import compute_maxsim_score |
|
|
|
|
| class ColBERTScoringHfRunner(torch.nn.Module): |
| def __init__(self, model_name, linear_weights_key): |
| super().__init__() |
|
|
| self.tokenizer = AutoTokenizer.from_pretrained(model_name) |
|
|
| self.device = torch.device("cuda" if torch.cuda.is_available() else "cpu") |
|
|
| extra = {} |
| if self.device.type == "cpu": |
| extra["attn_implementation"] = "eager" |
|
|
| self.model = AutoModel.from_pretrained( |
| model_name, |
| **extra, |
| ).to(self.device) |
| self.model.eval() |
|
|
| path = hf_hub_download(model_name, filename="model.safetensors") |
| weights = load_file(path) |
|
|
| self.linear_weight = weights[linear_weights_key].to(self.device).float() |
|
|
| @torch.inference_mode() |
| def forward(self, texts): |
| embeddings = [] |
| for text in texts: |
| inputs = self.tokenizer(text, return_tensors="pt").to(self.device) |
| hidden = self.model(**inputs).last_hidden_state.float() |
| projected = F.linear(hidden, self.linear_weight.float()) |
| normalised = F.normalize(projected, p=2, dim=-1) |
| embeddings.append(normalised.squeeze(0).cpu()) |
| return embeddings |
|
|
| @torch.inference_mode() |
| def predict(self, prompts: list[list[str]], *args, **kwargs): |
| hf_embeddings = [self(prompt) for prompt in prompts] |
| hf_outputs = [ |
| compute_maxsim_score(*map(torch.tensor, pair)).item() |
| for pair in hf_embeddings |
| ] |
| return torch.as_tensor(hf_outputs) |
|
|
|
|
| class EncoderScoringHfRunner(HfRunner): |
| def __init__(self, *args, **kwargs): |
| super().__init__(*args, **kwargs, is_sentence_transformer=True) |
|
|
| @torch.inference_mode() |
| def predict(self, prompts: list[list[str]], *args, **kwargs): |
| hf_embeddings = [self.encode(prompt) for prompt in prompts] |
| hf_outputs = [ |
| F.cosine_similarity(*map(torch.tensor, pair), dim=0) |
| for pair in hf_embeddings |
| ] |
| return torch.as_tensor(hf_outputs) |
|
|
|
|
| def make_base64_image( |
| width: int = 64, height: int = 64, color: tuple[int, int, int] = (255, 0, 0) |
| ) -> str: |
| """Create a small solid-color PNG image and return its base64 data URI.""" |
| img = Image.new("RGB", (width, height), color) |
| buf = BytesIO() |
| img.save(buf, format="PNG") |
| b64 = base64.b64encode(buf.getvalue()).decode() |
| return f"data:image/png;base64,{b64}" |
|
|
|
|
| def make_image_mm_param( |
| image_uri: str, |
| text: str | None = None, |
| ) -> ScoreMultiModalParam: |
| """Build a ScoreMultiModalParam containing an image (and optional text).""" |
| content: list = [ |
| ChatCompletionContentPartImageParam( |
| type="image_url", |
| image_url={"url": image_uri}, |
| ), |
| ] |
| if text is not None: |
| content.append( |
| ChatCompletionContentPartTextParam(type="text", text=text), |
| ) |
| return ScoreMultiModalParam(content=content) |
|
|
