From d79c485fbd6c027e90864ab088b535332b7663ac Mon Sep 17 00:00:00 2001 From: "shengzhe.li" Date: Wed, 25 Jun 2025 18:07:25 +0900 Subject: [PATCH 1/4] Adapt cls and mean pooling in Signed-off-by: shengzhe.li --- vllm/model_executor/models/modernbert.py | 8 +++++++- 1 file changed, 7 insertions(+), 1 deletion(-) diff --git a/vllm/model_executor/models/modernbert.py b/vllm/model_executor/models/modernbert.py index 7c1f889e8f3..6f390860c12 100644 --- a/vllm/model_executor/models/modernbert.py +++ b/vllm/model_executor/models/modernbert.py @@ -258,6 +258,7 @@ def __init__(self, config: ModernBertConfig): super().__init__() self.dense = nn.Linear(config.hidden_size, config.hidden_size, config.classifier_bias) + self.pooling_type = config.classifier_pooling self.act = nn.GELU() self.norm = nn.LayerNorm(config.hidden_size, eps=config.norm_eps, @@ -265,7 +266,12 @@ def __init__(self, config: ModernBertConfig): def forward(self, hidden_states: torch.Tensor) -> torch.Tensor: pooled_output = hidden_states - pooled_output = pooled_output.mean(dim=0, keepdim=False) + if self.pooling_type == "mean": + pooled_output = pooled_output.mean(dim=0, keepdim=False) + elif self.pooling_type == "cls": + pooled_output = pooled_output[0, :] + else: + raise ValueError(f"Pooling type should be either `cls` or `mean`, but got {self.pooling_type}") pooled_output = self.norm(self.act(self.dense(pooled_output))) return pooled_output From bf343f944e5c943914eaef81832e4d38728a2af6 Mon Sep 17 00:00:00 2001 From: "shengzhe.li" Date: Wed, 25 Jun 2025 18:33:48 +0900 Subject: [PATCH 2/4] reformat Signed-off-by: shengzhe.li --- vllm/model_executor/models/modernbert.py | 4 +++- 1 file changed, 3 insertions(+), 1 deletion(-) diff --git a/vllm/model_executor/models/modernbert.py b/vllm/model_executor/models/modernbert.py index 6f390860c12..91bff4f941e 100644 --- a/vllm/model_executor/models/modernbert.py +++ b/vllm/model_executor/models/modernbert.py @@ -271,7 +271,9 @@ def forward(self, hidden_states: torch.Tensor) -> torch.Tensor: elif self.pooling_type == "cls": pooled_output = pooled_output[0, :] else: - raise ValueError(f"Pooling type should be either `cls` or `mean`, but got {self.pooling_type}") + raise ValueError( + f"Pooling type should be either `cls` or `mean`, but got {self.pooling_type}" + ) pooled_output = self.norm(self.act(self.dense(pooled_output))) return pooled_output From cbfa28dacc3f1a4118614d43d81d43067d859742 Mon Sep 17 00:00:00 2001 From: "shengzhe.li" Date: Wed, 25 Jun 2025 18:52:01 +0900 Subject: [PATCH 3/4] reformat Signed-off-by: shengzhe.li --- vllm/model_executor/models/modernbert.py | 3 ++- 1 file changed, 2 insertions(+), 1 deletion(-) diff --git a/vllm/model_executor/models/modernbert.py b/vllm/model_executor/models/modernbert.py index 91bff4f941e..f3e6b6da868 100644 --- a/vllm/model_executor/models/modernbert.py +++ b/vllm/model_executor/models/modernbert.py @@ -272,7 +272,8 @@ def forward(self, hidden_states: torch.Tensor) -> torch.Tensor: pooled_output = pooled_output[0, :] else: raise ValueError( - f"Pooling type should be either `cls` or `mean`, but got {self.pooling_type}" + "Pooling type should be either `cls` or `mean`, " + f"but got {self.pooling_type}" ) pooled_output = self.norm(self.act(self.dense(pooled_output))) return pooled_output From 643b0905a5d95133142959cef9f389c87ee39ff0 Mon Sep 17 00:00:00 2001 From: "shengzhe.li" Date: Wed, 25 Jun 2025 19:04:41 +0900 Subject: [PATCH 4/4] reformat Signed-off-by: shengzhe.li --- vllm/model_executor/models/modernbert.py | 6 ++---- 1 file changed, 2 insertions(+), 4 deletions(-) diff --git a/vllm/model_executor/models/modernbert.py b/vllm/model_executor/models/modernbert.py index f3e6b6da868..9d619b38d38 100644 --- a/vllm/model_executor/models/modernbert.py +++ b/vllm/model_executor/models/modernbert.py @@ -271,10 +271,8 @@ def forward(self, hidden_states: torch.Tensor) -> torch.Tensor: elif self.pooling_type == "cls": pooled_output = pooled_output[0, :] else: - raise ValueError( - "Pooling type should be either `cls` or `mean`, " - f"but got {self.pooling_type}" - ) + raise ValueError("Pooling type should be either `cls` or `mean`, " + f"but got {self.pooling_type}") pooled_output = self.norm(self.act(self.dense(pooled_output))) return pooled_output