[Mypy Fix] Mypy fix for "vllm/model_executor/models/[cC][dD]"

tools/pre_commit/mypy.py

@@ -111,10 +111,7 @@
     "tests/v1/worker",
 ]
 
-# TODO(woosuk): Include the code from Megatron and HuggingFace.
 EXCLUDE = [
-    r"vllm/model_executor/models/[cC]",
-    r"vllm/model_executor/models/[dD]",
     r"vllm/model_executor/models/[eE]",
     r"vllm/model_executor/models/[fF]",
     r"vllm/model_executor/models/[gG]",

vllm/model_executor/models/clip.py

@@ -142,6 +142,7 @@ def get_num_image_tokens(
 
         pooler_config = self.ctx.model_config.pooler_config
         assert pooler_config is not None
+        assert pooler_config.seq_pooling_type is not None
 
         return get_num_selected_vision_tokens(
             vision_encoder_info.get_num_image_tokens(
@@ -854,6 +855,7 @@ def get_image_features(
         feature_select_strategy: VisionFeatureSelectStrategy | None = None,
     ) -> torch.Tensor:
         if feature_select_strategy is None:
+            assert self.pooler_config.seq_pooling_type is not None
             feature_select_strategy = _get_vision_feature_select_strategy(
                 self.pooler_config.seq_pooling_type
             )
@@ -961,6 +963,7 @@ def forward(
             raise RuntimeError("PP is not supported for this model")
 
         # Multimodal inputs
+        assert inputs_embeds is not None
         if not self._is_text_input:
             return inputs_embeds
 

vllm/model_executor/models/cohere2_moe.py

@@ -178,6 +178,7 @@ def __init__(
         self.max_position_embeddings = getattr(
             config, "model_max_length", None
         ) or getattr(config, "max_position_embeddings", 8192)
+        assert isinstance(self.max_position_embeddings, int)
         self.qkv_proj = QKVParallelLinear(
             self.hidden_size,
             self.head_dim,
@@ -283,6 +284,7 @@ def __init__(
             prefix=f"{prefix}.gate",
         )
 
+        self.shared_experts: Cohere2MoeMLP | None
         if hasattr(config, "num_shared_experts") and config.num_shared_experts > 0:
             self.shared_experts = Cohere2MoeMLP(
                 config=config,

vllm/model_executor/models/cohere2_vision.py

@@ -416,7 +416,7 @@ def _parse_and_validate_image_input(
         )
 
     def _patch_quant_config(
-        self, config: PretrainedConfig, quant_config: QuantizationConfig
+        self, config: PretrainedConfig, quant_config: QuantizationConfig | None
     ):
         # the awq models from OpenGVLab missing `modules_to_not_convert`
         # patch the quant_config to add `modules_to_not_convert` back

vllm/model_executor/models/cohere_asr.py

@@ -55,6 +55,7 @@
     CohereASRFeatureExtractor,
     CohereASRProcessor,
 )
+from vllm.utils.collection_utils import is_list_of
 from vllm.v1.attention.backend import (
     AttentionType,
 )
@@ -183,6 +184,7 @@ def _init_qkv(
     def forward(
         self,
         hidden_states: torch.Tensor,
+        encoder_hidden_states: torch.Tensor | None = None,
     ) -> torch.Tensor:
         qkv, _ = self.qkv_proj(hidden_states)
         q, k, v = qkv.split([self.q_size, self.kv_size, self.kv_size], dim=-1)
@@ -241,7 +243,7 @@ def _init_qkv(
     def forward(
         self,
         hidden_states: torch.Tensor,
-        encoder_hidden_states: torch.Tensor | None,
+        encoder_hidden_states: torch.Tensor | None = None,
     ) -> torch.Tensor:
         q, _ = self.q_proj(hidden_states)
 
@@ -837,14 +839,16 @@ def __init__(
         out_channels: int,
         kernel_size: int,
         stride: int = 1,
-        padding: str | int = 0,
+        padding: str | int | list[int] | None = 0,
         dilation: int = 1,
         groups: int = 1,
         bias: bool = True,
         padding_mode: str = "zeros",
         device=None,
         dtype=None,
     ) -> None:
+        self._left_padding: int
+        self._right_padding: int
         if padding is None:
             self._left_padding = kernel_size - 1
             self._right_padding = stride - 1
@@ -903,7 +907,7 @@ def __init__(
         d_model: int,
         kernel_size: int,
         norm_type: str = "batch_norm",
-        conv_context_size: int | None = None,
+        conv_context_size: int | list[int] | None = None,
         pointwise_activation: str = "glu_",
         use_bias: bool = True,
     ) -> None:
@@ -1158,6 +1162,7 @@ def forward(
         q, k, v = self.forward_qkv(query, key, value)
         q = q.transpose(1, 2)  # (batch, time1, head, d_k)
 
+        assert pos_emb is not None
         n_batch_pos = pos_emb.size(0)
         p = self.linear_pos(pos_emb).view(n_batch_pos, -1, self.h, self.d_k)
         p = p.transpose(1, 2)  # (batch, head, time1, d_k)
@@ -1213,7 +1218,7 @@ def __init__(
         n_heads: int = 4,
         conv_kernel_size: int = 31,
         conv_norm_type: str = "batch_norm",
-        conv_context_size: int | None = None,
+        conv_context_size: int | list[int] | None = None,
         pos_bias_u: nn.Parameter | torch.Tensor | None = None,
         pos_bias_v: nn.Parameter | torch.Tensor | None = None,
         att_context_size: list[int] | None = None,
@@ -1382,6 +1387,7 @@ def __init__(self, *, vllm_config: VllmConfig):
             conv_kernel_size=conv_kernel_size,
         )
 
+        self.xscale: float | None
         if xscaling:
             self.xscale = math.sqrt(d_model)
         else:
@@ -1633,9 +1639,12 @@ def _calc_context_sizes(
     ) -> tuple[list[list[int]], list[int], list[float], list[int]]:
         # convert att_context_size to a standard list of lists
         if att_context_size:
-            att_context_size_all = list(att_context_size)
-            if isinstance(att_context_size_all[0], int):
-                att_context_size_all = [att_context_size_all]
+            if is_list_of(att_context_size, int, check="all"):
+                att_context_size_all = [att_context_size]
+            elif is_list_of(att_context_size, list, check="all"):
+                att_context_size_all = att_context_size
+            else:
+                raise ValueError("att_context_size cannot mix nested and flat values")
             for i, att_cs in enumerate(att_context_size_all):
                 if att_context_style == "chunked_limited":
                     if att_cs[0] > 0 and att_cs[0] % (att_cs[1] + 1) > 0:
@@ -1684,6 +1693,7 @@ def _calc_context_sizes(
             if conv_context_size == "causal":
                 conv_context_size = [conv_kernel_size - 1, 0]
             else:
+                assert isinstance(conv_context_size, list)
                 total = conv_context_size[0] + conv_context_size[1] + 1
                 if total != conv_kernel_size:
                     raise ValueError(
@@ -1940,7 +1950,7 @@ def _call_hf_processor(
     ):
         if mm_data:
             feature_extractor = self.info.get_feature_extractor(**mm_kwargs)
-            mm_data = dict(audio=mm_data.pop("audios"))
+            mm_data = dict(audio=dict(mm_data).pop("audios"))
             mm_kwargs = dict(
                 **mm_kwargs,
                 sampling_rate=feature_extractor.sampling_rate,
@@ -2045,15 +2055,13 @@ def get_generation_prompt(cls, stt_params: SpeechToTextParams) -> PromptType:
         # (which is different from "_") and encode("▁ABC") ignores the first token
         # so the prompt_text is unreliable. However, prompt_token_ids can be used
         # to get prompt_text but it wont have the first token "▁".
-        prompt_text = None
         prompt_token_ids = cls._get_default_prompt_token_ids(
             tokenizer,
             model_config,
             language,
         )
 
         return TokensPrompt(
-            prompt=prompt_text,
             prompt_token_ids=prompt_token_ids,
             multi_modal_data={"audio": (audio, stt_config.sample_rate)},
         )
@@ -2199,7 +2207,9 @@ def embed_multimodal(self, **kwargs: object) -> MultiModalEmbeddings:
         audio_input, seq_lens = self._parse_and_validate_audio_input(**kwargs)
 
         if hasattr(audio_input, "input_features"):
-            out = self.model.get_encoder_outputs(audio_input["input_features"])
+            out = self.model.get_encoder_outputs(
+                audio_input["input_features"], seq_lens
+            )
         else:
             out = self.model.get_encoder_outputs(audio_input, seq_lens)
 
@@ -2219,12 +2229,17 @@ def _parse_and_validate_audio_input(
                 f"Incorrect type of audio features. Got type: {type(input_features)}"
             )
 
+        if not isinstance(length, torch.Tensor):
+            raise ValueError("Audio feature lengths must be a tensor.")
+
         if isinstance(input_features, torch.Tensor):
             seq_lens = length.reshape(-1)
         else:
+            if not is_list_of(input_features, torch.Tensor, check="all"):
+                raise ValueError("All audio features must be tensors.")
             input_features = [
-                feat.to(self.dtype).squeeze(0).transpose(1, 0)
-                for feat in input_features
+                feature.to(self.dtype).squeeze(0).transpose(1, 0)
+                for feature in input_features
             ]
             seq_lens = length.reshape(-1)
             input_features = torch.nn.utils.rnn.pad_sequence(

vllm/model_executor/models/cohere_eagle.py

@@ -58,6 +58,7 @@ def __init__(
         start_layer_id: int = 0,
     ) -> None:
         super().__init__()
+        assert vllm_config.speculative_config is not None
         self.config = vllm_config.speculative_config.draft_model_config.hf_config
         self.quant_config = get_draft_quant_config(vllm_config)
 
@@ -137,6 +138,7 @@ def forward(
 class EagleCohereForCausalLM(CohereForCausalLM):
     def __init__(self, *, vllm_config: VllmConfig, prefix: str = ""):
         nn.Module.__init__(self)
+        assert vllm_config.speculative_config is not None
         self.config = vllm_config.speculative_config.draft_model_config.hf_config
         # Flags checked by the speculative proposer to decide whether to share
         # embed_tokens / lm_head with the target model. Cohere EAGLE checkpoints
@@ -160,7 +162,7 @@ def __init__(self, *, vllm_config: VllmConfig, prefix: str = ""):
     def embed_input_ids(self, input_ids: torch.Tensor) -> torch.Tensor:
         return self.model.embed_input_ids(input_ids)
 
-    def forward(
+    def forward(  # type: ignore[override]
         self,
         input_ids: torch.Tensor,
         positions: torch.Tensor,

vllm/model_executor/models/colmodernvbert.py

@@ -33,6 +33,7 @@
     PromptUpdate,
 )
 from vllm.sequence import IntermediateTensors
+from vllm.tokenizers.hf import HfTokenizer
 from vllm.transformers_utils.configs.colmodernvbert import ColModernVBertConfig
 
 from .interfaces import (
@@ -161,6 +162,7 @@ def _call_hf_processor(
         tok_kwargs: Mapping[str, object],
     ) -> BatchFeature:
         tokenizer = self.info.get_tokenizer()
+        assert isinstance(tokenizer, HfTokenizer)
         text_encoding = tokenizer(
             prompt,
             return_tensors="pt",

vllm/model_executor/models/colpali.py

@@ -25,6 +25,7 @@
 
 from vllm.config import VllmConfig
 from vllm.model_executor.layers.pooler.tokwise import pooler_for_token_embed
+from vllm.model_executor.layers.pooler.tokwise.heads import TokenPoolerHead
 from vllm.model_executor.model_loader.weight_utils import default_weight_loader
 from vllm.multimodal import MULTIMODAL_REGISTRY
 
@@ -232,7 +233,9 @@ def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]:
                         loaded.add(f"custom_text_proj.{param_name}")
 
             # Update pooler projector for the lazy-creation path
-            self.pooler.head.projector = self.custom_text_proj
+            head = self.pooler.head
+            assert isinstance(head, TokenPoolerHead)
+            head.projector = self.custom_text_proj
 
         # Mark pooler projector params as loaded
         if hasattr(self, "pooler") and hasattr(self.pooler, "head"):

vllm/model_executor/models/commandr.py

@@ -154,6 +154,7 @@ def __init__(
         self.max_position_embeddings = getattr(
             config, "model_max_length", None
         ) or getattr(config, "max_position_embeddings", 8192)
+        assert isinstance(self.max_position_embeddings, int)
         self.use_qk_norm = getattr(config, "use_qk_norm", False)
         self.qkv_proj = QKVParallelLinear(
             self.hidden_size,

vllm/model_executor/models/config.py

@@ -9,6 +9,7 @@
     from transformers import PretrainedConfig
 
     from vllm.config import CacheConfig, ModelConfig, VllmConfig
+    from vllm.config.cache import MambaDType
 
 
 logger = init_logger(__name__)
@@ -452,6 +453,7 @@ class JambaForSequenceClassificationConfig(VerifyAndUpdateConfig):
     @staticmethod
     def verify_and_update_model_config(model_config: "ModelConfig") -> None:
         pooler_config = model_config.pooler_config
+        assert pooler_config is not None
         if pooler_config.use_activation is None:
             pooler_config.use_activation = False
 
@@ -518,6 +520,7 @@ def verify_and_update_model_config(model_config: "ModelConfig") -> None:
         config = model_config.hf_config
         config.num_labels = 1
         pooler_config = model_config.pooler_config
+        assert pooler_config is not None
         if pooler_config.logit_mean is None:
             pooler_config.logit_mean = 2.65
 
@@ -540,6 +543,7 @@ def verify_and_update_model_config(model_config: "ModelConfig") -> None:
         if pooling_type is None:
             raise ValueError(f"pool_type {hf_config.pooling!r} not supported")
 
+        assert model_config.pooler_config is not None
         model_config.pooler_config.seq_pooling_type = pooling_type
 
 
@@ -572,10 +576,13 @@ def verify_and_update_model_config(model_config: "ModelConfig") -> None:
         if pooling is None and hasattr(hf_config, "llm_config"):
             pooling = getattr(hf_config.llm_config, "pooling", "avg")
 
+        assert isinstance(pooling, str)
+
         pooling_type = pooling_type_map.get(pooling)
         if pooling_type is None:
             raise ValueError(f"pool_type {pooling!r} not supported")
 
+        assert model_config.pooler_config is not None
         model_config.pooler_config.seq_pooling_type = pooling_type
 
 
@@ -631,7 +638,7 @@ def verify_and_update_config(cls, vllm_config: "VllmConfig") -> None:
 
 
 class NemotronHForCausalLMConfig(VerifyAndUpdateConfig):
-    DEFAULT_MAMBA_SSM_CACHE_DTYPE = "float32"
+    DEFAULT_MAMBA_SSM_CACHE_DTYPE: "MambaDType" = "float32"
     """Only `float32` is known to have no accuracy issues by default."""
 
     @classmethod
@@ -643,7 +650,7 @@ def update_mamba_ssm_cache_dtype(
         `float32` if not specified.
         """
         if cache_config.mamba_ssm_cache_dtype == "auto":
-            mamba_ssm_cache_dtype = getattr(
+            mamba_ssm_cache_dtype: MambaDType = getattr(
                 hf_config, "mamba_ssm_cache_dtype", cls.DEFAULT_MAMBA_SSM_CACHE_DTYPE
             )
             logger.info(
@@ -729,6 +736,7 @@ class Qwen2ForProcessRewardModelConfig(VerifyAndUpdateConfig):
     @staticmethod
     def verify_and_update_model_config(model_config: "ModelConfig") -> None:
         pooler_config = model_config.pooler_config
+        assert pooler_config is not None
 
         if pooler_config.step_tag_id is None:
             pooler_config.step_tag_id = 151651
@@ -738,6 +746,7 @@ class Qwen2ForRewardModelConfig(VerifyAndUpdateConfig):
     @staticmethod
     def verify_and_update_model_config(model_config: "ModelConfig") -> None:
         pooler_config = model_config.pooler_config
+        assert pooler_config is not None
 
         if pooler_config.use_activation is None:
             pooler_config.use_activation = False