registry.py 43.6 KB
Newer Older
1
# SPDX-License-Identifier: Apache-2.0
2
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3

4
from collections.abc import Mapping, Set
5
from dataclasses import dataclass, field
6
from typing import Any, Literal, Optional
7
8

import pytest
9
import torch
10
11
from packaging.version import Version
from transformers import __version__ as TRANSFORMERS_VERSION
12

13
from vllm.config import ModelDType, TokenizerMode
14

15
16
17
18
19
20
21
22
23
24
25
26

@dataclass(frozen=True)
class _HfExamplesInfo:
    default: str
    """The default model to use for testing this architecture."""

    extras: Mapping[str, str] = field(default_factory=dict)
    """Extra models to use for testing this architecture."""

    tokenizer: Optional[str] = None
    """Set the tokenizer to load for this architecture."""

27
    tokenizer_mode: TokenizerMode = "auto"
28
29
30
31
32
33
34
35
    """Set the tokenizer type for this architecture."""

    speculative_model: Optional[str] = None
    """
    The default model to use for testing this architecture, which is only used
    for speculative decoding.
    """

36
37
38
39
40
    min_transformers_version: Optional[str] = None
    """
    The minimum version of HF Transformers that is required to run this model.
    """

41
42
43
44
45
46
47
48
49
50
    max_transformers_version: Optional[str] = None
    """
    The maximum version of HF Transformers that this model runs on.
    """

    transformers_version_reason: Optional[str] = None
    """
    The reason for the minimum/maximum version requirement.
    """

51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
    skip_tokenizer_init: bool = False
    """
    If true, skip initialization of tokenizer and detokenizer. 
    """

    dtype: ModelDType = "auto"
    """
    The data type for the model weights and activations.
    """

    enforce_eager: bool = False
    """
    Whether to enforce eager execution. If True, we will
    disable CUDA graph and always execute the model in eager mode.
    If False, we will use CUDA graph and eager execution in hybrid.
    """

68
69
70
71
72
73
74
75
76
77
78
    is_available_online: bool = True
    """
    Set this to ``False`` if the name of this architecture no longer exists on
    the HF repo. To maintain backwards compatibility, we have not removed them
    from the main model registry, so without this flag the registry tests will
    fail.
    """

    trust_remote_code: bool = False
    """The ``trust_remote_code`` level required to load the model."""

79
80
81
    v0_only: bool = False
    """The model is only available with the vLLM V0 engine."""

82
83
84
    hf_overrides: dict[str, Any] = field(default_factory=dict)
    """The ``hf_overrides`` required to load the model."""

85
86
87
88
89
90
    max_model_len: Optional[int] = None
    """
    The maximum model length to use for this model. Some models default to a
    length that is too large to fit into memory in CI.
    """

91
92
93
94
95
96
    revision: Optional[str] = None
    """
    The specific revision (commit hash, tag, or branch) to use for the model.
    If not specified, the default revision will be used.
    """

97
98
99
    max_num_seqs: Optional[int] = None
    """Maximum number of sequences to be processed in a single iteration."""

100
101
102
103
104
105
    use_original_num_layers: bool = False
    """
    If True, use the original number of layers from the model config 
    instead of minimal layers for testing.
    """

106
107
108
    def check_transformers_version(
        self,
        *,
109
        on_fail: Literal["error", "skip", "return"],
110
111
        check_min_version: bool = True,
        check_max_version: bool = True,
112
    ) -> Optional[str]:
113
114
115
116
        """
        If the installed transformers version does not meet the requirements,
        perform the given action.
        """
117
118
        if (self.min_transformers_version is None
                and self.max_transformers_version is None):
119
            return None
120
121

        current_version = TRANSFORMERS_VERSION
122
        cur_base_version = Version(current_version).base_version
123
124
125
        min_version = self.min_transformers_version
        max_version = self.max_transformers_version
        msg = f"`transformers=={current_version}` installed, but `transformers"
126
127
        # Only check the base version for the min/max version, otherwise preview
        # models cannot be run because `x.yy.0.dev0`<`x.yy.0`
128
129
        if (check_min_version and min_version
                and Version(cur_base_version) < Version(min_version)):
130
            msg += f">={min_version}` is required to run this model."
131
132
        elif (check_max_version and max_version
              and Version(cur_base_version) > Version(max_version)):
133
134
            msg += f"<={max_version}` is required to run this model."
        else:
135
            return None
136

137
138
139
140
141
        if self.transformers_version_reason:
            msg += f" Reason: {self.transformers_version_reason}"

        if on_fail == "error":
            raise RuntimeError(msg)
142
        elif on_fail == "skip":
143
            pytest.skip(msg)
144

145
146
        return msg

147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
    def check_available_online(
        self,
        *,
        on_fail: Literal["error", "skip"],
    ) -> None:
        """
        If the model is not available online, perform the given action.
        """
        if not self.is_available_online:
            msg = "Model is not available online"

            if on_fail == "error":
                raise RuntimeError(msg)
            else:
                pytest.skip(msg)

163
164
165
166

