[Model] Add RobertaForTokenClassification / XLMRobertaForTokenClassification (#47991)
Signed-off-by: krishy91 <crgkc.r@gmail.com>
Gopala-Krishna Char committed
Jul 15, 2026 at 16:30 UTC
5810e884f11875b46fbae0871754d855d8891105
5 files changed
+118
-1
docs/models/pooling_models/token_classify.md
+2
@@ -47,6 +47,8 @@ The BAAI/bge-m3 model leverages token classification for sparse retrieval. For m
47
| `ModernBertForTokenClassification` | ModernBERT-based | `disham993/electrical-ner-ModernBERT-base` | | |
48
| `OpenAIPrivacyFilterForTokenClassification` | gpt-oss-based encoder | `openai/privacy-filter` | | |
49
| `Qwen3ForTokenClassification`<sup>C</sup> | Qwen3-based | `bd2lcco/Qwen3-0.6B-finetuned` | | |
50
+| `RobertaForTokenClassification` | RoBERTa-based | `Jean-Baptiste/roberta-large-ner-english` | | |
51
+| `XLMRobertaForTokenClassification` | XLM-RoBERTa-based | `Davlan/xlm-roberta-base-ner-hrl` | | |
52
| `*Model`<sup>C</sup>, `*ForCausalLM`<sup>C</sup>, etc. | Generative models | N/A | \* | \* |
53
54
<sup>C</sup> Automatically converted into a classification model via `--convert classify`. ([details](./README.md#model-conversion))
tests/models/language/pooling/test_token_classification.py
+40
@@ -117,6 +117,46 @@ def test_modernbert_models(
117
torch.testing.assert_close(hf_output, vllm_output, atol=3.2e-2, rtol=1e-3)
118
119
120
+@pytest.mark.parametrize("model", ["Davlan/xlm-roberta-base-ner-hrl"])
121
+@pytest.mark.parametrize("dtype", ["float"])
122
+@torch.inference_mode
123
+def test_xlm_roberta_models(
124
+ hf_runner,
125
+ vllm_runner,
126
+ example_prompts,
127
+ model: str,
128
+ dtype: str,
129
+) -> None:
130
+ with vllm_runner(model, max_model_len=None, dtype=dtype) as vllm_model:
131
+ vllm_outputs = vllm_model.token_classify(example_prompts)
132
+
133
+ # Use eager attention on ROCm to avoid HF Transformers flash attention
134
+ # accuracy issues: https://github.com/vllm-project/vllm/issues/30167
135
+ hf_model_kwargs = {}
136
+ if current_platform.is_rocm():
137
+ hf_model_kwargs["attn_implementation"] = "eager"
138
+
139
+ with hf_runner(
140
+ model,
141
+ dtype=dtype,
142
+ auto_cls=AutoModelForTokenClassification,
143
+ model_kwargs=hf_model_kwargs,
144
+ ) as hf_model:
145
+ tokenizer = hf_model.tokenizer
146
+ hf_outputs = []
147
+ for prompt in example_prompts:
148
+ inputs = tokenizer([prompt], return_tensors="pt")
149
+ inputs = hf_model.wrap_device(inputs)
150
+ output = hf_model.model(**inputs)
151
+ hf_outputs.append(softmax(output.logits[0]))
152
+
153
+ # check logits difference
154
+ for hf_output, vllm_output in zip(hf_outputs, vllm_outputs):
155
+ hf_output = hf_output.detach().clone().cpu().float()
156
+ vllm_output = vllm_output.detach().clone().cpu().float()
157
+ torch.testing.assert_close(hf_output, vllm_output, atol=3.2e-2, rtol=1e-3)
158
+
159
+
160
PRIVACY_FILTER_PROMPTS = [
161
"My name is Harry Potter.",
162
"Email me at harry.potter@hogwarts.edu.",
tests/models/registry.py
+6
@@ -713,6 +713,12 @@ _TOKEN_CLASSIFICATION_EXAMPLE_MODELS = {
713
"openai/privacy-filter",
714
min_transformers_version="5.6.0.dev0",
715
),
716
+ "RobertaForTokenClassification": _HfExamplesInfo(
717
+ "Jean-Baptiste/roberta-large-ner-english"
718
+ ),
719
+ "XLMRobertaForTokenClassification": _HfExamplesInfo(
720
+ "Davlan/xlm-roberta-base-ner-hrl"
721
+ ),
722
}
723
724
_SEQUENCE_CLASSIFICATION_EXAMPLE_MODELS = {
vllm/model_executor/models/registry.py
+5
@@ -297,6 +297,11 @@ _TOKEN_CLASSIFICATION_MODELS = {
297
"qwen3_asr_forced_aligner",
298
"Qwen3ASRForcedAlignerForTokenClassification",
299
),
300
+ "RobertaForTokenClassification": ("roberta", "RobertaForTokenClassification"),
301
+ "XLMRobertaForTokenClassification": (
302
+ "roberta",
303
+ "RobertaForTokenClassification",
304
+ ),
305
}
306
307
_SEQUENCE_CLASSIFICATION_MODELS = {
vllm/model_executor/models/roberta.py
+65
-1
@@ -42,7 +42,7 @@ from vllm.sequence import IntermediateTensors
42
43
from .bert_with_rope import BertWithRope, JinaRobertaModel
44
from .interfaces import SupportsCrossEncoding
45
-from .interfaces_base import default_pooling_type
45
+from .interfaces_base import attn_type, default_pooling_type
46
47
48
class RobertaEmbedding(nn.Module):
@@ -332,3 +332,67 @@ class RobertaForSequenceClassification(nn.Module, SupportsCrossEncoding):
332
inputs_embeds=inputs_embeds,
333
intermediate_tensors=intermediate_tensors,
334
)
335
+
336
+
337
+@attn_type("encoder_only")
338
+@default_pooling_type(tok_pooling_type="ALL")
339
+class RobertaForTokenClassification(nn.Module):
340
+ """A model that uses Roberta to provide token classification.
341
+
342
+ Mirrors BertForTokenClassification, swapping in RobertaEmbedding for the
343
+ RoBERTa/XLM-RoBERTa position-embedding offset and weight layout.
344
+
345
+ Also registered as XLMRobertaForTokenClassification since XLM-RoBERTa
346
+ checkpoints share RoBERTa's architecture and `roberta.*` weight prefix.
347
+ """
348
+
349
+ is_pooling_model = True
350
+
351
+ def __init__(self, *, vllm_config: VllmConfig, prefix: str = ""):
352
+ super().__init__()
353
+ config = vllm_config.model_config.hf_config
354
+ self.head_dtype = vllm_config.model_config.head_dtype
355
+ self.num_labels = config.num_labels
356
+ self.roberta = BertModel(
357
+ vllm_config=vllm_config,
358
+ prefix=maybe_prefix(prefix, "roberta"),
359
+ embedding_class=RobertaEmbedding,
360
+ )
361
+ self.classifier = nn.Linear(
362
+ config.hidden_size, config.num_labels, dtype=self.head_dtype
363
+ )
364
+
365
+ pooler_config = vllm_config.model_config.pooler_config
366
+ assert pooler_config is not None
367
+
368
+ self.pooler = pooler_for_token_classify(pooler_config)
369
+
370
+ def embed_input_ids(self, input_ids: torch.Tensor) -> torch.Tensor:
371
+ return self.roberta.embed_input_ids(input_ids)
372
+
373
+ def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]):
374
+ loader = AutoWeightsLoader(self)
375
+ return loader.load_weights(weights)
376
+
377
+ def forward(
378
+ self,
379
+ input_ids: torch.Tensor | None,
380
+ positions: torch.Tensor,
381
+ intermediate_tensors: IntermediateTensors | None = None,
382
+ inputs_embeds: torch.Tensor | None = None,
383
+ token_type_ids: torch.Tensor | None = None,
384
+ ) -> torch.Tensor:
385
+ if token_type_ids is not None:
386
+ assert self.roberta.config.vocab_size < (1 << TOKEN_TYPE_SHIFT)
387
+ assert input_ids is not None
388
+ _encode_token_type_ids(input_ids, token_type_ids)
389
+
390
+ hidden_states = self.roberta(
391
+ input_ids=input_ids,
392
+ positions=positions,
393
+ inputs_embeds=inputs_embeds,
394
+ intermediate_tensors=intermediate_tensors,
395
+ )
396
+
397
+ hidden_states = hidden_states.to(self.head_dtype)
398
+ return self.classifier(hidden_states)