| 1 | # SPDX-License-Identifier: Apache-2.0 |
| 2 | # SPDX-FileCopyrightText: Copyright contributors to the vLLM project |
| 3 | # ruff: noqa: E501 |
| 4 | from typing import Any |
| 5 | |
| 6 | import mteb |
| 7 | import numpy as np |
| 8 | import pytest |
| 9 | import torch |
| 10 | from torch.utils.data import DataLoader |
| 11 | |
| 12 | from tests.conftest import HfRunner |
| 13 | from tests.models.utils import RerankModelInfo |
| 14 | from tests.utils import multi_gpu_test |
| 15 | |
| 16 | from .mteb_score_utils import MtebCrossEncoderMixin, mteb_test_rerank_models |
| 17 | |
| 18 | qwen3_reranker_hf_overrides = { |
| 19 | "architectures": ["Qwen3ForSequenceClassification"], |
| 20 | "classifier_from_token": ["no", "yes"], |
| 21 | "is_original_qwen3_reranker": True, |
| 22 | } |
| 23 | |
| 24 | RERANK_MODELS = [ |
| 25 | RerankModelInfo( |
| 26 | "Qwen/Qwen3-Reranker-0.6B", |
| 27 | architecture="Qwen3ForSequenceClassification", |
| 28 | hf_overrides=qwen3_reranker_hf_overrides, |
| 29 | chat_template_name="qwen3_reranker.jinja", |
| 30 | seq_pooling_type="LAST", |
| 31 | attn_type="decoder", |
| 32 | is_prefix_caching_supported=True, |
| 33 | is_chunked_prefill_supported=True, |
| 34 | mteb_score=0.33459, |
| 35 | mteb_tol=1e-2, |
| 36 | enable_test=True, |
| 37 | ), |
| 38 | RerankModelInfo( |
| 39 | "Qwen/Qwen3-Reranker-4B", |
| 40 | architecture="Qwen3ForSequenceClassification", |
| 41 | chat_template_name="qwen3_reranker.jinja", |
| 42 | hf_overrides=qwen3_reranker_hf_overrides, |
| 43 | enable_test=False, |
| 44 | ), |
| 45 | ] |
| 46 | |
| 47 | |
| 48 | class Qwen3RerankerHfRunner(MtebCrossEncoderMixin, HfRunner): |
| 49 | def __init__( |
| 50 | self, model_name: str, dtype: str = "auto", *args: Any, **kwargs: Any |
| 51 | ) -> None: |
| 52 | from transformers import AutoModelForCausalLM, AutoTokenizer |
| 53 | |
| 54 | HfRunner.__init__( |
| 55 | self, |
| 56 | model_name=model_name, |
| 57 | auto_cls=AutoModelForCausalLM, |
| 58 | dtype=dtype, |
| 59 | **kwargs, |
| 60 | ) |
| 61 | |
| 62 | self.tokenizer = AutoTokenizer.from_pretrained(model_name, padding_side="left") |
| 63 | self.token_false_id = self.tokenizer.convert_tokens_to_ids("no") |
| 64 | self.token_true_id = self.tokenizer.convert_tokens_to_ids("yes") |
| 65 | self.max_length = 40960 |
| 66 | |
| 67 | @torch.no_grad |
| 68 | def predict( |
| 69 | self, |
| 70 | inputs1: DataLoader[mteb.types.BatchedInput], |
| 71 | inputs2: DataLoader[mteb.types.BatchedInput], |
| 72 | *args, |
| 73 | **kwargs, |
| 74 | ) -> np.ndarray: |
| 75 | queries = [text for batch in inputs1 for text in batch["text"]] |
| 76 | corpus = [text for batch in inputs2 for text in batch["text"]] |
| 77 | |
| 78 | tokenizer = self.tokenizer |
| 79 | prompts = [] |
| 80 | for query, document in zip(queries, corpus): |
| 81 | conversation = [ |
| 82 | {"role": "query", "content": query}, |
| 83 | {"role": "document", "content": document}, |
| 84 | ] |
| 85 | |
| 86 | prompt = tokenizer.apply_chat_template( |
| 87 | conversation=conversation, |
| 88 | tools=None, |
| 89 | chat_template=self.chat_template, |
| 90 | tokenize=False, |
| 91 | ) |
| 92 | prompts.append(prompt) |
| 93 | |
| 94 | def compute_logits(inputs): |
| 95 | batch_scores = self.model(**inputs).logits[:, -1, :] |
| 96 | true_vector = batch_scores[:, self.token_true_id] |
| 97 | false_vector = batch_scores[:, self.token_false_id] |
| 98 | batch_scores = torch.stack([false_vector, true_vector], dim=1) |
| 99 | batch_scores = torch.nn.functional.log_softmax(batch_scores, dim=1) |
| 100 | scores = batch_scores[:, 1].exp() |
| 101 | return scores |
| 102 | |
| 103 | scores = [] |
| 104 | for prompt in prompts: |
| 105 | inputs = tokenizer([prompt], return_tensors="pt") |
| 106 | inputs = self.wrap_device(inputs) |
| 107 | score = compute_logits(inputs) |
| 108 | scores.append(score[0].item()) |
| 109 | return torch.Tensor(scores) |
| 110 | |
| 111 | |
| 112 | @pytest.mark.flaky(reruns=2) |
| 113 | @pytest.mark.parametrize("model_info", RERANK_MODELS) |
| 114 | def test_rerank_models_mteb(vllm_runner, model_info: RerankModelInfo) -> None: |
| 115 | mteb_test_rerank_models(vllm_runner, model_info, hf_runner=Qwen3RerankerHfRunner) |
| 116 | |
| 117 | |
| 118 | @pytest.mark.flaky(reruns=2) |
| 119 | @pytest.mark.parametrize("model_info", RERANK_MODELS) |
| 120 | @multi_gpu_test(num_gpus=2) |
| 121 | def test_rerank_models_mteb_tp(vllm_runner, model_info: RerankModelInfo) -> None: |
| 122 | assert model_info.architecture == "Qwen3ForSequenceClassification" |
| 123 | |
| 124 | vllm_extra_kwargs: dict[str, Any] = { |
| 125 | "tensor_parallel_size": 2, |
| 126 | } |
| 127 | |
| 128 | mteb_test_rerank_models( |
| 129 | vllm_runner, |
| 130 | model_info, |
| 131 | vllm_extra_kwargs=vllm_extra_kwargs, |
| 132 | hf_runner=Qwen3RerankerHfRunner, |
| 133 | ) |