[Mypy] Fix typing for J models (#54130)

Signed-off-by: Taneem Ibrahim <taneem.ibrahim@gmail.com>

tools/pre_commit/mypy.py

@@ -112,7 +112,6 @@ SEPARATE_GROUPS = [
 ]
 
 EXCLUDE = [
-    r"vllm/model_executor/models/[jJ]",
     r"vllm/model_executor/models/[kK]",
     r"vllm/model_executor/models/[lL]",
     r"vllm/model_executor/models/[mM]",

vllm/model_executor/models/jamba.py

@@ -120,7 +120,7 @@ class JambaMambaDecoderLayer(nn.Module):
         model_config: ModelConfig | None = None,
         cache_config: CacheConfig | None = None,
         quant_config: QuantizationConfig | None = None,
-        is_lora_enabled: bool | None = False,
+        is_lora_enabled: bool = False,
         prefix: str = "",
         **kwargs,
     ) -> None:

vllm/model_executor/models/jina.py

@@ -183,7 +183,9 @@ def _load_adapter(
     return adapter_config, adapter_weights
 
 
-def _build_lora_pairs(adapter_weights: dict) -> dict:
+def _build_lora_pairs(
+    adapter_weights: dict[str, torch.Tensor],
+) -> dict[str, dict[str, torch.Tensor]]:
     """Group raw adapter tensors into {base_key: {"A": tensor, "B": tensor}} pairs.
 
     Transforms adapter keys like:
@@ -191,7 +193,7 @@ def _build_lora_pairs(adapter_weights: dict) -> dict:
     Into base keys like:
         layers.0.self_attn.q_proj.weight
     """
-    lora_pairs = defaultdict(dict)
+    lora_pairs: defaultdict[str, dict[str, torch.Tensor]] = defaultdict(dict)
     for key, tensor in adapter_weights.items():
         clean_key = key
         if clean_key.startswith("base_model.model."):
@@ -232,7 +234,7 @@ def _load_jina_v5_weights(
     model: nn.Module, weights: Iterable[tuple[str, torch.Tensor]]
 ) -> set[str]:
     """Shared loader: merge the selected task LoRA adapter into the base weights."""
-    lora_pairs: dict = {}
+    lora_pairs: dict[str, dict[str, torch.Tensor]] = {}
     scaling = 1.0
 
     result = _load_adapter(model._model_name, model._task, model._revision)