# yapf: disable
_TEXT_GENERATION_EXAMPLE_MODELS = {
    # [Decoder-only]
167
    "ApertusForCausalLM": _HfExamplesInfo("swiss-ai/Apertus-8B-2509",
168
169
                                          min_transformers_version="4.56.0",
                                          trust_remote_code=True),
170
171
172
173
    "AquilaModel": _HfExamplesInfo("BAAI/AquilaChat-7B",
                                   trust_remote_code=True),
    "AquilaForCausalLM": _HfExamplesInfo("BAAI/AquilaChat2-7B",
                                         trust_remote_code=True),
174
    "ArceeForCausalLM": _HfExamplesInfo("arcee-ai/AFM-4.5B-Base"),
175
176
177
178
179
180
    "ArcticForCausalLM": _HfExamplesInfo("Snowflake/snowflake-arctic-instruct",
                                         trust_remote_code=True),
    "BaiChuanForCausalLM": _HfExamplesInfo("baichuan-inc/Baichuan-7B",
                                         trust_remote_code=True),
    "BaichuanForCausalLM": _HfExamplesInfo("baichuan-inc/Baichuan2-7B-chat",
                                         trust_remote_code=True),
181
182
    "BailingMoeForCausalLM": _HfExamplesInfo("inclusionAI/Ling-lite-1.5",
                                         trust_remote_code=True),
ant-yy's avatar
ant-yy committed
183
184
    "BailingMoeV2ForCausalLM": _HfExamplesInfo("inclusionAI/Ling-mini-2.0",
                                         trust_remote_code=True),
185
    "BambaForCausalLM": _HfExamplesInfo("ibm-ai-platform/Bamba-9B-v1",
186
                                        min_transformers_version="4.55.3",
187
                                        extras={"tiny": "hmellor/tiny-random-BambaForCausalLM"}),  # noqa: E501
188
189
    "BloomForCausalLM": _HfExamplesInfo("bigscience/bloom-560m",
                                        {"1b": "bigscience/bloomz-1b1"}),
190
    "ChatGLMModel": _HfExamplesInfo("zai-org/chatglm3-6b",
191
                                    trust_remote_code=True,
192
                                    max_transformers_version="4.48"),
193
194
    "ChatGLMForConditionalGeneration": _HfExamplesInfo("thu-coai/ShieldLM-6B-chatglm3",  # noqa: E501
                                                       trust_remote_code=True),
195
196
    "CohereForCausalLM": _HfExamplesInfo("CohereForAI/c4ai-command-r-v01",
                                         trust_remote_code=True),
197
198
    "Cohere2ForCausalLM": _HfExamplesInfo("CohereForAI/c4ai-command-r7b-12-2024", # noqa: E501
                                         trust_remote_code=True),
199
200
201
    "CwmForCausalLM": _HfExamplesInfo("facebook/cwm", # noqa: E501
                                      trust_remote_code=True,
                                      is_available_online=False),
202
    "DbrxForCausalLM": _HfExamplesInfo("databricks/dbrx-instruct"),
203
    "DeciLMForCausalLM": _HfExamplesInfo("nvidia/Llama-3_3-Nemotron-Super-49B-v1", # noqa: E501
204
205
206
207
                                         trust_remote_code=True),
    "DeepseekForCausalLM": _HfExamplesInfo("deepseek-ai/deepseek-llm-7b-chat"),
    "DeepseekV2ForCausalLM": _HfExamplesInfo("deepseek-ai/DeepSeek-V2-Lite-Chat",  # noqa: E501
                                         trust_remote_code=True),
Robert Shaw's avatar
Robert Shaw committed
208
209
    "DeepseekV3ForCausalLM": _HfExamplesInfo("deepseek-ai/DeepSeek-V3",  # noqa: E501
                                         trust_remote_code=True),
210
    "DeepseekV32ForCausalLM": _HfExamplesInfo("deepseek-ai/DeepSeek-V3.2-Exp"),
211
    "Ernie4_5ForCausalLM": _HfExamplesInfo("baidu/ERNIE-4.5-0.3B-PT",
212
                                            min_transformers_version="4.54"),
213
    "Ernie4_5_MoeForCausalLM": _HfExamplesInfo("baidu/ERNIE-4.5-21B-A3B-PT",
214
                                               min_transformers_version="4.54"),
215
216
217
218
    "ExaoneForCausalLM": _HfExamplesInfo("LGAI-EXAONE/EXAONE-3.0-7.8B-Instruct",
                                         trust_remote_code=True),
    "Exaone4ForCausalLM": _HfExamplesInfo("LGAI-EXAONE/EXAONE-4.0-32B",
                                          min_transformers_version="4.54"),
219
    "Fairseq2LlamaForCausalLM": _HfExamplesInfo("mgleize/fairseq2-dummy-Llama-3.2-1B"),  # noqa: E501
220
    "FalconForCausalLM": _HfExamplesInfo("tiiuae/falcon-7b"),
221
    "FalconH1ForCausalLM":_HfExamplesInfo("tiiuae/Falcon-H1-0.5B-Base"),
222
    "GemmaForCausalLM": _HfExamplesInfo("google/gemma-1.1-2b-it"),
223
    "Gemma2ForCausalLM": _HfExamplesInfo("google/gemma-2-9b"),
224
    "Gemma3ForCausalLM": _HfExamplesInfo("google/gemma-3-1b-it"),
Nicolò Lucchesi's avatar
Nicolò Lucchesi committed
225
    "Gemma3nForCausalLM": _HfExamplesInfo("google/gemma-3n-E2B-it",
Robert Shaw's avatar
Robert Shaw committed
226
                                          min_transformers_version="4.53"),
227
228
    "GlmForCausalLM": _HfExamplesInfo("zai-org/glm-4-9b-chat-hf"),
    "Glm4ForCausalLM": _HfExamplesInfo("zai-org/GLM-4-9B-0414"),
229
230
    "Glm4MoeForCausalLM": _HfExamplesInfo("zai-org/GLM-4.5",
                                          min_transformers_version="4.54"),   # noqa: E501
231
232
233
    "GPT2LMHeadModel": _HfExamplesInfo("openai-community/gpt2",
                                       {"alias": "gpt2"}),
    "GPTBigCodeForCausalLM": _HfExamplesInfo("bigcode/starcoder",
234
                                             extras={"tiny": "bigcode/tiny_starcoder_py"},  # noqa: E501
235
236
                                             min_transformers_version="4.55.1",
                                             transformers_version_reason="HF model broken in 4.55.0"),  # noqa: E501
237
238
239
240
    "GPTJForCausalLM": _HfExamplesInfo("Milos/slovak-gpt-j-405M",
                                       {"6b": "EleutherAI/gpt-j-6b"}),
    "GPTNeoXForCausalLM": _HfExamplesInfo("EleutherAI/pythia-70m",
                                          {"1b": "EleutherAI/pythia-1.4b"}),
241
    "GptOssForCausalLM": _HfExamplesInfo("lmsys/gpt-oss-20b-bf16"),
242
243
    "GraniteForCausalLM": _HfExamplesInfo("ibm/PowerLM-3b"),
    "GraniteMoeForCausalLM": _HfExamplesInfo("ibm/PowerMoE-3b"),
244
245
    "GraniteMoeHybridForCausalLM": _HfExamplesInfo("ibm-granite/granite-4.0-tiny-preview", # noqa: E501
                                                   min_transformers_version="4.55.3"),
246
    "GraniteMoeSharedForCausalLM": _HfExamplesInfo("ibm-research/moe-7b-1b-active-shared-experts"),  # noqa: E501
Michael Goin's avatar
Michael Goin committed
247
248
    "Grok1ModelForCausalLM": _HfExamplesInfo("hpcai-tech/grok-1",
                                             trust_remote_code=True),
249
250
    "HunYuanMoEV1ForCausalLM": _HfExamplesInfo("tencent/Hunyuan-A13B-Instruct",
                                               trust_remote_code=True),
251
    # TODO: Remove is_available_online once their config.json is fixed
252
    "HunYuanDenseV1ForCausalLM":_HfExamplesInfo("tencent/Hunyuan-7B-Instruct-0124",
253
254
                                                trust_remote_code=True,
                                                is_available_online=False),
255
256
257
258
259
260
    "InternLMForCausalLM": _HfExamplesInfo("internlm/internlm-chat-7b",
                                           trust_remote_code=True),
    "InternLM2ForCausalLM": _HfExamplesInfo("internlm/internlm2-chat-7b",
                                            trust_remote_code=True),
    "InternLM2VEForCausalLM": _HfExamplesInfo("OpenGVLab/Mono-InternVL-2B",
                                              trust_remote_code=True),
261
262
    "InternLM3ForCausalLM": _HfExamplesInfo("internlm/internlm3-8b-instruct",
                                            trust_remote_code=True),
263
    "JAISLMHeadModel": _HfExamplesInfo("inceptionai/jais-13b-chat"),
264
    "JambaForCausalLM": _HfExamplesInfo("ai21labs/AI21-Jamba-1.5-Mini",
265
                                        min_transformers_version="4.55.3",
266
267
268
269
                                        extras={
                                            "tiny": "ai21labs/Jamba-tiny-dev",
                                            "random": "ai21labs/Jamba-tiny-random",  # noqa: E501
                                        }),
270
271
    "Lfm2ForCausalLM": _HfExamplesInfo("LiquidAI/LFM2-1.2B",
                                       min_transformers_version="4.54"),
272
273
    "LlamaForCausalLM": _HfExamplesInfo("meta-llama/Llama-3.2-1B-Instruct",
                                        extras={"guard": "meta-llama/Llama-Guard-3-1B",  # noqa: E501
274
275
                                                "hermes": "NousResearch/Hermes-3-Llama-3.1-8B", # noqa: E501
                                                "fp8": "RedHatAI/Meta-Llama-3.1-8B-Instruct-FP8"}),  # noqa: E501
276
277
    "LLaMAForCausalLM": _HfExamplesInfo("decapoda-research/llama-7b-hf",
                                        is_available_online=False),
278
279
    "Llama4ForCausalLM": _HfExamplesInfo("meta-llama/Llama-4-Scout-17B-16E-Instruct", # noqa: E501
                                         is_available_online=False),
XuruiYang's avatar
XuruiYang committed
280
281
    "LongcatFlashForCausalLM": _HfExamplesInfo
                ("meituan-longcat/LongCat-Flash-Chat", trust_remote_code=True),
282
    "MambaForCausalLM": _HfExamplesInfo("state-spaces/mamba-130m-hf"),
283
284
285
286
287
    "Mamba2ForCausalLM": _HfExamplesInfo("mistralai/Mamba-Codestral-7B-v0.1",
                                         min_transformers_version="4.55.3",
                                         extras={
                                            "random": "yujiepan/mamba2-codestral-v0.1-tiny-random", # noqa: E501
                                         }),
288
289
290
291
292
    "FalconMambaForCausalLM": _HfExamplesInfo("tiiuae/falcon-mamba-7b-instruct"),  # noqa: E501
    "MiniCPMForCausalLM": _HfExamplesInfo("openbmb/MiniCPM-2B-sft-bf16",
                                         trust_remote_code=True),
    "MiniCPM3ForCausalLM": _HfExamplesInfo("openbmb/MiniCPM3-4B",
                                         trust_remote_code=True),
293
    "MiniMaxForCausalLM": _HfExamplesInfo("MiniMaxAI/MiniMax-Text-01-hf"),
294
    "MiniMaxText01ForCausalLM": _HfExamplesInfo("MiniMaxAI/MiniMax-Text-01",
295
296
                                                trust_remote_code=True,
                                                revision="a59aa9cbc53b9fb8742ca4e9e1531b9802b6fdc3"),  # noqa: E501
297
298
    "MiniMaxM1ForCausalLM": _HfExamplesInfo("MiniMaxAI/MiniMax-M1-40k",
                                            trust_remote_code=True),
299
    "MistralForCausalLM": _HfExamplesInfo("mistralai/Mistral-7B-Instruct-v0.1"),
300
    "MixtralForCausalLM": _HfExamplesInfo("mistralai/Mixtral-8x7B-Instruct-v0.1",  # noqa: E501
301
                                          {"tiny": "TitanML/tiny-mixtral"}),  # noqa: E501
302
303
304
    "MotifForCausalLM": _HfExamplesInfo("Motif-Technologies/Motif-2.6B",
                                        trust_remote_code=True,
                                        v0_only=True),
305
306
307
    "MptForCausalLM": _HfExamplesInfo("mpt", is_available_online=False),
    "MPTForCausalLM": _HfExamplesInfo("mosaicml/mpt-7b"),
    "NemotronForCausalLM": _HfExamplesInfo("nvidia/Minitron-8B-Base"),
Luis Vega's avatar
Luis Vega committed
308
309
    "NemotronHForCausalLM": _HfExamplesInfo("nvidia/Nemotron-H-8B-Base-8K",
                                            trust_remote_code=True),
310
    "OlmoForCausalLM": _HfExamplesInfo("allenai/OLMo-1B-hf"),
311
    "Olmo2ForCausalLM": _HfExamplesInfo("allenai/OLMo-2-0425-1B"),
312
    "Olmo3ForCausalLM": _HfExamplesInfo("shanearora/2025-sep-a-base-model"),
313
    "OlmoeForCausalLM": _HfExamplesInfo("allenai/OLMoE-1B-7B-0924-Instruct"),
314
315
    "OPTForCausalLM": _HfExamplesInfo("facebook/opt-125m",
                                      {"1b": "facebook/opt-iml-max-1.3b"}),
316
317
318
    "OrionForCausalLM": _HfExamplesInfo("OrionStarAI/Orion-14B-Chat",
                                        trust_remote_code=True),
    "PersimmonForCausalLM": _HfExamplesInfo("adept/persimmon-8b-chat"),
319
    "PhiForCausalLM": _HfExamplesInfo("microsoft/phi-2"),
320
321
322
    "Phi3ForCausalLM": _HfExamplesInfo("microsoft/Phi-3-mini-4k-instruct"),
    "PhiMoEForCausalLM": _HfExamplesInfo("microsoft/Phi-3.5-MoE-instruct",
                                         trust_remote_code=True),
Shinichi Hemmi's avatar
Shinichi Hemmi committed
323
    "Plamo2ForCausalLM": _HfExamplesInfo("pfnet/plamo-2-1b",
324
325
326
                                         max_transformers_version="4.55.4",
                                         transformers_version_reason="HF model uses remote code that is not compatible with latest Transformers",  # noqa: E501
                                         trust_remote_code=True),
327
    "QWenLMHeadModel": _HfExamplesInfo("Qwen/Qwen-7B-Chat",
328
329
                                       max_transformers_version="4.53",
                                       transformers_version_reason="HF model uses remote code that is not compatible with latest Transformers",  # noqa: E501
330
                                       trust_remote_code=True),
331
332
    "Qwen2ForCausalLM": _HfExamplesInfo("Qwen/Qwen2-0.5B-Instruct",
                                        extras={"2.5": "Qwen/Qwen2.5-0.5B-Instruct"}), # noqa: E501
333
    "Qwen2MoeForCausalLM": _HfExamplesInfo("Qwen/Qwen1.5-MoE-A2.7B-Chat"),
Jee Jee Li's avatar
Jee Jee Li committed
334
335
    "Qwen3ForCausalLM": _HfExamplesInfo("Qwen/Qwen3-8B"),
    "Qwen3MoeForCausalLM": _HfExamplesInfo("Qwen/Qwen3-30B-A3B"),
336
    "Qwen3NextForCausalLM": _HfExamplesInfo("Qwen/Qwen3-Next-80B-A3B-Instruct",
337
338
                                            extras={"tiny-random": "tiny-random/qwen3-next-moe"}, # noqa: E501
                                            min_transformers_version="4.56.3"),
339
    "RWForCausalLM": _HfExamplesInfo("tiiuae/falcon-40b"),
340
341
342
    "SeedOssForCausalLM": _HfExamplesInfo("ByteDance-Seed/Seed-OSS-36B-Instruct", # noqa: E501
                                          trust_remote_code=True,
                                          is_available_online=False),
343
    "SmolLM3ForCausalLM": _HfExamplesInfo("HuggingFaceTB/SmolLM3-3B"),
344
345
    "StableLMEpochForCausalLM": _HfExamplesInfo("stabilityai/stablelm-zephyr-3b"),  # noqa: E501
    "StableLmForCausalLM": _HfExamplesInfo("stabilityai/stablelm-3b-4e1t"),
346
    "Starcoder2ForCausalLM": _HfExamplesInfo("bigcode/starcoder2-3b"),
Song's avatar
Song committed
347
    "Step3TextForCausalLM": _HfExamplesInfo("stepfun-ai/step3",
348
                                            trust_remote_code=True),
349
350
    "SolarForCausalLM": _HfExamplesInfo("upstage/solar-pro-preview-instruct",
                                        trust_remote_code=True),
351
352
    "TeleChat2ForCausalLM": _HfExamplesInfo("Tele-AI/TeleChat2-3B",
                                            trust_remote_code=True),
353
354
    "TeleFLMForCausalLM": _HfExamplesInfo("CofeAI/FLM-2-52B-Instruct-2407",
                                            trust_remote_code=True),
355
    "XverseForCausalLM": _HfExamplesInfo("xverse/XVERSE-7B-Chat",
356
                                         tokenizer="meta-llama/Llama-2-7b",
357
                                         trust_remote_code=True),
358
    "Zamba2ForCausalLM": _HfExamplesInfo("Zyphra/Zamba2-7B-instruct"),
359
360
    "MiMoForCausalLM": _HfExamplesInfo("XiaomiMiMo/MiMo-7B-RL",
                                        trust_remote_code=True),
361
    "Dots1ForCausalLM": _HfExamplesInfo("rednote-hilab/dots.llm1.inst"),
362
363
364
365
}

