Skip to content
Merged
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
12 changes: 8 additions & 4 deletions tests/models/language/pooling/test_splade_sparse_pooler.py
Original file line number Diff line number Diff line change
@@ -1,8 +1,6 @@
# SPDX-License-Identifier: Apache-2.0
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project

import types

import pytest
import torch
import torch.nn as nn
Expand All @@ -11,6 +9,8 @@
BertMLMHead,
SPLADESparsePooler,
)
from vllm.pooling_params import PoolingParams
from vllm.v1.pool.metadata import PoolingMetadata, PoolingStates

# ---------------------------------------------------------------------
# Functional test: SPLADE formula correctness (no HF download needed)
Expand Down Expand Up @@ -38,8 +38,12 @@ def test_splade_pooler_matches_reference_formula(B, T, H, V):
],
dtype=torch.long,
)
meta = types.SimpleNamespace(
prompt_lens=prompt_lens_tenser, prompt_token_ids=token_ids
meta = PoolingMetadata(
prompt_lens=prompt_lens_tenser,
prompt_token_ids=token_ids,
prompt_token_ids_cpu=token_ids,
pooling_params=[PoolingParams(task="embed")] * B,
Comment thread
haosdent marked this conversation as resolved.
pooling_states=[PoolingStates() for _ in range(B)],
)

# MLM head (prefer BertMLMHead, fallback to Linear if unavailable)
Expand Down
Loading