configs refactor 2/3: registry-based audio extractor hooks - #310
Merged
Conversation
_configs.py grew to 2591 lines (9 sub-configs, 19 per-model configs,
3 mega-extractor switch functions, a 530-line ArchitectureConfig.
from_transformers). This is the first of three mechanical refactors
that carve it into a scalable package layout. No behavior change in
this PR — every public name is still importable from mobius._configs.
src/mobius/_configs/
├── __init__.py # re-exports everything that was in _configs.py
├── _sub_configs.py # pure-data dataclasses (RoPE/Vision/Audio/Codec/TTS)
├── _quantization.py # QuantizationConfig + from_transformers
└── _base.py # BaseModelConfig, ArchitectureConfig, per-model
# subclasses, and the _extract_* helpers
Follow-up PRs in this series:
Part 2/3 — convert the _extract_audio_config / _extract_vision_config
model_type switches into a decorator-registered dispatch
so new models add a file under per_model/ instead of a
branch in the central function.
Part 3/3 — move per-model config subclasses (Gemma2Config,
MllamaConfig, NemotronHConfig, ...) into per_model/ and
carve up ArchitectureConfig.from_transformers.
Tests: 2769 passed (full src/ + tests/build_graph_test.py + cli_test.py).
Ruff: clean.
Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com>
Signed-off-by: Justin Chu <justinchu@microsoft.com>
Replaces the 135-line model_type switch in _extract_audio_config with
a tiny plugin registry. Each model now contributes a hook in its own
file under src/mobius/_configs/per_model/, mutating an audio_fields
dict or short-circuiting with a fully-formed sub-config. Adding a new
audio-capable architecture no longer requires editing _base.py.
New module:
src/mobius/_configs/_extractors.py # register_audio_hook + dispatch
Per-model files (one per former branch):
per_model/_audio_default.py # audio_processor, embd_layer, speech_lora
per_model/_phi4mm_audio.py # phi4mm audio_token_id
per_model/_qwen3_asr_audio.py # thinker_config.audio_config + token ids
per_model/_gemma4_audio.py # short-circuit to Gemma4AudioConfig
per_model/_sensevoice_audio.py # encoder_conf + frontend_conf mapping
_extract_audio_config in _base.py shrinks to a 5-line shim that
triggers the per_model side-effect import and calls the dispatcher.
No behavior change: every existing audio-capable model still produces
the same AudioConfig from the same HF config.
Vision-side conversion (also a mega-switch) is intentionally deferred
to a follow-up so reviewers can verify the registry pattern on the
smaller surface first.
Tests:
pytest src/ tests/build_graph_test.py tests/cli_test.py -n auto
2769 passed, 41 skipped
pytest tests/arch_validation_test.py -k 'sensevoice_small or phi4mm
or qwen3_asr or gemma4'
15 passed
Ruff: clean.
Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com>
Signed-off-by: Justin Chu <justinchu@microsoft.com>
Performance Comparison
|
Codecov Report❌ Patch coverage is 📢 Thoughts on this report? Let us know! |
Before: every hook body had to open with
if model_type != "phi4mm":
return None
That's noisy and easy to get wrong when copy-pasting a hook for a new
model. Make the decorator accept an optional filter and have the
dispatcher skip hooks whose filter doesn't match the current
model_type:
@register_audio_hook # runs for every model_type
def _default(...): ...
@register_audio_hook("phi4mm") # phi4mm only
def _phi4mm(...): ...
@register_audio_hook("gemma4", "gemma4_text")
def _gemma4_simple_case(...): ...
Hooks that also need to look at parent_config (Gemma4 audio,
Qwen3-ASR via thinker_config) keep the manual if-guard inside the
body, since the decorator filter can't reach beyond model_type.
Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com>
Signed-off-by: Justin Chu <justinchu@microsoft.com>
Two concerns covered: 1. Mechanism — register_audio_hook works as a bare decorator and as a parameterised factory; the dispatcher applies model_type filters; a hook can short-circuit by returning a dict. 2. Cross-contamination — every per-model audio hook is verified to NOT fire for unrelated model_types. For filtered hooks (phi4mm, sensevoice), the dispatcher's filter must skip non-matching types even when the input config has matching shape (e.g. an audio_config dict that would normally trigger phi4mm). Bare hooks (default, qwen3_asr, gemma4) must return cleanly on a vanilla text-model config. Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com> Signed-off-by: Justin Chu <justinchu@microsoft.com>
Signed-off-by: Justin Chu <justinchuby@users.noreply.github.com>
justinchuby
added a commit
that referenced
this pull request
May 27, 2026
## Part 3 of 3 — stacked on #310 Mirrors part 2/3 for the vision side: replaces the 185-line `_extract_vision_config` switch with a per-model hook registry. The vision dispatcher carries slightly more logic than audio because vision also lifts a fixed set of shared fields (`image_token_id`, `mrope_section`, `spatial_merge_size`, …) up to the top-level of the returned dict. ## New per_model files | File | Purpose | |---|---| | `_vision_default.py` | canonical HF `vision_config` + LoRA / embd_layer / mrope_section | | `_phi4mm_vision.py` | hard-coded SigLIP encoder dims (not in JSON) | | `_hunyuan_vl_mot_vision.py` | InternViT-style ViT, flat config.json | | `_internvl_vision.py` | default `image_token_id` for InternVL chain | ## End state After this 3-PR series, **every model-type switch in the original `_extract_*_config` functions has been moved into discoverable per-model files**. New audio- or vision-capable models add a single file under `src/mobius/_configs/per_model/` instead of editing `_base.py`. The follow-on refactor (out of scope here) is to do the same for per-model config subclasses (`Gemma2Config`, `MllamaConfig`, etc.) and to carve up `ArchitectureConfig.from_transformers` itself. ## Tests | Suite | Result | |---|---| | `pytest src/ tests/build_graph_test.py tests/cli_test.py -n auto` | 2769 passed, 41 skipped | | `pytest tests/arch_validation_test.py -k 'phi4mm or hunyuan_vl_mot or qwen2_vl or qwen3_vl or gemma4'` | 30 passed | | Ruff | clean | --------- Signed-off-by: Justin Chu <justinchu@microsoft.com> Signed-off-by: Justin Chu <justinchuby@users.noreply.github.com> Signed-off-by: justinchuby <11205048+justinchuby@users.noreply.github.com> Signed-off-by: justinchuby <11205safetensors048+justinchuby@users.noreply.github.com> Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com> Co-authored-by: justinchuby <11205048+justinchuby@users.noreply.github.com> Co-authored-by: justinchuby <11205safetensors048+justinchuby@users.noreply.github.com>
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
1 participant
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
Part 2 of 3 — stacked on #309
Replaces the 135-line
model_typeswitch in_extract_audio_configwith a tiny plugin registry. Each model now contributes a hook in its own file undersrc/mobius/_configs/per_model/, mutating anaudio_fieldsdict or short-circuiting with a fully-formed sub-config.Adding a new audio-capable architecture no longer requires editing
_base.py.Layout
_extract_audio_configin_base.pyshrinks to a 5-line shim that triggers theper_modelside-effect import and calls the dispatcher.No behavior change: every existing audio-capable model still produces the same
AudioConfigfrom the same HF config.Hook protocol
Deferred
Vision-side conversion (also a mega-switch) is intentionally deferred to a follow-up so reviewers can verify the registry pattern on the smaller surface first.
Tests
pytest src/ tests/build_graph_test.py tests/cli_test.py -n autopytest tests/arch_validation_test.py -k 'sensevoice_small or phi4mm or qwen3_asr or gemma4'