@tduncan / Nvidia-Optimization-Stack-VLLM Mirror · read-only / commits / ca4eb82bcb

Fork 0

A high-throughput and memory-efficient inference and serving engine for LLMs

Mirrored from https://github.com/tduncan-lum/Nvidia-Optimization-Stack---VLLM.git

Synced Nov 15, 01:00 UTC

[Model] Re-add the implicit conversion feature for as_seq_cls_model (#21103)

Signed-off-by: wang.yuqi <noooop@126.com>

wang.yuqi committed Jul 18, 2025 at 15:15 UTC ca4eb82bcba97f4fb0a377287ff4d36e19c3d33e

11 files changed +165 -75

tests/models/registry.py

+20 -12

@@ -265,7 +265,6 @@ _TEXT_GENERATION_EXAMPLE_MODELS = {

265 "Qwen2MoeForCausalLM": _HfExamplesInfo("Qwen/Qwen1.5-MoE-A2.7B-Chat"),

266 "Qwen3ForCausalLM": _HfExamplesInfo("Qwen/Qwen3-8B"),

267 "Qwen3MoeForCausalLM": _HfExamplesInfo("Qwen/Qwen3-30B-A3B"),

268 - "Qwen3ForSequenceClassification": _HfExamplesInfo("tomaarsen/Qwen3-Reranker-0.6B-seq-cls"), # noqa: E501

268 "RWForCausalLM": _HfExamplesInfo("tiiuae/falcon-40b"),

269 "StableLMEpochForCausalLM": _HfExamplesInfo("stabilityai/stablelm-zephyr-3b"), # noqa: E501

270 "StableLmForCausalLM": _HfExamplesInfo("stabilityai/stablelm-3b-4e1t"),

@@ -292,7 +291,6 @@ _EMBEDDING_EXAMPLE_MODELS = {

291 # [Text-only]

292 "BertModel": _HfExamplesInfo("BAAI/bge-base-en-v1.5", v0_only=True),

293 "Gemma2Model": _HfExamplesInfo("BAAI/bge-multilingual-gemma2", v0_only=True), # noqa: E501

295 - "GPT2ForSequenceClassification": _HfExamplesInfo("nie3e/sentiment-polish-gpt2-small"), # noqa: E501

294 "GritLM": _HfExamplesInfo("parasail-ai/GritLM-7B-vllm"),

295 "GteModel": _HfExamplesInfo("Snowflake/snowflake-arctic-embed-m-v2.0",

296 trust_remote_code=True),

@@ -311,7 +309,6 @@ _EMBEDDING_EXAMPLE_MODELS = {

309 "Qwen2Model": _HfExamplesInfo("ssmits/Qwen2-7B-Instruct-embed-base"),

310 "Qwen2ForRewardModel": _HfExamplesInfo("Qwen/Qwen2.5-Math-RM-72B"),

311 "Qwen2ForProcessRewardModel": _HfExamplesInfo("Qwen/Qwen2.5-Math-PRM-7B"),

314 - "Qwen2ForSequenceClassification": _HfExamplesInfo("jason9693/Qwen2.5-1.5B-apeach"), # noqa: E501

312 "RobertaModel": _HfExamplesInfo("sentence-transformers/stsb-roberta-base-v2", v0_only=True), # noqa: E501

313 "RobertaForMaskedLM": _HfExamplesInfo("sentence-transformers/all-roberta-large-v1", v0_only=True), # noqa: E501

314 "XLMRobertaModel": _HfExamplesInfo("intfloat/multilingual-e5-small", v0_only=True), # noqa: E501

@@ -324,20 +321,29 @@ _EMBEDDING_EXAMPLE_MODELS = {

321 is_available_online=False), # noqa: E501

322 }

323

327 -_CROSS_ENCODER_EXAMPLE_MODELS = {

328 - # [Text-only]

324 +_SEQUENCE_CLASSIFICATION_EXAMPLE_MODELS = {

325 + # [Decoder-only]

326 + "GPT2ForSequenceClassification": _HfExamplesInfo("nie3e/sentiment-polish-gpt2-small"), # noqa: E501

327 +

328 + # [Cross-encoder]

329 "BertForSequenceClassification": _HfExamplesInfo("cross-encoder/ms-marco-MiniLM-L-6-v2", v0_only=True), # noqa: E501

330 - "GemmaForSequenceClassification": _HfExamplesInfo("BAAI/bge-reranker-v2-gemma", # noqa: E501

331 - v0_only=True,

332 - hf_overrides={"architectures": ["GemmaForSequenceClassification"], # noqa: E501

333 - "classifier_from_token": ["Yes"], # noqa: E501

334 - "method": "no_post_processing"}), # noqa: E501

335 - "LlamaForSequenceClassification": _HfExamplesInfo("Skywork/Skywork-Reward-V2-Llama-3.2-1B"), # noqa: E501

330 "ModernBertForSequenceClassification": _HfExamplesInfo("Alibaba-NLP/gte-reranker-modernbert-base", v0_only=True), # noqa: E501

331 "RobertaForSequenceClassification": _HfExamplesInfo("cross-encoder/quora-roberta-base", v0_only=True), # noqa: E501

332 "XLMRobertaForSequenceClassification": _HfExamplesInfo("BAAI/bge-reranker-v2-m3", v0_only=True), # noqa: E501

333 }

334

335 +_AUTOMATIC_CONVERTED_MODELS = {

336 + # Use as_seq_cls_model for automatic conversion

337 + "GemmaForSequenceClassification": _HfExamplesInfo("BAAI/bge-reranker-v2-gemma", # noqa: E501

338 + v0_only=True,

339 + hf_overrides={"architectures": ["GemmaForSequenceClassification"], # noqa: E501

340 + "classifier_from_token": ["Yes"], # noqa: E501

341 + "method": "no_post_processing"}), # noqa: E501

342 + "LlamaForSequenceClassification": _HfExamplesInfo("Skywork/Skywork-Reward-V2-Llama-3.2-1B"), # noqa: E501

343 + "Qwen2ForSequenceClassification": _HfExamplesInfo("jason9693/Qwen2.5-1.5B-apeach"), # noqa: E501

344 + "Qwen3ForSequenceClassification": _HfExamplesInfo("tomaarsen/Qwen3-Reranker-0.6B-seq-cls"), # noqa: E501

345 +}

346 +

347 _MULTIMODAL_EXAMPLE_MODELS = {

348 # [Decoder-only]

349 "AriaForConditionalGeneration": _HfExamplesInfo("rhymes-ai/Aria"),

@@ -449,6 +455,7 @@ _MULTIMODAL_EXAMPLE_MODELS = {

455 "JinaVLForRanking": _HfExamplesInfo("jinaai/jina-reranker-m0"), # noqa: E501

456 }

457

458 +

459 _SPECULATIVE_DECODING_EXAMPLE_MODELS = {

460 "EAGLEModel": _HfExamplesInfo("JackFram/llama-68m",

461 speculative_model="abhigoyal/vllm-eagle-llama-68m-random"), # noqa: E501

@@ -489,7 +496,7 @@ _TRANSFORMERS_MODELS = {

496 _EXAMPLE_MODELS = {

497 **_TEXT_GENERATION_EXAMPLE_MODELS,

498 **_EMBEDDING_EXAMPLE_MODELS,

492 - **_CROSS_ENCODER_EXAMPLE_MODELS,

499 + **_SEQUENCE_CLASSIFICATION_EXAMPLE_MODELS,

500 **_MULTIMODAL_EXAMPLE_MODELS,

501 **_SPECULATIVE_DECODING_EXAMPLE_MODELS,

502 **_TRANSFORMERS_MODELS,

@@ -522,3 +529,4 @@ class HfExampleModels:

529

530

531 HF_EXAMPLE_MODELS = HfExampleModels(_EXAMPLE_MODELS)

532 +AUTO_EXAMPLE_MODELS = HfExampleModels(_AUTOMATIC_CONVERTED_MODELS)

tests/models/test_initialization.py

+21 -8

@@ -13,20 +13,21 @@ from vllm.v1.core.kv_cache_utils import get_kv_cache_config

13 from vllm.v1.engine.core import EngineCore as V1EngineCore

15 from ..utils import create_new_process_for_each_test

16 -from .registry import HF_EXAMPLE_MODELS

16 +from .registry import AUTO_EXAMPLE_MODELS, HF_EXAMPLE_MODELS, HfExampleModels

19 -@pytest.mark.parametrize("model_arch", HF_EXAMPLE_MODELS.get_supported_archs())

19 @create_new_process_for_each_test()

21 -def test_can_initialize(model_arch: str, monkeypatch: pytest.MonkeyPatch):

22 - """The reason for using create_new_process_for_each_test is to avoid

23 - the WARNING:

24 - "We must use the 'spawn' multiprocessing start method. Overriding

20 +def can_initialize(model_arch: str, monkeypatch: pytest.MonkeyPatch,

21 + EXAMPLE_MODELS: HfExampleModels):

22 + """The reason for using create_new_process_for_each_test is to avoid

23 + the WARNING:

24 + "We must use the 'spawn' multiprocessing start method. Overriding

25 VLLM_WORKER_MULTIPROC_METHOD to 'spawn'."

26 - The spawn process causes the _initialize_kv_caches_v1 function below to

26 + The spawn process causes the _initialize_kv_caches_v1 function below to

27 become ineffective.

28 """

29 - model_info = HF_EXAMPLE_MODELS.get_hf_info(model_arch)

29 +

30 + model_info = EXAMPLE_MODELS.get_hf_info(model_arch)

31 model_info.check_available_online(on_fail="skip")

32 model_info.check_transformers_version(on_fail="skip")

@@ -127,3 +128,15 @@ def test_can_initialize(model_arch: str, monkeypatch: pytest.MonkeyPatch):

128 load_format="dummy",

129 hf_overrides=hf_overrides,

130 )

131 +

132 +

133 +@pytest.mark.parametrize("model_arch", HF_EXAMPLE_MODELS.get_supported_archs())

134 +def test_can_initialize(model_arch: str, monkeypatch: pytest.MonkeyPatch):

135 + can_initialize(model_arch, monkeypatch, HF_EXAMPLE_MODELS)

136 +

137 +

138 +@pytest.mark.parametrize("model_arch",

139 + AUTO_EXAMPLE_MODELS.get_supported_archs())

140 +def test_implicit_converted_models(model_arch: str,

141 + monkeypatch: pytest.MonkeyPatch):

142 + can_initialize(model_arch, monkeypatch, AUTO_EXAMPLE_MODELS)

tests/models/test_transformers.py

+35

@@ -138,3 +138,38 @@ def test_quantization(

138 name_0="transformers",

139 name_1="vllm",

140 )

141 +

142 +

143 +@pytest.mark.parametrize(

144 + "model",

145 + ["jason9693/Qwen2.5-1.5B-apeach"],

146 +)

147 +@pytest.mark.parametrize("dtype", ["half"])

148 +def test_classify(

149 + hf_runner,

150 + vllm_runner,

151 + example_prompts,

152 + model: str,

153 + dtype: str,

154 + monkeypatch,

155 +) -> None:

156 + import torch

157 + from transformers import AutoModelForSequenceClassification

158 +

159 + with vllm_runner(model,

160 + max_model_len=512,

161 + dtype=dtype,

162 + model_impl="transformers") as vllm_model:

163 + vllm_outputs = vllm_model.classify(example_prompts)

164 +

165 + with hf_runner(model,

166 + dtype=dtype,

167 + auto_cls=AutoModelForSequenceClassification) as hf_model:

168 + hf_outputs = hf_model.classify(example_prompts)

169 +

170 + for hf_output, vllm_output in zip(hf_outputs, vllm_outputs):

171 + hf_output = torch.tensor(hf_output)

172 + vllm_output = torch.tensor(vllm_output)

173 +

174 + assert torch.allclose(hf_output, vllm_output,

175 + 1e-3 if dtype == "float" else 1e-2)

vllm/config.py

+25 -21

@@ -551,7 +551,7 @@ class ModelConfig:

551 # For pooling models, self.task is used to indicate the

552 # user-selected task

553 if self.task == "score":

554 - if self.registry.is_cross_encoder_model(self.architectures):

554 + if self._is_classify_task(self.architectures):

555 self.task = "classify"

556 else:

557 self.task = "embed"

@@ -806,6 +806,12 @@ class ModelConfig:

806 f"one of {get_args(TokenizerMode)}.")

807 self.tokenizer_mode = tokenizer_mode

808

809 + def _is_classify_task(self, architectures: list[str]):

810 + for arch in architectures:

811 + if arch.endswith("ForSequenceClassification"):

812 + return True

813 + return self.registry.is_cross_encoder_model(architectures)

814 +

815 def _get_preferred_pooling_task(

816 self,

817 architectures: list[str],

@@ -813,14 +819,11 @@ class ModelConfig:

819 model_id = self.model

820 if get_pooling_config(model_id, self.revision):

821 return "embed"

816 - if self.registry.is_cross_encoder_model(architectures):

817 - return "classify"

822 if self.registry.is_transcription_model(architectures):

823 return "transcription"

824

825 suffix_to_preferred_task: list[tuple[str, _ResolvedTask]] = [

826 # Other models follow this pattern

823 - ("ForSequenceClassification", "classify"),

827 ("EmbeddingModel", "embed"),

828 ("RewardModel", "reward"),

829 ]

@@ -878,11 +881,14 @@ class ModelConfig:

881 self,

882 task_option: TaskOption,

883 ) -> dict[RunnerType, list[_ResolvedTask]]:

881 - return {

882 - "generate": self._get_supported_generation_tasks(task_option),

883 - "pooling": self._get_supported_pooling_tasks(task_option),

884 - "draft": ["draft"]

885 - }

884 + if self._is_classify_task(self.architectures):

885 + return {"generate": [], "pooling": ["classify"], "draft": []}

886 + else:

887 + return {

888 + "generate": self._get_supported_generation_tasks(task_option),

889 + "pooling": self._get_supported_pooling_tasks(task_option),

890 + "draft": ["draft"]

891 + }

892

893 def _get_supported_runner_types(

894 self,

@@ -925,12 +931,16 @@ class ModelConfig:

931 f"Available tasks for runner={task_runner!r}: "

932 f"{supported_tasks[task_runner]}")

933

934 + if "classify" in supported_tasks.get("pooling", []):

935 + # When multiple pooling tasks are present, default to

936 + # pooling (eg cross-encoder) for non-standard architectures.

937 + return "pooling"

938 +

939 suffix_to_preferred_runner: list[tuple[str, RunnerType]] = [

940 ("ForCausalLM", "generate"),

941 ("ForConditionalGeneration", "generate"),

942 ("ChatModel", "generate"),

943 ("LMHeadModel", "generate"),

933 - ("ForSequenceClassification", "pooling"),

944 ("EmbeddingModel", "pooling"),

945 ("RewardModel", "pooling"),

946 ]

@@ -940,10 +950,6 @@ class ModelConfig:

950 if arch.endswith(suffix) and pref_runner in supported_runner_types:

951 return pref_runner

952

943 - if "classify" in supported_tasks.get("pooling", []):

944 - # When multiple pooling tasks are present, default to

945 - # pooling (eg cross-encoder) for non-standard architectures.

946 - return "pooling"

953 if "generate" in supported_runner_types:

954 return "generate"

955 if "pooling" in supported_runner_types:

@@ -1525,7 +1531,7 @@ class ModelConfig:

1531

1532 @property

1533 def is_matryoshka(self) -> bool:

1528 - return (hasattr(self.hf_config, "matryoshka_dimensions")

1534 + return (bool(getattr(self.hf_config, "matryoshka_dimensions", None))

1535 or getattr(self.hf_config, "is_matryoshka", False))

1536

1537 @property

@@ -1539,13 +1545,11 @@ class ModelConfig:

1545 return getattr(self.hf_config, "use_pad_token", True)

1546

1547 def get_and_verify_max_len(self, max_model_len: int):

1542 - # For pooling models, the tokenizer's `model_max_length` is often a

1543 - # reliable source for the maximum sequence length. However, for

1544 - # generative models, this can be incorrect and unduly limit the

1545 - # context window (e.g., DeepSeek-R1). Therefore, we only consider

1546 - # tokenizer_config for pooling models.

1548 + # Consider max_model_len in tokenizer_config only when

1549 + # pooling models use absolute position_embedding.

1550 tokenizer_config = None

1548 - if self.runner_type == "pooling":

1551 + if (self.runner_type == "pooling" and getattr(

1552 + self.hf_config, "position_embedding_type", "") == "absolute"):

1553 tokenizer_config = try_get_tokenizer_config(

1554 self.tokenizer,

1555 trust_remote_code=self.trust_remote_code,

vllm/model_executor/model_loader/utils.py

+26 -4

@@ -22,7 +22,8 @@ from vllm.model_executor.layers.quantization.base_config import (

22 QuantizationConfig, QuantizeMethodBase)

23 from vllm.model_executor.models import ModelRegistry

24 from vllm.model_executor.models.adapters import (as_embedding_model,

25 - as_reward_model)

25 + as_reward_model,

26 + as_seq_cls_model)

27 from vllm.model_executor.models.interfaces import SupportsQuant

28 from vllm.utils import is_pin_memory_available

@@ -238,9 +239,29 @@ def get_model_architecture(

239 vllm_supported_archs = ModelRegistry.get_supported_archs()

240 vllm_not_supported = not any(arch in vllm_supported_archs

241 for arch in architectures)

242 +

243 + if vllm_not_supported:

244 + # try automatic conversion in adapters.py

245 + for arch in architectures:

246 + if not arch.endswith("ForSequenceClassification"):

247 + continue

248 +

249 + assert model_config.task == "classify"

250 + causal_lm_arch = arch.replace("ForSequenceClassification",

251 + "ForCausalLM")

252 + causal_lm_arch_vllm_supported = (causal_lm_arch

253 + in vllm_supported_archs)

254 + if not causal_lm_arch_vllm_supported:

255 + continue

256 +

257 + architectures = [causal_lm_arch]

258 + vllm_not_supported = False

259 + break

260 +

261 if (model_config.model_impl == ModelImpl.TRANSFORMERS or

262 model_config.model_impl != ModelImpl.VLLM and vllm_not_supported):

263 architectures = resolve_transformers_arch(model_config, architectures)

264 + logger.debug_once("Resolve transformers arch %s", str(architectures))

265 elif (model_config.quantization is not None

266 and model_config.quantization not in mixtral_supported

267 and "MixtralForCausalLM" in architectures):

@@ -248,12 +269,13 @@ def get_model_architecture(

269

270 model_cls, arch = ModelRegistry.resolve_model_cls(architectures)

271 if model_config.task == "embed":

272 + logger.debug_once("Automatic conversion using `as_embedding_model`.")

273 model_cls = as_embedding_model(model_cls)

274 elif model_config.task == "classify":

253 - # Cannot automatically run as_seq_cls_model,

254 - # otherwise it will cause a circular reference on is_cross_encoder_model

255 - pass

275 + logger.debug_once("Automatic conversion using `as_seq_cls_model`.")

276 + model_cls = as_seq_cls_model(model_cls)

277 elif model_config.task == "reward":

278 + logger.debug_once("Automatic conversion using `as_reward_model`.")

279 model_cls = as_reward_model(model_cls)

280

281 return model_cls, arch

vllm/model_executor/models/adapters.py

+9 -6

@@ -331,13 +331,13 @@ def load_weights_using_from_2_way_softmax(

331

332 false_id = tokenizer.convert_tokens_to_ids(tokens[0])

333 true_id = tokenizer.convert_tokens_to_ids(tokens[1])

334 - weight = model.lm_head.weight.data[[true_id]].to(

334 + score_weight = model.lm_head.weight.data[[true_id]].to(

335 torch.float32) - model.lm_head.weight.data[[false_id]].to(

336 torch.float32)

337

338 param = model.score.weight

339 weight_loader = getattr(param, "weight_loader", default_weight_loader)

340 - weight_loader(param, weight)

340 + weight_loader(param, score_weight)

341

342 del model.lm_head

343 loaded_weights.add("score.weight")

@@ -350,6 +350,8 @@ def load_weights_no_post_processing(model,

350 torch.Tensor]]):

351 from vllm.model_executor.layers.vocab_parallel_embedding import (

352 ParallelLMHead)

353 + from vllm.model_executor.model_loader.weight_utils import (

354 + default_weight_loader)

355 from vllm.model_executor.models.utils import AutoWeightsLoader

356

357 model_config = model.vllm_config.model_config

@@ -357,8 +359,6 @@ def load_weights_no_post_processing(model,

359 tokens = cast(list[int], tokens)

360 assert len(tokens) > 0

361

360 - device = model.score.weight.device

361 -

362 if model.config.tie_word_embeddings:

363 model.lm_head = model.model.embed_tokens

364 else:

@@ -376,8 +376,11 @@ def load_weights_no_post_processing(model,

376 trust_remote_code=model_config.trust_remote_code)

377

378 token_ids = [tokenizer.convert_tokens_to_ids(t) for t in tokens]

379 - score_weight = model.lm_head.weight.data[token_ids].to(device)

380 - model.score.weight.data.copy_(score_weight)

379 + score_weight = model.lm_head.weight.data[token_ids]

380 +

381 + param = model.score.weight

382 + weight_loader = getattr(param, "weight_loader", default_weight_loader)

383 + weight_loader(param, score_weight)

384

385 del model.lm_head

386 loaded_weights.add("score.weight")

vllm/model_executor/models/gemma.py

-4

@@ -43,7 +43,6 @@ from vllm.model_executor.model_loader.weight_utils import default_weight_loader

43 from vllm.model_executor.sampling_metadata import SamplingMetadata

44 from vllm.sequence import IntermediateTensors

46 -from .adapters import as_seq_cls_model

46 from .interfaces import SupportsLoRA, SupportsPP

47 from .utils import (AutoWeightsLoader, is_pp_missing_parameter,

48 make_empty_intermediate_tensors_factory, make_layers,

@@ -426,6 +425,3 @@ class GemmaForCausalLM(nn.Module, SupportsLoRA, SupportsPP):

425 if self.config.tie_word_embeddings else None),

426 )

427 return loader.load_weights(weights)

429 -

430 -

431 -GemmaForSequenceClassification = as_seq_cls_model(GemmaForCausalLM)

vllm/model_executor/models/llama.py

-4

@@ -49,7 +49,6 @@ from vllm.model_executor.model_loader.weight_utils import (

49 from vllm.model_executor.sampling_metadata import SamplingMetadata

50 from vllm.sequence import IntermediateTensors

52 -from .adapters import as_seq_cls_model

52 from .interfaces import SupportsLoRA, SupportsPP

53 from .utils import (AutoWeightsLoader, PPMissingLayer, extract_layer_index,

54 is_pp_missing_parameter,

@@ -646,6 +645,3 @@ class LlamaForCausalLM(nn.Module, SupportsLoRA, SupportsPP):

645 name = name.replace(item, mapping[item])

646

647 return name, loaded_weight

649 -

650 -

651 -LlamaForSequenceClassification = as_seq_cls_model(LlamaForCausalLM)

vllm/model_executor/models/qwen2.py

-4

@@ -50,7 +50,6 @@ from vllm.model_executor.model_loader.weight_utils import (

50 from vllm.model_executor.sampling_metadata import SamplingMetadata

51 from vllm.sequence import IntermediateTensors

53 -from .adapters import as_seq_cls_model

53 from .interfaces import SupportsLoRA, SupportsPP

54 from .utils import (AutoWeightsLoader, PPMissingLayer, extract_layer_index,

55 is_pp_missing_parameter,

@@ -496,6 +495,3 @@ class Qwen2ForCausalLM(nn.Module, SupportsLoRA, SupportsPP):

495 if self.config.tie_word_embeddings else None),

496 )

497 return loader.load_weights(weights)

499 -

500 -

501 -Qwen2ForSequenceClassification = as_seq_cls_model(Qwen2ForCausalLM)

vllm/model_executor/models/qwen3.py

-4

@@ -44,7 +44,6 @@ from vllm.model_executor.layers.vocab_parallel_embedding import ParallelLMHead

44 from vllm.model_executor.sampling_metadata import SamplingMetadata

45 from vllm.sequence import IntermediateTensors

47 -from .adapters import as_seq_cls_model

47 from .interfaces import SupportsLoRA, SupportsPP

48 from .qwen2 import Qwen2MLP as Qwen3MLP

49 from .qwen2 import Qwen2Model

@@ -320,6 +319,3 @@ class Qwen3ForCausalLM(nn.Module, SupportsLoRA, SupportsPP):

319 if self.config.tie_word_embeddings else None),

320 )

321 return loader.load_weights(weights)

323 -

324 -

325 -Qwen3ForSequenceClassification = as_seq_cls_model(Qwen3ForCausalLM)

vllm/model_executor/models/registry.py

+29 -8

@@ -12,7 +12,7 @@ import sys

12 import tempfile

13 from abc import ABC, abstractmethod

14 from collections.abc import Set

15 -from dataclasses import dataclass, field

15 +from dataclasses import asdict, dataclass, field

16 from functools import lru_cache

17 from typing import Callable, Optional, TypeVar, Union

@@ -181,10 +181,6 @@ _CROSS_ENCODER_MODELS = {

181 "ModernBertForSequenceClassification": ("modernbert",

182 "ModernBertForSequenceClassification"),

183 # [Auto-converted (see adapters.py)]

184 - "GemmaForSequenceClassification": ("gemma", "GemmaForSequenceClassification"), # noqa: E501

185 - "Qwen2ForSequenceClassification": ("qwen2", "Qwen2ForSequenceClassification"), # noqa: E501

186 - "Qwen3ForSequenceClassification": ("qwen3", "Qwen3ForSequenceClassification"), # noqa: E501

187 - "LlamaForSequenceClassification": ("llama", "LlamaForSequenceClassification"), # noqa: E501

184 "JinaVLForRanking": ("jina_vl", "JinaVLForSequenceClassification"), # noqa: E501,

185 }

186

@@ -462,10 +458,26 @@ class _ModelRegistry:

458 return _try_load_model_cls(model_arch, self.models[model_arch])

459

460 def _try_inspect_model_cls(self, model_arch: str) -> Optional[_ModelInfo]:

465 - if model_arch not in self.models:

466 - return None

461 + if model_arch in self.models:

462 + return _try_inspect_model_cls(model_arch, self.models[model_arch])

463 +

464 + if model_arch.endswith("ForSequenceClassification"):

465 + causal_lm_arch = model_arch.replace("ForSequenceClassification",

466 + "ForCausalLM")

467 + if causal_lm_arch not in self.models:

468 + return None

469 +

470 + info = _try_inspect_model_cls(causal_lm_arch,

471 + self.models[causal_lm_arch])

472

468 - return _try_inspect_model_cls(model_arch, self.models[model_arch])

473 + info = _ModelInfo(**dict(

474 + asdict(info), **{

475 + "architecture": model_arch,

476 + "supports_cross_encoding": True

477 + }))

478 + return info

479 +

480 + return None

481

482 def _normalize_archs(

483 self,

@@ -480,6 +492,15 @@ class _ModelRegistry:

492 normalized_arch = list(

493 filter(lambda model: model in self.models, architectures))

494

495 + # try automatic conversion in adapters.py

496 + for arch in architectures:

497 + if not arch.endswith("ForSequenceClassification"):

498 + continue

499 + causal_lm_arch = arch.replace("ForSequenceClassification",

500 + "ForCausalLM")

501 + if causal_lm_arch in self.models:

502 + normalized_arch.append(arch)

503 +

504 # make sure Transformers backend is put at the last as a fallback

505 if len(normalized_arch) != len(architectures):

506 normalized_arch.append("TransformersForCausalLM")