configs refactor 2/3: registry-based audio extractor hooks

src/mobius/_configs/_base.py

@@ -16,7 +16,6 @@
     CodecDecoderConfig,
     CodecEncoderConfig,
     CodePredictorConfig,
-    Gemma4AudioConfig,
     RoPEConfig,
     SpeakerEncoderConfig,
     TTSConfig,
@@ -453,138 +452,14 @@ def _extract_vision_config(config, parent_config, model_type: str) -> dict:
 def _extract_audio_config(config, parent_config, model_type: str) -> dict:
     """Extract audio sub-config from a HuggingFace config.
 
-    Builds an :class:`AudioConfig` (if audio fields are found) and
-    returns a dict of options to merge into :class:`ArchitectureConfig`
-    kwargs.
+    Thin shim that delegates to the per-model registry. The actual
+    hooks live under :mod:`mobius._configs.per_model` and are
+    registered with :mod:`mobius._configs._extractors` at import time.
     """
-    audio_fields: dict = {}
-    audio_processor = getattr(config, "audio_processor", None)
-    if isinstance(audio_processor, dict) and "config" in audio_processor:
-        ac = audio_processor["config"]
-        nemo = ac.get("nemo_conv_settings", {})
-        rel_bias = ac.get("relative_attention_bias_args", {})
-        audio_fields.update(
-            attention_dim=ac.get("attention_dim"),
-            attention_heads=ac.get("attention_heads"),
-            num_blocks=ac.get("num_blocks"),
-            linear_units=ac.get("linear_units"),
-            kernel_size=ac.get("kernel_size"),
-            input_size=ac.get("input_size"),
-            conv_channels=nemo.get("conv_channels", ac.get("attention_dim")),
-            t5_bias_max_distance=rel_bias.get("t5_bias_max_distance"),
-        )
+    from mobius._configs import per_model  # noqa: F401  - side-effect import
+    from mobius._configs._extractors import extract_audio_config as _dispatch
 