_EMBEDDING_EXAMPLE_MODELS = {
    # [Text-only]
366
367
    "BertModel": _HfExamplesInfo("BAAI/bge-base-en-v1.5"),
    "Gemma2Model": _HfExamplesInfo("BAAI/bge-multilingual-gemma2"),  # noqa: E501
368
    "Gemma3TextModel": _HfExamplesInfo("google/embeddinggemma-300m"),
369
    "GritLM": _HfExamplesInfo("parasail-ai/GritLM-7B-vllm"),
370
371
    "GteModel": _HfExamplesInfo("Snowflake/snowflake-arctic-embed-m-v2.0",
                                               trust_remote_code=True),
372
373
    "GteNewModel": _HfExamplesInfo("Alibaba-NLP/gte-base-en-v1.5",
                                   trust_remote_code=True,
374
                                   hf_overrides={"architectures": ["GteNewModel"]}),  # noqa: E501
375
376
    "InternLM2ForRewardModel": _HfExamplesInfo("internlm/internlm2-1_8b-reward",
                                               trust_remote_code=True),
377
    "JambaForSequenceClassification": _HfExamplesInfo("ai21labs/Jamba-tiny-reward-dev"),  # noqa: E501
378
    "LlamaModel": _HfExamplesInfo("llama", is_available_online=False),
379
    "MistralModel": _HfExamplesInfo("intfloat/e5-mistral-7b-instruct"),
380
    "ModernBertModel": _HfExamplesInfo("Alibaba-NLP/gte-modernbert-base",
381
                                trust_remote_code=True),
382
    "NomicBertModel": _HfExamplesInfo("nomic-ai/nomic-embed-text-v2-moe",
383
                                               trust_remote_code=True),  # noqa: E501
384
    "Qwen2Model": _HfExamplesInfo("ssmits/Qwen2-7B-Instruct-embed-base"),
385
386
387
388
389
390
    "Qwen2ForRewardModel": _HfExamplesInfo("Qwen/Qwen2.5-Math-RM-72B",
                                           max_transformers_version="4.53",
                                           transformers_version_reason="HF model uses remote code that is not compatible with latest Transformers"),  # noqa: E501
    "Qwen2ForProcessRewardModel": _HfExamplesInfo("Qwen/Qwen2.5-Math-PRM-7B",
                                                  max_transformers_version="4.53",
                                                  transformers_version_reason="HF model uses remote code that is not compatible with latest Transformers"),  # noqa: E501
391
392
393
    "RobertaModel": _HfExamplesInfo("sentence-transformers/stsb-roberta-base-v2"),  # noqa: E501
    "RobertaForMaskedLM": _HfExamplesInfo("sentence-transformers/all-roberta-large-v1"),  # noqa: E501
    "XLMRobertaModel": _HfExamplesInfo("intfloat/multilingual-e5-small"),  # noqa: E501
394
395
396
397
398
    # [Multimodal]
    "LlavaNextForConditionalGeneration": _HfExamplesInfo("royokong/e5-v"),
    "Phi3VForCausalLM": _HfExamplesInfo("TIGER-Lab/VLM2Vec-Full",
                                         trust_remote_code=True),
    "Qwen2VLForConditionalGeneration": _HfExamplesInfo("MrLight/dse-qwen2-2b-mrl-v1"), # noqa: E501
399
    "PrithviGeoSpatialMAE": _HfExamplesInfo("ibm-nasa-geospatial/Prithvi-EO-2.0-300M-TL-Sen1Floods11", # noqa: E501
400
401
402
403
404
405
406
                                            dtype=torch.float16,
                                            enforce_eager=True,
                                            skip_tokenizer_init=True,
                                            # This is to avoid the model
                                            # going OOM in CI
                                            max_num_seqs=32,
                                            ),
407
    "Terratorch": _HfExamplesInfo("ibm-nasa-geospatial/Prithvi-EO-2.0-300M-TL-Sen1Floods11", # noqa: E501
408
409
410
411
412
413
                                  dtype=torch.float16,
                                  enforce_eager=True,
                                  skip_tokenizer_init=True,
                                  # This is to avoid the model going OOM in CI
                                  max_num_seqs=32,
                                  ),
414
415
}

