[Mypy] Fix typing for J models (#54130)
Signed-off-by: Taneem Ibrahim <taneem.ibrahim@gmail.com>
Taneem Ibrahim committed
Aug 29, 2026 at 07:50 UTC
fd5d3aea9470bb92376eb2d9f5d64cd8f23de31b
3 files changed
+6
-5
tools/pre_commit/mypy.py
-1
@@ -112,7 +112,6 @@ SEPARATE_GROUPS = [
112
]
113
114
EXCLUDE = [
115
- r"vllm/model_executor/models/[jJ]",
115
r"vllm/model_executor/models/[kK]",
116
r"vllm/model_executor/models/[lL]",
117
r"vllm/model_executor/models/[mM]",
vllm/model_executor/models/jamba.py
+1
-1
@@ -120,7 +120,7 @@ class JambaMambaDecoderLayer(nn.Module):
120
model_config: ModelConfig | None = None,
121
cache_config: CacheConfig | None = None,
122
quant_config: QuantizationConfig | None = None,
123
- is_lora_enabled: bool | None = False,
123
+ is_lora_enabled: bool = False,
124
prefix: str = "",
125
**kwargs,
126
) -> None:
vllm/model_executor/models/jina.py
+5
-3
@@ -183,7 +183,9 @@ def _load_adapter(
183
return adapter_config, adapter_weights
184
185
186
-def _build_lora_pairs(adapter_weights: dict) -> dict:
186
+def _build_lora_pairs(
187
+ adapter_weights: dict[str, torch.Tensor],
188
+) -> dict[str, dict[str, torch.Tensor]]:
189
"""Group raw adapter tensors into {base_key: {"A": tensor, "B": tensor}} pairs.
190
191
Transforms adapter keys like:
@@ -191,7 +193,7 @@ def _build_lora_pairs(adapter_weights: dict) -> dict:
193
Into base keys like:
194
layers.0.self_attn.q_proj.weight
195
"""
194
- lora_pairs = defaultdict(dict)
196
+ lora_pairs: defaultdict[str, dict[str, torch.Tensor]] = defaultdict(dict)
197
for key, tensor in adapter_weights.items():
198
clean_key = key
199
if clean_key.startswith("base_model.model."):
@@ -232,7 +234,7 @@ def _load_jina_v5_weights(
234
model: nn.Module, weights: Iterable[tuple[str, torch.Tensor]]
235
) -> set[str]:
236
"""Shared loader: merge the selected task LoRA adapter into the base weights."""
235
- lora_pairs: dict = {}
237
+ lora_pairs: dict[str, dict[str, torch.Tensor]] = {}
238
scaling = 1.0
239
240
result = _load_adapter(model._model_name, model._task, model._revision)