+ ping_pong_test_score_models for test_rerank_models_correctness

noooop · noooop · commit 6dc55bab9d98 · 2025-06-18T16:31:53.000+08:00
diff --git a/examples/offline_inference/converting2seq_cls_models.py b/examples/offline_inference/converting2seq_cls_models.py
@@ -9,8 +9,9 @@
 import transformers
 
 
-def from_2_way_softmax(causal_lm, seq_cls_model, tokenizer,
-                       classifier_from_tokens, device):
+def from_2_way_softmax(
+    causal_lm, seq_cls_model, tokenizer, classifier_from_tokens, device
+):
     # for Qwen3-Reranker
     assert len(classifier_from_tokens) == 2
 
@@ -20,32 +21,53 @@ def from_2_way_softmax(causal_lm, seq_cls_model, tokenizer,
     b = tokenizer.convert_tokens_to_ids(classifier_from_tokens[1])
 
     score_weight = lm_head_weights[b].to(torch.float32).to(device).to(
-        torch.float32) - lm_head_weights[a].to(device)
+        torch.float32
+    ) - lm_head_weights[a].to(device)
 
     with torch.no_grad():
         seq_cls_model.score.weight.copy_(score_weight.unsqueeze(0))
         if seq_cls_model.score.bias is not None:
             seq_cls_model.score.bias.zero_()
 
 
-method_map = {function.__name__: function for function in [from_2_way_softmax]}
+def from_1_way_sigmoid(
+    causal_lm, seq_cls_model, tokenizer, classifier_from_tokens, device
+):
+    # for BAAI/bge-reranker-v2-gemma
+    assert len(classifier_from_tokens) == 1
+
+    lm_head_weights = causal_lm.lm_head.weight
+
+    a = tokenizer.convert_tokens_to_ids(classifier_from_tokens[0])
+
+    score_weight = lm_head_weights[a].to(device)
+
+    with torch.no_grad():
+        seq_cls_model.score.weight.copy_(score_weight.unsqueeze(0))
+        if seq_cls_model.score.bias is not None:
+            seq_cls_model.score.bias.zero_()
+
+
+method_map = {
+    function.__name__: function for function in [from_2_way_softmax, from_1_way_sigmoid]
+}
 
 
 def converting(model_name, classifier_from_tokens, path, method, device="cpu"):
     assert method in method_map
 
     tokenizer = transformers.AutoTokenizer.from_pretrained(model_name)
     causal_lm = transformers.AutoModelForCausalLM.from_pretrained(
-        model_name, device_map=device)
+        model_name, device_map=device
+    )
 
     seq_cls_model = transformers.AutoModelForSequenceClassification.from_pretrained(
-        model_name,
-        num_labels=1,
-        ignore_mismatched_sizes=True,
-        device_map=device)
+        model_name, num_labels=1, ignore_mismatched_sizes=True, device_map=device
+    )
 
-    method_map[method](causal_lm, seq_cls_model, tokenizer,
-                       classifier_from_tokens, device)
+    method_map[method](
+        causal_lm, seq_cls_model, tokenizer, classifier_from_tokens, device
+    )
 
     seq_cls_model.config.pad_token_id = tokenizer.pad_token_id
 
@@ -55,31 +77,35 @@ def converting(model_name, classifier_from_tokens, path, method, device="cpu"):
 
 def parse_args():
     parser = argparse.ArgumentParser(
-        description=
-        "Converting *ForCausalLM models to *ForSequenceClassification models.")
-    parser.add_argument("--model_name",
-                        type=str,
-                        default="Qwen/Qwen3-Reranker-0.6B",
-                        help="Model name")
-    parser.add_argument("--classifier_from_tokens",
-                        type=str,
-                        default='["no", "yes"]',
-                        help="classifier from tokens")
-    parser.add_argument("--method",
-                        type=str,
-                        default='from_2_way_softmax',
-                        help="Converting converting")
-    parser.add_argument("--path",
-                        type=str,
-                        default="./converted_model",
-                        help="Path to save converted model")
+        description="Converting *ForCausalLM models to *ForSequenceClassification models."
+    )
+    parser.add_argument(
+        "--model_name", type=str, default="Qwen/Qwen3-Reranker-0.6B", help="Model name"
+    )
+    parser.add_argument(
+        "--classifier_from_tokens",
+        type=str,
+        default='["no", "yes"]',
+        help="classifier from tokens",
+    )
+    parser.add_argument(
+        "--method", type=str, default="from_2_way_softmax", help="Converting converting"
+    )
+    parser.add_argument(
+        "--path",
+        type=str,
+        default="./converted_model",
+        help="Path to save converted model",
+    )
     return parser.parse_args()
 
 
 if __name__ == "__main__":
     args = parse_args()
 