416
417
418
419
420
_SEQUENCE_CLASSIFICATION_EXAMPLE_MODELS = {
    # [Decoder-only]
    "GPT2ForSequenceClassification": _HfExamplesInfo("nie3e/sentiment-polish-gpt2-small"),  # noqa: E501

    # [Cross-encoder]
421
    "BertForSequenceClassification": _HfExamplesInfo("cross-encoder/ms-marco-MiniLM-L-6-v2"),  # noqa: E501
422
    "BertForTokenClassification": _HfExamplesInfo("boltuix/NeuroBERT-NER"),
423
424
425
426
    "GteNewForSequenceClassification": _HfExamplesInfo("Alibaba-NLP/gte-multilingual-reranker-base",  # noqa: E501
                                                       trust_remote_code=True,
                                                       hf_overrides={
                                                           "architectures": ["GteNewForSequenceClassification"]}),# noqa: E501
427
428
429
    "ModernBertForSequenceClassification": _HfExamplesInfo("Alibaba-NLP/gte-reranker-modernbert-base"), # noqa: E501
    "RobertaForSequenceClassification": _HfExamplesInfo("cross-encoder/quora-roberta-base"),  # noqa: E501
    "XLMRobertaForSequenceClassification": _HfExamplesInfo("BAAI/bge-reranker-v2-m3"),  # noqa: E501
430
431
}

