pyink

Author: hengtaoguo

src/maxtext/checkpoint_conversion/utils/param_mapping.py

@@ -1042,7 +1042,7 @@ def QWEN3_5_MAXTEXT_TO_HF_PARAM_MAPPING(config, maxtext_config, scan_layers=Fals
               ): f"model.language_model.layers.{i}.mlp.experts.gate_up_proj",
           }
       )
-  
+
   # Vision mapping for Qwen3.5
   if maxtext_config.use_multimodal and "vision_config" in config:
     vision_config = config["vision_config"]
@@ -1094,15 +1094,11 @@ def QWEN3_5_MAXTEXT_TO_HF_PARAM_MAPPING(config, maxtext_config, scan_layers=Fals
     mapping["params-vision_encoder-Qwen3_5MoeVisionProjector_0-merger-mlp_0-kernel"] = (
         "model.visual.merger.linear_fc1.weight"
     )
-    mapping["params-vision_encoder-Qwen3_5MoeVisionProjector_0-merger-mlp_0-bias"] = (
-        "model.visual.merger.linear_fc1.bias"
-    )
+    mapping["params-vision_encoder-Qwen3_5MoeVisionProjector_0-merger-mlp_0-bias"] = "model.visual.merger.linear_fc1.bias"
     mapping["params-vision_encoder-Qwen3_5MoeVisionProjector_0-merger-mlp_2-kernel"] = (
         "model.visual.merger.linear_fc2.weight"
     )
-    mapping["params-vision_encoder-Qwen3_5MoeVisionProjector_0-merger-mlp_2-bias"] = (
-        "model.visual.merger.linear_fc2.bias"
-    )
+    mapping["params-vision_encoder-Qwen3_5MoeVisionProjector_0-merger-mlp_2-bias"] = "model.visual.merger.linear_fc2.bias"
 
   return mapping
 

src/maxtext/multimodal/processor.py

@@ -224,6 +224,7 @@ def get_bidirectional_mask_vision(config, decoder_input_tokens, is_video: bool =
     bidirectional_mask_vision = decoder_input_tokens == LLAMA4_PATCH_TOKEN
   elif config.model_name in ["qwen3-omni-30b-a3b", "qwen3.5-35b-a3b", "qwen3.5-397b-a17b"]:
     from maxtext.multimodal.processor_qwen3_omni import QwenTokens  # pylint: disable=import-outside-toplevel
+
     tokens = QwenTokens(config)
 
     if is_video:
@@ -238,8 +239,9 @@ def get_bidirectional_mask_audio(config, decoder_input_tokens):
   bidirectional_mask_audio = None
   if config.model_name in ["qwen3-omni-30b-a3b"]:
     from maxtext.multimodal.processor_qwen3_omni import QwenTokens  # pylint: disable=import-outside-toplevel
+
     tokens = QwenTokens(config)
 
     # Create bidirectional_mask for audio token merging
-    bidirectional_mask_audio = (decoder_input_tokens == tokens.audio_pad)
+    bidirectional_mask_audio = decoder_input_tokens == tokens.audio_pad
   return bidirectional_mask_audio

src/maxtext/multimodal/processor_qwen3_omni.py

@@ -96,7 +96,10 @@ def __init__(self, config=None):
     # If config is None, will fall back to default Qwen3-Omni tokens.
     self.model_name = getattr(config, "model_name", None) or self._DEFAULT_MODEL
     # Match by prefix (e.g. "qwen3.5" covers qwen3.5 family), fall back to default.
-    token_config = next((v for k, v in QWEN_SPECIAL_TOKEN_CONFIGS.items() if self.model_name.startswith(k)), QWEN_SPECIAL_TOKEN_CONFIGS[self._DEFAULT_MODEL])
+    token_config = next(
+        (v for k, v in QWEN_SPECIAL_TOKEN_CONFIGS.items() if self.model_name.startswith(k)),
+        QWEN_SPECIAL_TOKEN_CONFIGS[self._DEFAULT_MODEL],
+    )
     self.__dict__.update(token_config)