-    converting(model_name=args.model_name,
-               classifier_from_tokens=json.loads(args.classifier_from_tokens),
-               method=args.method,
-               path=args.path)
+    converting(
+        model_name=args.model_name,
+        classifier_from_tokens=json.loads(args.classifier_from_tokens),
+        method=args.method,
+        path=args.path,
+    )
diff --git a/tests/models/language/pooling/score_utils.py b/tests/models/language/pooling/score_utils.py
@@ -0,0 +1,56 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+import math
+
+import pytest
+
+from tests.models.utils import RerankModelInfo
+
+
+def ping_pong_test_score_models(hf_runner,
+                                vllm_runner,
+                                model_info: RerankModelInfo,
+                                vllm_extra_kwargs=None,
+                                hf_model_callback=None):
+    if not model_info.enable_test:
+        # A model family has many models with the same architecture,
+        # and we don't need to test each one.
+        pytest.skip("Skipping test.")
+
+    sentences = []
+
+    vllm_extra_kwargs = vllm_extra_kwargs or {}
+    # This test must use float32 to pass.
+    vllm_extra_kwargs["dtype"] = "float32"
+
+    with vllm_runner(model_info.name,
+                     task="score",
+                     max_model_len=None,
+                     **vllm_extra_kwargs) as vllm_model:
+
+        max_model_len = vllm_model.model.llm_engine.model_config.max_model_len
+
+        for i in range(0, int(math.log2(max_model_len - 1))):
+            sentences.append(("ping", "pong" * 2**i))
+
+        text_1 = [x[0] for x in sentences]
+        text_2 = [x[1] for x in sentences]
+        vllm_outputs = vllm_model.score(text_1=text_1, text_2=text_2)
+
+    with hf_runner(
+            model_info.name,
+            dtype="float32",
+            is_cross_encoder=True,
+    ) as hf_model:
+
+        if hf_model_callback is not None:
+            hf_model_callback(hf_model)
+
+        # use batchsize = 1 to avoid oom
+        hf_outputs = [
+            hf_model.predict([sentences[i]])[0] for i in range(len(sentences))
+        ]
+
+    for i in range(len(sentences)):
+        assert float(hf_outputs[i]) == pytest.approx(float(vllm_outputs[i]), rel=0.01), \
+            f"Test failed at #{i}, vllm: {vllm_outputs[i]}, st: {hf_outputs[i]}"
diff --git a/tests/models/language/pooling/test_baai.py b/tests/models/language/pooling/test_baai.py
@@ -5,6 +5,7 @@
 from ...utils import EmbedModelInfo, RerankModelInfo
 from .embed_utils import correctness_test_embed_models
 from .mteb_utils import mteb_test_embed_models, mteb_test_rerank_models
+from .score_utils import ping_pong_test_score_models
 
 MODELS = [
     ########## BertModel
@@ -91,3 +92,9 @@ def test_embed_models_correctness(hf_runner, vllm_runner,
 def test_rerank_models_mteb(hf_runner, vllm_runner,
                             model_info: RerankModelInfo) -> None:
     mteb_test_rerank_models(hf_runner, vllm_runner, model_info)
+
+
+@pytest.mark.parametrize("model_info", RERANK_MODELS)
+def test_rerank_models_correctness(hf_runner, vllm_runner,
+                                   model_info: RerankModelInfo) -> None:
+    ping_pong_test_score_models(hf_runner, vllm_runner, model_info)
diff --git a/tests/models/language/pooling/test_cross_encoder.py b/tests/models/language/pooling/test_cross_encoder.py
@@ -3,6 +3,7 @@
 import pytest
 
 from .mteb_utils import RerankModelInfo, mteb_test_rerank_models
+from .score_utils import ping_pong_test_score_models
 
 RERANK_MODELS = [
     RerankModelInfo("cross-encoder/ms-marco-TinyBERT-L-2-v2",
@@ -16,3 +17,9 @@
 def test_rerank_models_mteb(hf_runner, vllm_runner,
                             model_info: RerankModelInfo) -> None:
     mteb_test_rerank_models(hf_runner, vllm_runner, model_info)
+
+
+@pytest.mark.parametrize("model_info", RERANK_MODELS)
+def test_rerank_models_correctness(hf_runner, vllm_runner,
+                                   model_info: RerankModelInfo) -> None:
+    ping_pong_test_score_models(hf_runner, vllm_runner, model_info)
diff --git a/tests/models/language/pooling/test_gte.py b/tests/models/language/pooling/test_gte.py
@@ -7,6 +7,7 @@
 from ...utils import RerankModelInfo
 from .embed_utils import EmbedModelInfo, correctness_test_embed_models
 from .mteb_utils import mteb_test_embed_models, mteb_test_rerank_models
+from .score_utils import ping_pong_test_score_models
 
 MODELS = [
     ########## BertModel
@@ -93,3 +94,9 @@ def test_embed_models_correctness(hf_runner, vllm_runner,
 def test_rerank_models_mteb(hf_runner, vllm_runner,
                             model_info: RerankModelInfo) -> None:
     mteb_test_rerank_models(hf_runner, vllm_runner, model_info)
+
+
+@pytest.mark.parametrize("model_info", RERANK_MODELS)
+def test_rerank_models_correctness(hf_runner, vllm_runner,
+                                   model_info: RerankModelInfo) -> None:
+    ping_pong_test_score_models(hf_runner, vllm_runner, model_info)
diff --git a/tests/models/language/pooling/test_jina.py b/tests/models/language/pooling/test_jina.py
@@ -10,6 +10,7 @@
 from .embed_utils import (check_embeddings_close,
                           correctness_test_embed_models, matryoshka_fy)
 from .mteb_utils import mteb_test_embed_models, mteb_test_rerank_models
+from .score_utils import ping_pong_test_score_models
 
 EMBEDDING_MODELS = [
     EmbedModelInfo("jinaai/jina-embeddings-v3",
@@ -60,6 +61,12 @@ def test_rerank_models_mteb(hf_runner, vllm_runner,
     mteb_test_rerank_models(hf_runner, vllm_runner, model_info)
 
 
+@pytest.mark.parametrize("model_info", RERANK_MODELS)
+def test_rerank_models_correctness(hf_runner, vllm_runner,
+                                   model_info: RerankModelInfo) -> None:
+    ping_pong_test_score_models(hf_runner, vllm_runner, model_info)
+
+
 @pytest.mark.parametrize("model_info", EMBEDDING_MODELS)
 @pytest.mark.parametrize("dtype", ["half"])
 @pytest.mark.parametrize("dimensions", [16, 32])