432
433
434
435
436
437
438
439
440
441
442
_AUTOMATIC_CONVERTED_MODELS = {
    # Use as_seq_cls_model for automatic conversion
    "GemmaForSequenceClassification": _HfExamplesInfo("BAAI/bge-reranker-v2-gemma",  # noqa: E501
                                                      hf_overrides={"architectures": ["GemmaForSequenceClassification"], # noqa: E501
                                                                    "classifier_from_token": ["Yes"],  # noqa: E501
                                                                    "method": "no_post_processing"}),  # noqa: E501
    "LlamaForSequenceClassification": _HfExamplesInfo("Skywork/Skywork-Reward-V2-Llama-3.2-1B"),  # noqa: E501
    "Qwen2ForSequenceClassification": _HfExamplesInfo("jason9693/Qwen2.5-1.5B-apeach"),  # noqa: E501
    "Qwen3ForSequenceClassification": _HfExamplesInfo("tomaarsen/Qwen3-Reranker-0.6B-seq-cls"),  # noqa: E501
}

443
444
_MULTIMODAL_EXAMPLE_MODELS = {
    # [Decoder-only]
445
    "AriaForConditionalGeneration": _HfExamplesInfo("rhymes-ai/Aria"),
Jennifer Zhao's avatar
Jennifer Zhao committed
446
    "AyaVisionForConditionalGeneration": _HfExamplesInfo("CohereForAI/aya-vision-8b"), # noqa: E501
447
    "Blip2ForConditionalGeneration": _HfExamplesInfo("Salesforce/blip2-opt-2.7b",  # noqa: E501
448
                                                     extras={"6b": "Salesforce/blip2-opt-6.7b"}),  # noqa: E501
449
    "ChameleonForConditionalGeneration": _HfExamplesInfo("facebook/chameleon-7b"),  # noqa: E501
450
    "Cohere2VisionForConditionalGeneration": _HfExamplesInfo("CohereLabs/command-a-vision-07-2025"), # noqa: E501
451
    "DeepseekVLV2ForCausalLM": _HfExamplesInfo("deepseek-ai/deepseek-vl2-tiny",  # noqa: E501
452
453
454
                                                extras={"fork": "Isotr0py/deepseek-vl2-tiny"},  # noqa: E501
                                                max_transformers_version="4.48",  # noqa: E501
                                                transformers_version_reason="HF model is not compatible.",  # noqa: E501
455
                                                hf_overrides={"architectures": ["DeepseekVLV2ForCausalLM"]}),  # noqa: E501
Roger Wang's avatar
Roger Wang committed
456
457
    "DotsOCRForCausalLM": _HfExamplesInfo("rednote-hilab/dots.ocr",
                                          trust_remote_code=True),
458
    "Emu3ForConditionalGeneration": _HfExamplesInfo("BAAI/Emu3-Chat-hf"),
459
460
    "Ernie4_5_VLMoeForConditionalGeneration": _HfExamplesInfo("baidu/ERNIE-4.5-VL-28B-A3B-PT",  # noqa: E501
                                                              trust_remote_code=True),
461
    "FuyuForCausalLM": _HfExamplesInfo("adept/fuyu-8b"),
462
    "Gemma3ForConditionalGeneration": _HfExamplesInfo("google/gemma-3-4b-it"),
Nicolò Lucchesi's avatar
Nicolò Lucchesi committed
463
464
    "Gemma3nForConditionalGeneration": _HfExamplesInfo("google/gemma-3n-E2B-it",    # noqa: E501
                                        min_transformers_version="4.53"),
465
    "GraniteSpeechForConditionalGeneration": _HfExamplesInfo("ibm-granite/granite-speech-3.3-2b"),  # noqa: E501
466
    "GLM4VForCausalLM": _HfExamplesInfo("zai-org/glm-4v-9b",
467
468
                                        trust_remote_code=True,
                                        hf_overrides={"architectures": ["GLM4VForCausalLM"]}),  # noqa: E501
469
    "Glm4vForConditionalGeneration": _HfExamplesInfo("zai-org/GLM-4.1V-9B-Thinking"),  # noqa: E501
470
    "Glm4vMoeForConditionalGeneration": _HfExamplesInfo("zai-org/GLM-4.5V",
471
                                                        min_transformers_version="4.56"),  # noqa: E501
472
    "H2OVLChatModel": _HfExamplesInfo("h2oai/h2ovl-mississippi-800m",
473
                                      trust_remote_code=True,
474
475
476
                                      extras={"2b": "h2oai/h2ovl-mississippi-2b"},  # noqa: E501
                                      max_transformers_version="4.48",  # noqa: E501
                                      transformers_version_reason="HF model is not compatible."),  # noqa: E501
477
478
    "HCXVisionForCausalLM": _HfExamplesInfo("naver-hyperclovax/HyperCLOVAX-SEED-Vision-Instruct-3B",  # noqa: E501
                                            trust_remote_code=True),
479
480
    "Idefics3ForConditionalGeneration": _HfExamplesInfo("HuggingFaceM4/Idefics3-8B-Llama3",  # noqa: E501
                                                        {"tiny": "HuggingFaceTB/SmolVLM-256M-Instruct"},    # noqa: E501
481
482
                                                        min_transformers_version="4.56",
                                                        transformers_version_reason="HF model broken in 4.55"),  # noqa: E501
483
484
    "InternS1ForConditionalGeneration": _HfExamplesInfo("internlm/Intern-S1",
                                                        trust_remote_code=True),  # noqa: E501
485
    "InternVLChatModel": _HfExamplesInfo("OpenGVLab/InternVL2-1B",
486
                                         extras={"2B": "OpenGVLab/InternVL2-2B",
487
488
489
490
                                                 "3.0": "OpenGVLab/InternVL3-1B",   # noqa: E501
                                                 "3.5-qwen3": "OpenGVLab/InternVL3_5-1B",   # noqa: E501
                                                 "3.5-qwen3moe": "OpenGVLab/InternVL3_5-30B-A3B",   # noqa: E501
                                                 "3.5-gptoss": "OpenGVLab/InternVL3_5-GPT-OSS-20B-A4B-Preview"},  # noqa: E501
491
                                         trust_remote_code=True),
492
    "InternVLForConditionalGeneration": _HfExamplesInfo("OpenGVLab/InternVL3-1B-hf"),    # noqa: E501
493
494
    "KeyeForConditionalGeneration": _HfExamplesInfo("Kwai-Keye/Keye-VL-8B-Preview", # noqa: E501
                                                    trust_remote_code=True),
495
496
    "KeyeVL1_5ForConditionalGeneration": _HfExamplesInfo("Kwai-Keye/Keye-VL-1_5-8B", # noqa: E501
                                                         trust_remote_code=True),
497
498
    "KimiVLForConditionalGeneration": _HfExamplesInfo("moonshotai/Kimi-VL-A3B-Instruct",  # noqa: E501
                                                      extras={"thinking": "moonshotai/Kimi-VL-A3B-Thinking"},  # noqa: E501
499
                                                      trust_remote_code=True),
500
    "Llama4ForConditionalGeneration": _HfExamplesInfo("meta-llama/Llama-4-Scout-17B-16E-Instruct",   # noqa: E501
501
                                                      max_model_len=10240,
502
                                                      extras={"llama-guard-4": "meta-llama/Llama-Guard-4-12B"},  # noqa: E501
503
                                                      ),
504
    "LlavaForConditionalGeneration": _HfExamplesInfo("llava-hf/llava-1.5-7b-hf",
505
506
                                                     extras={"mistral": "mistral-community/pixtral-12b", # noqa: E501
                                                             "mistral-fp8": "nm-testing/pixtral-12b-FP8-dynamic"}),  # noqa: E501
507
508
509
    "LlavaNextForConditionalGeneration": _HfExamplesInfo("llava-hf/llava-v1.6-mistral-7b-hf"),  # noqa: E501
    "LlavaNextVideoForConditionalGeneration": _HfExamplesInfo("llava-hf/LLaVA-NeXT-Video-7B-hf"),  # noqa: E501
    "LlavaOnevisionForConditionalGeneration": _HfExamplesInfo("llava-hf/llava-onevision-qwen2-0.5b-ov-hf"),  # noqa: E501
510
    "MantisForConditionalGeneration": _HfExamplesInfo("TIGER-Lab/Mantis-8B-siglip-llama3",  # noqa: E501
511
512
                                                      max_transformers_version="4.48",  # noqa: E501
                                                      transformers_version_reason="HF model is not compatible.",  # noqa: E501
513
                                                      hf_overrides={"architectures": ["MantisForConditionalGeneration"]}),  # noqa: E501
514
515
    "MiDashengLMModel": _HfExamplesInfo("mispeech/midashenglm-7b",
                            trust_remote_code=True),
516
517
    "MiniCPMO": _HfExamplesInfo("openbmb/MiniCPM-o-2_6",
                                trust_remote_code=True),
518
    "MiniCPMV": _HfExamplesInfo("openbmb/MiniCPM-Llama3-V-2_5",
tc-mb's avatar
tc-mb committed
519
                                extras={"2.6": "openbmb/MiniCPM-V-2_6", "4.0": "openbmb/MiniCPM-V-4", "4.5": "openbmb/MiniCPM-V-4_5"},  # noqa: E501
520
                                trust_remote_code=True),
521
    "MiniMaxVL01ForConditionalGeneration": _HfExamplesInfo("MiniMaxAI/MiniMax-VL-01", # noqa: E501
522
523
                                              trust_remote_code=True,
                                              v0_only=True),
524
525
    "Mistral3ForConditionalGeneration": _HfExamplesInfo("mistralai/Mistral-Small-3.1-24B-Instruct-2503",  # noqa: E501
                                                        extras={"fp8": "nm-testing/Mistral-Small-3.1-24B-Instruct-2503-FP8-dynamic"}),  # noqa: E501
526
    "MolmoForCausalLM": _HfExamplesInfo("allenai/Molmo-7B-D-0924",
527
                                        max_transformers_version="4.48",
528
                                        transformers_version_reason="Incorrectly-detected `tensorflow` import.",  # noqa: E501
529
                                        extras={"olmo": "allenai/Molmo-7B-O-0924"},  # noqa: E501
530
531
532
                                        trust_remote_code=True),
    "NVLM_D": _HfExamplesInfo("nvidia/NVLM-D-72B",
                              trust_remote_code=True),
533
534
    "Llama_Nemotron_Nano_VL" : _HfExamplesInfo("nvidia/Llama-3.1-Nemotron-Nano-VL-8B-V1", # noqa: E501
                                                     trust_remote_code=True),
535
    "NemotronH_Nano_VL_V2": _HfExamplesInfo("nano_vl_dummy",
536
537
                                          is_available_online=False,
                                          trust_remote_code=True),
538
    "Ovis": _HfExamplesInfo("AIDC-AI/Ovis2-1B", trust_remote_code=True,
539
540
                            max_transformers_version="4.53",
                            transformers_version_reason="HF model is not compatible",  # noqa: E501
541
542
                            extras={"1.6-llama": "AIDC-AI/Ovis1.6-Llama3.2-3B",
                                    "1.6-gemma": "AIDC-AI/Ovis1.6-Gemma2-9B"}),  # noqa: E501
543
    "Ovis2_5": _HfExamplesInfo("AIDC-AI/Ovis2.5-2B",
544
                               trust_remote_code=True),
545
546
    "PaliGemmaForConditionalGeneration": _HfExamplesInfo("google/paligemma-3b-mix-224",  # noqa: E501
                                                         extras={"v2": "google/paligemma2-3b-ft-docci-448"}),  # noqa: E501
547
548
    "Phi3VForCausalLM": _HfExamplesInfo("microsoft/Phi-3-vision-128k-instruct",
                                        trust_remote_code=True,
549
550
                                        max_transformers_version="4.48",
                                        transformers_version_reason="Use of deprecated imports which have been removed.",  # noqa: E501
551
                                        extras={"phi3.5": "microsoft/Phi-3.5-vision-instruct"}),  # noqa: E501
552
553
    "Phi4MMForCausalLM": _HfExamplesInfo("microsoft/Phi-4-multimodal-instruct",
                                        trust_remote_code=True),
554
555
    "Phi4MultimodalForCausalLM": _HfExamplesInfo("microsoft/Phi-4-multimodal-instruct",  # noqa: E501
                                                 revision="refs/pr/70"),
556
    "PixtralForConditionalGeneration": _HfExamplesInfo("mistralai/Pixtral-12B-2409",  # noqa: E501
557
                                                       tokenizer_mode="mistral"),
558
559
560
561
    "QwenVLForConditionalGeneration": _HfExamplesInfo("Qwen/Qwen-VL",
                                                      extras={"chat": "Qwen/Qwen-VL-Chat"},  # noqa: E501
                                                      trust_remote_code=True,
                                                      hf_overrides={"architectures": ["QwenVLForConditionalGeneration"]}),  # noqa: E501
562
563
    "Qwen2AudioForConditionalGeneration": _HfExamplesInfo("Qwen/Qwen2-Audio-7B-Instruct"),  # noqa: E501
    "Qwen2VLForConditionalGeneration": _HfExamplesInfo("Qwen/Qwen2-VL-2B-Instruct"),  # noqa: E501
564
565
    "Qwen2_5_VLForConditionalGeneration": _HfExamplesInfo("Qwen/Qwen2.5-VL-3B-Instruct", # noqa: E501
                                                          max_model_len=4096),
566
567
    "Qwen2_5OmniModel": _HfExamplesInfo("Qwen/Qwen2.5-Omni-3B"),
    "Qwen2_5OmniForConditionalGeneration": _HfExamplesInfo("Qwen/Qwen2.5-Omni-7B-AWQ"),  # noqa: E501
568
569
    "Qwen3VLForConditionalGeneration": _HfExamplesInfo("Qwen/Qwen3-VL-4B-Instruct", # noqa: E501
                                                        max_model_len=4096,
570
571
                                                        min_transformers_version="4.57",
                                                        is_available_online=False),
572
    "Qwen3VLMoeForConditionalGeneration": _HfExamplesInfo("Qwen/Qwen3-VL-30B-A3B-Instruct", # noqa: E501
573
574
575
                                                          max_model_len=4096,
                                                          min_transformers_version="4.57",
                                                          is_available_online=False),
576
577
    "RForConditionalGeneration": _HfExamplesInfo("YannQi/R-4B",
                                                 trust_remote_code=True),
578
579
    "SkyworkR1VChatModel": _HfExamplesInfo("Skywork/Skywork-R1V-38B",
                                           trust_remote_code=True),
580
    "SmolVLMForConditionalGeneration": _HfExamplesInfo("HuggingFaceTB/SmolVLM2-2.2B-Instruct",  # noqa: E501
581
582
                                                       min_transformers_version="4.56",
                                                       transformers_version_reason="HF model broken in 4.55"),  # noqa: E501
Song's avatar
Song committed
583
    "Step3VLForConditionalGeneration": _HfExamplesInfo("stepfun-ai/step3",
584
                                                        trust_remote_code=True),
585
    "UltravoxModel": _HfExamplesInfo("fixie-ai/ultravox-v0_5-llama-3_2-1b",  # noqa: E501
586
                                     trust_remote_code=True),
587
    "TarsierForConditionalGeneration": _HfExamplesInfo("omni-research/Tarsier-7b"),  # noqa: E501
588
589
    "Tarsier2ForConditionalGeneration": _HfExamplesInfo("omni-research/Tarsier2-Recap-7b",  # noqa: E501
                                                        hf_overrides={"architectures": ["Tarsier2ForConditionalGeneration"]}),  # noqa: E501
590
591
592
593
594
595
    "VoxtralForConditionalGeneration": _HfExamplesInfo(
        "mistralai/Voxtral-Mini-3B-2507",
        min_transformers_version="4.54",
        # disable this temporarily until we support HF format
        is_available_online=False,
    ),
596
    # [Encoder-decoder]
597
    "WhisperForConditionalGeneration": _HfExamplesInfo("openai/whisper-large-v3"),  # noqa: E501
598
599
    # [Cross-encoder]
    "JinaVLForRanking": _HfExamplesInfo("jinaai/jina-reranker-m0"),   # noqa: E501
600
601
}