-    embd_layer = getattr(config, "embd_layer", None)
-    if isinstance(embd_layer, dict):
-        audio_fields["projection_hidden_size"] = config.hidden_size
-
-    # Phi4MM audio token ID
-    if model_type == "phi4mm":
-        audio_config_dict = getattr(config, "audio_config", None)
-        if audio_config_dict is not None:
-            ac_dict = (
-                audio_config_dict
-                if isinstance(audio_config_dict, dict)
-                else vars(audio_config_dict)
-            )
-            audio_fields["token_id"] = ac_dict.get("audio_token_id")
-
-    speech_lora = getattr(config, "speech_lora", None)
-    if speech_lora is not None:
-        audio_fields["lora"] = (
-            speech_lora if isinstance(speech_lora, dict) else vars(speech_lora)
-        )
-
-    # Qwen3-ASR audio config (from thinker_config)
-    thinker_config_source = parent_config or config
-    hf_thinker_config = getattr(thinker_config_source, "thinker_config", None)
-    if hf_thinker_config is not None:
-        tc = (
-            hf_thinker_config
-            if not isinstance(hf_thinker_config, dict)
-            else type("TC", (), hf_thinker_config)()
-        )
-        hf_audio_config = getattr(tc, "audio_config", None)
-        if hf_audio_config is not None:
-            ac = (
-                hf_audio_config
-                if not isinstance(hf_audio_config, dict)
-                else type("AC", (), hf_audio_config)()
-            )
-            audio_fields.update(
-                d_model=getattr(ac, "d_model", None),
-                encoder_layers=getattr(ac, "encoder_layers", None),
-                encoder_attention_heads=getattr(ac, "encoder_attention_heads", None),
-                encoder_ffn_dim=getattr(ac, "encoder_ffn_dim", None),
-                num_mel_bins=getattr(ac, "num_mel_bins", None),
-                max_source_positions=getattr(ac, "max_source_positions", None),
-                downsample_hidden_size=getattr(ac, "downsample_hidden_size", None),
-                output_dim=getattr(ac, "output_dim", None),
-                activation_function=getattr(ac, "activation_function", "gelu"),
-                n_window=getattr(ac, "n_window", None),
-                n_window_infer=getattr(ac, "n_window_infer", None),
-            )
-        # Special tokens from thinker config
-        audio_fields["audio_token_id"] = getattr(tc, "audio_token_id", None)
-        audio_fields["audio_start_token_id"] = getattr(tc, "audio_start_token_id", None)
-        audio_fields["audio_end_token_id"] = getattr(tc, "audio_end_token_id", None)
-        audio_fields["classify_num"] = getattr(tc, "classify_num", None)
-
-    # Gemma4 audio config (from composite audio_config sub-config).
-    # model_type may be "gemma4_text" when build() resolves to the text sub-config;
-    # check parent_config to catch that case.
-    parent_model_type = getattr(parent_config, "model_type", "") if parent_config else ""
-    if model_type in ("gemma4", "gemma4_text") or parent_model_type == "gemma4":
-        composite = parent_config or config
-        hf_audio_config = getattr(composite, "audio_config", None)
-        if hf_audio_config is not None:
-            ac = (
-                hf_audio_config
-                if not isinstance(hf_audio_config, dict)
-                else type("AC", (), hf_audio_config)()
-            )
-            subsampling = getattr(ac, "subsampling_conv_channels", None)
-            return {
-                "audio": Gemma4AudioConfig(
-                    num_layers=getattr(ac, "num_hidden_layers", 12),
-                    hidden_size=getattr(ac, "hidden_size", 1024),
-                    subsampling_conv_channels=(
-                        list(subsampling) if subsampling is not None else None
-                    ),
-                    use_causal_chunked_attn=getattr(ac, "use_causal_chunked_attn", False),
-                    output_dim=getattr(ac, "output_dim", None),
-                    output_proj_dims=getattr(ac, "output_proj_dims", None),
-                    audio_token_id=getattr(composite, "audio_token_id", None),
-                )
-            }
-
-    # SenseVoice / FunASR-style: encoder config lives under "encoder_conf"
-    # and the mel-spec frontend under "frontend_conf". Top-level holds
-    # input_size + vocab_size. model_type is "sensevoice".
-    if model_type == "sensevoice":
-        encoder_conf = getattr(config, "encoder_conf", None) or {}
-        frontend_conf = getattr(config, "frontend_conf", None) or {}
-        if isinstance(encoder_conf, dict) and encoder_conf:
-            audio_fields.update(
-                attention_dim=encoder_conf.get("output_size"),
-                attention_heads=encoder_conf.get("attention_heads"),
-                num_blocks=encoder_conf.get("num_blocks"),
-                tp_num_blocks=encoder_conf.get("tp_blocks"),
-                linear_units=encoder_conf.get("linear_units"),
-                kernel_size=encoder_conf.get("kernel_size"),
-                input_size=getattr(config, "input_size", None),
-            )
-            if isinstance(frontend_conf, dict):
-                audio_fields["num_mel_bins"] = frontend_conf.get("n_mels")
-
-    # Build AudioConfig sub-config if any audio fields are set
-    has_audio = any(v is not None for v in audio_fields.values())
-
-    result: dict = {}
-    if has_audio:
-        result["audio"] = AudioConfig(**audio_fields)
-
-    return result
+    return _dispatch(config, parent_config, model_type)
 
 
 @dataclasses.dataclass

src/mobius/_configs/_extractors.py

