add io_process_plugin for sparse embedding (#34214)

Signed-off-by: augusto.yjh <augusto.yjh@antgroup.com>
Signed-off-by: Augusto Yao <augusto.yjh@antgroup.com>
Co-authored-by: Cyrus Leung <cyrus.tl.leung@gmail.com>

tests/plugins_tests/test_bge_m3_sparse_io_processor_plugins.py

@@ -0,0 +1,212 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+
+import json
+
+import pytest
+import requests
+
+# Test configuration for BGE-M3 sparse plugin
+from tests.utils import RemoteOpenAIServer
+from vllm.entrypoints.pooling.pooling.protocol import IOProcessorResponse
+
+model_config = {
+    "model_name": "BAAI/bge-m3",
+    "plugin": "bge_m3_sparse_plugin",
+    "test_input": "What is the capital of France?",
+    "hf_overrides": json.dumps(
+        {"architectures": ["BgeM3EmbeddingModel"], "head_dtype": "float16"}
+    ),
+}
+
+
+def _float_close(expected: object, result: object):
+    assert isinstance(expected, float) and isinstance(result, float), (
+        f"{expected=}  or {result=} is not float"
+    )
+    return (expected - result) < 1e-3 or abs(expected / result - 1) < 1e-3
+
+
+def _get_attr_or_val(obj: object | dict, key: str):
+    if isinstance(obj, dict) and key in obj:
+        return obj[key]
+    return getattr(obj, key, None)
+
+
+def _check_sparse_embedding(data, check_tokens=False):
+    expected_weights = [
+        {"token_id": 32, "weight": 0.0552978515625, "token": "?"},
+        {"token_id": 70, "weight": 0.09808349609375, "token": "the"},
+        {"token_id": 83, "weight": 0.08154296875, "token": "is"},
+        {"token_id": 111, "weight": 0.11810302734375, "token": "of"},
+        {"token_id": 4865, "weight": 0.1171875, "token": "What"},
+        {"token_id": 9942, "weight": 0.292236328125, "token": "France"},
+        {"token_id": 10323, "weight": 0.2802734375, "token": "capital"},
+    ]
+    expected_embed = {x["token_id"]: x for x in expected_weights}
+
+    assert len(data) == len(expected_embed)
+    for entry in data:
+        expected_val = expected_embed[_get_attr_or_val(entry, "token_id")]
+        assert _float_close(
+            expected_val["weight"], _get_attr_or_val(entry, "weight")
+        ), f"actual embed {entry} not equal to {expected_val}"
+        if check_tokens:
+            assert expected_val["token"] == _get_attr_or_val(entry, "token"), (
+                f"actual embed {entry} not equal to {expected_val}"
+            )
+        else:
+            assert _get_attr_or_val(entry, "token") is None, (
+                f"{entry} should not return token"
+            )
+
+
+@pytest.fixture(scope="function")
+def server():
+    args = [
+        "--runner",
+        "pooling",
+        "--enforce-eager",
+        "--max-num-seqs",
+        "32",
+        "--hf_overrides",
+        model_config["hf_overrides"],
+        "--io-processor-plugin",
+        model_config["plugin"],
+    ]
+
+    with RemoteOpenAIServer(model_config["model_name"], args) as remote_server:
+        yield remote_server
+
+
+@pytest.mark.asyncio
+@pytest.mark.parametrize(
+    "return_tokens",
+    [True, False],
+)
+async def test_bge_m3_sparse_plugin_online(
+    server: RemoteOpenAIServer, return_tokens: bool
+):
+    """Test BGE-M3 sparse plugin in online mode via API."""
+    request_payload = {
+        "model": model_config["model_name"],
+        "task": "token_classify",
+        "data": {"input": model_config["test_input"], "return_tokens": return_tokens},
+    }
+
+    ret = requests.post(
+        server.url_for("pooling"),
+        json=request_payload,
+    )
+
+    response = ret.json()
+
+    # Verify the request response is in the correct format
+    assert (parsed_response := IOProcessorResponse(**response).data)
+
+    # Verify the output is formatted as expected for this plugin
+    assert _get_attr_or_val(parsed_response, "data")
+    assert len(_get_attr_or_val(parsed_response, "data")) > 0
+
+    data_entry = _get_attr_or_val(parsed_response, "data")[0]
+    assert _get_attr_or_val(data_entry, "object") == "sparse-embedding"
+    assert _get_attr_or_val(data_entry, "sparse_embedding")
+
+    # Verify sparse embedding format
+    sparse_embedding = _get_attr_or_val(data_entry, "sparse_embedding")
+    assert isinstance(sparse_embedding, list)
+    _check_sparse_embedding(sparse_embedding, return_tokens)
+
+    # Verify usage information
+    usage = _get_attr_or_val(parsed_response, "usage")
+    assert usage, f"usage not found for {parsed_response}"
+    assert _get_attr_or_val(usage, "prompt_tokens") > 0
+    assert _get_attr_or_val(usage, "total_tokens") == _get_attr_or_val(
+        usage, "prompt_tokens"
+    )
+
+
+@pytest.mark.parametrize(
+    "return_tokens",
+    [True, False],
+)
+def test_bge_m3_sparse_plugin_offline(vllm_runner, return_tokens: bool):
+    """Test BGE-M3 sparse plugin in offline mode."""
+    prompt = {
+        "data": {
+            "input": model_config["test_input"],
+            "return_tokens": return_tokens,
+        }
+    }
+
+    with vllm_runner(
+        model_config["model_name"],
+        runner="pooling",
+        enforce_eager=True,
+        max_num_seqs=32,
+        io_processor_plugin=model_config["plugin"],
+        hf_overrides=json.loads(model_config["hf_overrides"]),
+        default_torch_num_threads=1,
+    ) as llm_runner:
+        llm = llm_runner.get_llm()
+        pooler_output = llm.encode(prompt, pooling_task="token_classify")
+
+    outputs = pooler_output[0]
+
+    # Verify output structure
+    assert hasattr(outputs, "outputs")
+    response = outputs.outputs
+    assert hasattr(response, "data")
+    assert len(response.data) == 1
+    # Verify response data
+    for i, output in enumerate(response.data):
+        # Each output should have sparse embeddings
+        sparse_embedding = output.sparse_embedding
+        assert isinstance(sparse_embedding, list)
+        _check_sparse_embedding(sparse_embedding, return_tokens)
+
+    # Verify usage
+    assert response.usage.prompt_tokens > 0
+    assert response.usage.total_tokens == response.usage.prompt_tokens
+
+
+def test_bge_m3_sparse_plugin_offline_multiple_inputs(vllm_runner):
+    """Test BGE-M3 sparse plugin with multiple inputs in offline mode."""
+    prompts = {
+        "data": {
+            "input": [
+                "What is the capital of France?",
+                "What is the capital of Germany?",
+                "What is the capital of Spain?",
+            ],
+            "return_tokens": True,
+        }
+    }
+
+    with vllm_runner(
+        model_config["model_name"],
+        runner="pooling",
+        enforce_eager=True,
+        max_num_seqs=32,
+        io_processor_plugin=model_config["plugin"],
+        hf_overrides=json.loads(model_config["hf_overrides"]),
+        default_torch_num_threads=1,
+    ) as llm_runner:
+        llm = llm_runner.get_llm()
+        pooler_output = llm.encode(prompts, pooling_task="token_classify")
+
+    outputs = pooler_output[0]
+
+    # Verify output structure
+    assert hasattr(outputs, "outputs")
+    response = outputs.outputs
+    assert hasattr(response, "data")
+    assert len(response.data) == 3
+    for i, output in enumerate(response.data):
+        # Each output should have sparse embeddings
+        sparse_embedding = output.sparse_embedding
+        assert isinstance(sparse_embedding, list)
+
+    # Verify usage
+    assert response.usage.prompt_tokens > 0
+    assert response.usage.total_tokens == response.usage.prompt_tokens

tests/plugins_tests/test_io_processor_plugins.py

@@ -39,7 +39,7 @@ def _compute_image_hash(base64_data: str) -> str:
 def test_loading_missing_plugin():
     vllm_config = VllmConfig()
     with pytest.raises(ValueError):
-        get_io_processor(vllm_config, "wrong_plugin")
+        get_io_processor(vllm_config, None, "wrong_plugin")
 
 
 @pytest.fixture(scope="function")