602

603
604
605
_SPECULATIVE_DECODING_EXAMPLE_MODELS = {
    "MedusaModel": _HfExamplesInfo("JackFram/llama-68m",
                                   speculative_model="abhigoyal/vllm-medusa-llama-68m-random"),  # noqa: E501
606
607
608
609
    # Temporarily disabled.
    # TODO(woosuk): Re-enable this once the MLP Speculator is supported in V1.
    # "MLPSpeculatorPreTrainedModel": _HfExamplesInfo("JackFram/llama-160m",
    #                                                 speculative_model="ibm-ai-platform/llama-160m-accelerator"),  # noqa: E501
610
611
612
    "DeepSeekMTPModel": _HfExamplesInfo("luccafong/deepseek_mtp_main_random",
                                        speculative_model="luccafong/deepseek_mtp_draft_random",  # noqa: E501
                                        trust_remote_code=True),
613
614
615
    "EagleDeepSeekMTPModel": _HfExamplesInfo("eagle618/deepseek-v3-random",
                                        speculative_model="eagle618/eagle-deepseek-v3-random",  # noqa: E501
                                        trust_remote_code=True),
616
    "EagleLlamaForCausalLM": _HfExamplesInfo("meta-llama/Meta-Llama-3-8B-Instruct", # noqa: E501
617
618
                                             trust_remote_code=True,
                                             speculative_model="yuhuili/EAGLE-LLaMA3-Instruct-8B",
619
620
                                             tokenizer="meta-llama/Meta-Llama-3-8B-Instruct"), # noqa: E501
    "Eagle3LlamaForCausalLM": _HfExamplesInfo("meta-llama/Llama-3.1-8B-Instruct",  # noqa: E501
621
                                            trust_remote_code=True,
622
623
624
625
626
                                            speculative_model="yuhuili/EAGLE3-LLaMA3.1-Instruct-8B", # noqa: E501
                                            tokenizer="meta-llama/Llama-3.1-8B-Instruct",
                                            use_original_num_layers=True,
                                            max_model_len=10240),
    "LlamaForCausalLMEagle3": _HfExamplesInfo("Qwen/Qwen3-8B",  # noqa: E501
627
628
                                            trust_remote_code=True,
                                            speculative_model="AngelSlim/Qwen3-8B_eagle3",   # noqa: E501
629
630
                                            tokenizer="Qwen/Qwen3-8B",
                                            use_original_num_layers=True),
zhiweiz's avatar
zhiweiz committed
631
632
633
634
635
    "EagleLlama4ForCausalLM": _HfExamplesInfo(
        "morgendave/EAGLE-Llama-4-Scout-17B-16E-Instruct",
        trust_remote_code=True,
        speculative_model="morgendave/EAGLE-Llama-4-Scout-17B-16E-Instruct",
        tokenizer="meta-llama/Llama-4-Scout-17B-16E-Instruct"),  # noqa: E501
636
637
638
639
640
    "EagleMiniCPMForCausalLM": _HfExamplesInfo("openbmb/MiniCPM-1B-sft-bf16",
                                            trust_remote_code=True,
                                            is_available_online=False,
                                            speculative_model="openbmb/MiniCPM-2B-sft-bf16",
                                            tokenizer="openbmb/MiniCPM-2B-sft-bf16"),
641
642
643
    "ErnieMTPModel": _HfExamplesInfo("baidu/ERNIE-4.5-21B-A3B-PT",
                                    trust_remote_code=True,
                                    speculative_model="baidu/ERNIE-4.5-21B-A3B-PT"),
Yuxuan Zhang's avatar
Yuxuan Zhang committed
644
645
    "Glm4MoeMTPModel": _HfExamplesInfo("zai-org/GLM-4.5",
                                        speculative_model="zai-org/GLM-4.5",
Yuxuan Zhang's avatar
Yuxuan Zhang committed
646
647
                                        min_transformers_version="4.54",
                                        is_available_online=False),
XuruiYang's avatar
XuruiYang committed
648
649
650
651
    "LongCatFlashMTPModel": _HfExamplesInfo(
        "meituan-longcat/LongCat-Flash-Chat",
        trust_remote_code=True,
        speculative_model="meituan-longcat/LongCat-Flash-Chat"),
652
653
    "MiMoMTPModel": _HfExamplesInfo("XiaomiMiMo/MiMo-7B-RL",
                                    trust_remote_code=True,
654
655
                                    speculative_model="XiaomiMiMo/MiMo-7B-RL"),
    "Qwen3NextMTP": _HfExamplesInfo("Qwen/Qwen3-Next-80B-A3B-Instruct",
656
                                     min_transformers_version="4.56.3"),
657
658
}

