Skip to content

[Feature]: Add classification Task with AutoModelForSequenceClassification and BertForSequenceClassification #10939

@yamamotolotation

Description

@yamamotolotation

🚀 The feature, motivation and pitch

Latest Version: v0.6.4.post1

Summary:
I would like to request support for using BERT and DistilBERT models for classification tasks. As of v0.6.4.post1, BertForSequenceClassification is only available for Sentence Pair Scoring tasks via the Score API.

Detail:
With a recent Pull Request, #9704 Qwen2ForSequenceClassification became available for text classification tasks. This is a fantastic improvement!

However, in text multi-classification tasks, models other than Qwen2 are often preferred due to the following reasons:

  • Limited multilingual support in Qwen2.
  • A larger number of parameters in Qwen2, leading to longer inference times.

BERT and DistilBERT, which have been extensively researched and fine-tuned for a wide range of languages, often achieve better classification accuracy and faster inference times. These models are frequently integrated into applications as local LLMs, making them an excellent fit for vLLM's role as a high-speed inference accelerator.

Currently, if we try to execute code for text classification using BERT or DistilBERT models in the same manner as the sample code for Qwen2, the following error is raised:

from vllm import LLM, SamplingParams

model = LLM(
    model="cl-tohoku/bert-base-japanese-v3.pt",
    tensor_parallel_size=tensor_parallel_size,
)
outputs = model.generate(
    prompts=list(texts),
    sampling_params=sampling_params,
)
ValueError: Model architectures ['BertForSequenceClassification'] are not supported for now. Supported architectures: dict_keys(['AquilaModel', 'AquilaForCausalLM', 'ArcticForCausalLM', 'BaiChuanForCausalLM', 'BaichuanForCausalLM', 'BloomForCausalLM', 'CohereForCausalLM', 'DbrxForCausalLM', 'DeciLMForCausalLM', 'DeepseekForCausalLM', 'DeepseekV2ForCausalLM', 'ExaoneForCausalLM', 'FalconForCausalLM', 'GemmaForCausalLM', 'Gemma2ForCausalLM', 'GPT2LMHeadModel', 'GPTBigCodeForCausalLM', 'GPTJForCausalLM', 'GPTNeoXForCausalLM', 'GraniteForCausalLM', 'GraniteMoeForCausalLM', 'InternLMForCausalLM', 'InternLM2ForCausalLM', 'InternLM2VEForCausalLM', 'JAISLMHeadModel', 'JambaForCausalLM', 'LlamaForCausalLM', 'LLaMAForCausalLM', 'MambaForCausalLM', 'FalconMambaForCausalLM', 'MiniCPMForCausalLM', 'MiniCPM3ForCausalLM', 'MistralForCausalLM', 'MixtralForCausalLM', 'QuantMixtralForCausalLM', 'MptForCausalLM', 'MPTForCausalLM', 'NemotronForCausalLM', 'OlmoForCausalLM', 'OlmoeForCausalLM', 'OPTForCausalLM', 'OrionForCausalLM', 'PersimmonForCausalLM', 'PhiForCausalLM', 'Phi3ForCausalLM', 'Phi3SmallForCausalLM', 'PhiMoEForCausalLM', 'Qwen2ForCausalLM', 'Qwen2MoeForCausalLM', 'RWForCausalLM', 'StableLMEpochForCausalLM', 'StableLmForCausalLM', 'Starcoder2ForCausalLM', 'SolarForCausalLM', 'XverseForCausalLM', 'BartModel', 'BartForConditionalGeneration', 'Florence2ForConditionalGeneration', 'BertModel', 'RobertaModel', 'XLMRobertaModel', 'Gemma2Model', 'LlamaModel', 'MistralModel', 'Qwen2Model', 'Qwen2ForRewardModel', 'Qwen2ForSequenceClassification', 'LlavaNextForConditionalGeneration', 'Phi3VForCausalLM', 'Qwen2VLForConditionalGeneration', 'Blip2ForConditionalGeneration', 'ChameleonForConditionalGeneration', 'ChatGLMModel', 'ChatGLMForConditionalGeneration', 'FuyuForCausalLM', 'H2OVLChatModel', 'InternVLChatModel', 'Idefics3ForConditionalGeneration', 'LlavaForConditionalGeneration', 'LlavaNextVideoForConditionalGeneration', 'LlavaOnevisionForConditionalGeneration', 'MiniCPMV', 'MolmoForCausalLM', 'NVLM_D', 'PaliGemmaForConditionalGeneration', 'PixtralForConditionalGeneration', 'QWenLMHeadModel', 'Qwen2AudioForConditionalGeneration', 'UltravoxModel', 'MllamaForConditionalGeneration', 'EAGLEModel', 'MedusaModel', 'MLPSpeculatorPreTrainedModel'])

Adding support for AutoModelForSequenceClassification, BertForSequenceClassification, and similar models would greatly enhance the utility of vLLM for a broader range of text classification tasks.

Alternatives

No response

Additional context

No response

Before submitting a new issue...

  • Make sure you already searched for relevant issues, and asked the chatbot living at the bottom right corner of the documentation page, which can answer lots of frequently asked questions.

Metadata

Metadata

Assignees

No one assigned

    Labels

    feature requestNew feature or requeststaleOver 90 days of inactivity

    Type

    No type

    Projects

    No projects

    Milestone

    No milestone

    Relationships

    None yet

    Development

    No branches or pull requests

    Issue actions