@@ -0,0 +1,142 @@
+# Copyright (c) Microsoft Corporation.
+# Licensed under the MIT License.
+
+"""Plugin-style registry for sub-config extractors.
+
+Each ``extract_*`` function used to be a single mega-switch inside
+:mod:`mobius._configs._base` with a chain of ``if model_type == "..."``
+branches. That meant every new architecture had to edit a shared file
+and risk merge conflicts with unrelated work.
+
+This module replaces those switches with a tiny registry. Hooks are
+plain functions registered via the :func:`register_audio_hook` /
+:func:`register_vision_hook` decorators. Each hook is invoked on every
+extraction and is responsible for guarding its own applicability
+(typically by checking ``model_type`` or for the presence of a specific
+HuggingFace field). A hook may:
+
+* mutate ``fields`` to contribute key/value pairs into the default
+  sub-config that the dispatcher will instantiate at the end, or
+* return a fully-formed ``dict`` (e.g. ``{"audio": Gemma4AudioConfig(...)}``)
+  to short-circuit — skipping all subsequent hooks and the default
+  instantiation. Use this when a model needs a non-default sub-config
+  subclass.
+
+New models live in :mod:`mobius._configs.per_model`. Importing that
+package is what populates the registries (each module registers its
+own hooks at import time).
+"""
+
+from __future__ import annotations
+
+from collections.abc import Callable
+from typing import Any
+
+Hook = Callable[[Any, Any, str, dict], dict | None]
+
+# Each registry entry is ``(model_type_filter, hook)`` where the filter is
+# either ``None`` (always run) or a frozenset of model_type strings.
+_AUDIO_HOOKS: list[tuple[frozenset[str] | None, Hook]] = []
+_VISION_HOOKS: list[tuple[frozenset[str] | None, Hook]] = []
+
+
+def _make_register(registry: list) -> Callable:
+    """Build a decorator that supports both bare and parameterised usage."""
+
+    def register(*model_types):
+        """Register a hook in *registry*.
+
+        Two usages are supported:
+
+        * Bare decorator — runs for every model_type. Use this for default
+          hooks that pull a generic HuggingFace field common to many models.
+
+          .. code-block:: python
+
+              @register_audio_hook
+              def _default(config, parent, mt, fields): ...
+
+        * Parameterised decorator — runs only when ``model_type`` matches one
+          of the supplied strings. The dispatcher filters before invocation
+          so hook bodies don't need to repeat the ``if model_type != ...``
+          guard. Hooks that *also* need to inspect ``parent_config`` (e.g.
+          Gemma4's text-config short-circuit) can still do that check inside.
+
+          .. code-block:: python
+
+              @register_audio_hook("phi4mm")
+              def _phi4mm(config, parent, mt, fields): ...
+
+              @register_audio_hook("gemma4", "gemma4_text")
+              def _gemma4(config, parent, mt, fields): ...
+        """
+        # Bare decorator: ``@register`` with a single callable arg
+        if (
+            len(model_types) == 1
+            and callable(model_types[0])
+            and not isinstance(model_types[0], str)
+        ):
+            fn = model_types[0]
+            registry.append((None, fn))
+            return fn
+
+        types_set = frozenset(model_types) if model_types else None
+
+        def deco(fn: Hook) -> Hook:
+            registry.append((types_set, fn))
+            return fn
+
+        return deco
+
+    return register
+
+
+register_audio_hook = _make_register(_AUDIO_HOOKS)
+register_vision_hook = _make_register(_VISION_HOOKS)
+
+
+def _run(hooks: list, config, parent_config, model_type: str, fields: dict):
+    """Apply each hook whose filter matches ``model_type``.
+
+    Short-circuits on the first hook that returns a non-None dict.
+    """
+    for filter_set, hook in hooks:
+        if filter_set is not None and model_type not in filter_set:
+            continue
+        result = hook(config, parent_config, model_type, fields)
+        if result is not None:
+            return result
+    return None
+
+
+def extract_audio_config(config, parent_config, model_type: str) -> dict:
+    """Run every applicable audio hook and assemble the result.
+
+    Each hook either contributes to ``fields`` (which become kwargs for
+    :class:`AudioConfig` at the end) or returns a dict that short-circuits
+    the dispatcher.
+    """
+    from mobius._configs._sub_configs import AudioConfig
+
+    fields: dict = {}
+    short_circuit = _run(_AUDIO_HOOKS, config, parent_config, model_type, fields)
+    if short_circuit is not None:
+        return short_circuit
+    if any(v is not None for v in fields.values()):
+        return {"audio": AudioConfig(**fields)}
+    return {}
+
+
+def extract_vision_config(config, parent_config, model_type: str) -> dict:
+    """Run every registered vision hook and assemble the result.
+
+    Vision hooks differ from audio hooks: there is no default
+    :class:`VisionConfig` autoinstantiation — vision sub-configs are
+    constructed by an always-applied "default" hook so that other hooks
+    can override its output.
+    """
+    fields: dict = {}
+    short_circuit = _run(_VISION_HOOKS, config, parent_config, model_type, fields)
+    if short_circuit is not None:
+        return short_circuit
+    return fields