659
_TRANSFORMERS_BACKEND_MODELS = {
660
    "TransformersModel": _HfExamplesInfo("Qwen/Qwen3-Embedding-0.6B"),
661
    "TransformersForCausalLM": _HfExamplesInfo("hmellor/Ilama-3.2-1B", trust_remote_code=True),  # noqa: E501
662
    "TransformersForMultimodalLM": _HfExamplesInfo("BAAI/Emu3-Chat-hf"),
663
664
}

665
666
667
_EXAMPLE_MODELS = {
    **_TEXT_GENERATION_EXAMPLE_MODELS,
    **_EMBEDDING_EXAMPLE_MODELS,
668
    **_SEQUENCE_CLASSIFICATION_EXAMPLE_MODELS,
669
670
    **_MULTIMODAL_EXAMPLE_MODELS,
    **_SPECULATIVE_DECODING_EXAMPLE_MODELS,
671
    **_TRANSFORMERS_BACKEND_MODELS,
672
673
674
675
676
677
678
679
680
}


class HfExampleModels:
    def __init__(self, hf_models: Mapping[str, _HfExamplesInfo]) -> None:
        super().__init__()

        self.hf_models = hf_models

681
    def get_supported_archs(self) -> Set[str]:
682
683
684
685
686
        return self.hf_models.keys()

    def get_hf_info(self, model_arch: str) -> _HfExamplesInfo:
        return self.hf_models[model_arch]

687
688
689
690
691
    def find_hf_info(self, model_id: str) -> _HfExamplesInfo:
        for info in self.hf_models.values():
            if info.default == model_id:
                return info

692
693
694
695
696
        # Fallback to extras
        for info in self.hf_models.values():
            if any(extra == model_id for extra in info.extras.values()):
                return info

697
698
        raise ValueError(f"No example model defined for {model_id}")

699

Patrick von Platen's avatar
Patrick von Platen committed
700
HF_EXAMPLE_MODELS = HfExampleModels(_EXAMPLE_MODELS)
701
AUTO_EXAMPLE_MODELS = HfExampleModels(_AUTOMATIC_CONVERTED_MODELS)