| 1 | # SPDX-License-Identifier: Apache-2.0 |
| 2 | # SPDX-FileCopyrightText: Copyright contributors to the vLLM project |
| 3 | from typing import Any |
| 4 | |
| 5 | import mteb |
| 6 | import numpy as np |
| 7 | import pytest |
| 8 | import torch |
| 9 | from torch.utils.data import DataLoader |
| 10 | |
| 11 | from tests.conftest import HfRunner |
| 12 | from tests.models.utils import RerankModelInfo |
| 13 | |
| 14 | from .mteb_score_utils import MtebCrossEncoderMixin, mteb_test_rerank_models |
| 15 | |
| 16 | mxbai_rerank_hf_overrides = { |
| 17 | "architectures": ["Qwen2ForSequenceClassification"], |
| 18 | "classifier_from_token": ["0", "1"], |
| 19 | "method": "from_2_way_softmax", |
| 20 | } |
| 21 | |
| 22 | RERANK_MODELS = [ |
| 23 | RerankModelInfo( |
| 24 | "mixedbread-ai/mxbai-rerank-base-v2", |
| 25 | architecture="Qwen2ForSequenceClassification", |
| 26 | hf_overrides=mxbai_rerank_hf_overrides, |
| 27 | seq_pooling_type="LAST", |
| 28 | attn_type="decoder", |
| 29 | is_prefix_caching_supported=True, |
| 30 | is_chunked_prefill_supported=True, |
| 31 | chat_template_name="mxbai_rerank_v2.jinja", |
| 32 | mteb_score=0.33651, |
| 33 | mteb_tol=1e-2, |
| 34 | enable_test=True, |
| 35 | ), |
| 36 | RerankModelInfo( |
| 37 | "mixedbread-ai/mxbai-rerank-large-v2", |
| 38 | architecture="Qwen2ForSequenceClassification", |
| 39 | hf_overrides=mxbai_rerank_hf_overrides, |
| 40 | chat_template_name="mxbai_rerank_v2.jinja", |
| 41 | enable_test=False, |
| 42 | ), |
| 43 | ] |
| 44 | |
| 45 | |
| 46 | class MxbaiRerankerHfRunner(MtebCrossEncoderMixin, HfRunner): |
| 47 | def __init__( |
| 48 | self, model_name: str, dtype: str = "auto", *args: Any, **kwargs: Any |
| 49 | ) -> None: |
| 50 | from transformers import AutoModelForCausalLM, AutoTokenizer |
| 51 | |
| 52 | HfRunner.__init__( |
| 53 | self, |
| 54 | model_name=model_name, |
| 55 | auto_cls=AutoModelForCausalLM, |
| 56 | dtype=dtype, |
| 57 | **kwargs, |
| 58 | ) |
| 59 | |
| 60 | self.tokenizer = AutoTokenizer.from_pretrained(model_name, padding_side="left") |
| 61 | self.yes_loc = self.tokenizer.convert_tokens_to_ids("1") |
| 62 | self.no_loc = self.tokenizer.convert_tokens_to_ids("0") |
| 63 | |
| 64 | @torch.no_grad |
| 65 | def predict( |
| 66 | self, |
| 67 | inputs1: DataLoader[mteb.types.BatchedInput], |
| 68 | inputs2: DataLoader[mteb.types.BatchedInput], |
| 69 | *args, |
| 70 | **kwargs, |
| 71 | ) -> np.ndarray: |
| 72 | queries = [text for batch in inputs1 for text in batch["text"]] |
| 73 | corpus = [text for batch in inputs2 for text in batch["text"]] |
| 74 | |
| 75 | tokenizer = self.tokenizer |
| 76 | prompts = [] |
| 77 | for query, document in zip(queries, corpus): |
| 78 | conversation = [ |
| 79 | {"role": "query", "content": query}, |
| 80 | {"role": "document", "content": document}, |
| 81 | ] |
| 82 | |
| 83 | prompt = tokenizer.apply_chat_template( |
| 84 | conversation=conversation, |
| 85 | tools=None, |
| 86 | chat_template=self.chat_template, |
| 87 | tokenize=False, |
| 88 | ) |
| 89 | prompts.append(prompt) |
| 90 | |
| 91 | def compute_logits(inputs): |
| 92 | logits = self.model(**inputs).logits[:, -1, :] |
| 93 | yes_logits = logits[:, self.yes_loc] |
| 94 | no_logits = logits[:, self.no_loc] |
| 95 | logits = yes_logits - no_logits |
| 96 | scores = logits.float().sigmoid() |
| 97 | return scores |
| 98 | |
| 99 | scores = [] |
| 100 | for prompt in prompts: |
| 101 | inputs = tokenizer([prompt], return_tensors="pt") |
| 102 | inputs = self.wrap_device(inputs) |
| 103 | score = compute_logits(inputs) |
| 104 | scores.append(score[0].item()) |
| 105 | return torch.Tensor(scores) |
| 106 | |
| 107 | |
| 108 | @pytest.mark.flaky(reruns=2) |
| 109 | @pytest.mark.parametrize("model_info", RERANK_MODELS) |
| 110 | def test_rerank_models_mteb(vllm_runner, model_info: RerankModelInfo) -> None: |
| 111 | mteb_test_rerank_models(vllm_runner, model_info, hf_runner=MxbaiRerankerHfRunner) |