Fix/ser wav2vec2 regression head#511
Merged
Merged
Conversation
Merge alpha: YAMNet backend, auto-resampling, pdoc cleanup
Promote scalene profiling wrapper to main (alpha→main)
![gemini-code-assist[bot]](https://avatars.githubusercontent.com/in/956858?s=60&v=4){kind=small}
Contributor
There was a problem hiding this comment.
Code Review
This pull request introduces a custom backend for models using the Wav2Vec2ForSpeechClassification architecture to ensure regression heads are loaded correctly rather than randomly initialized. Key additions include a custom regression head and model class, along with logic to detect and route these specific models during inference. Review feedback suggests replacing broad exception handling with more specific types and refactoring duplicated configuration loading logic into a shared helper function.
satra
reviewed
May 9, 2026
satra
left a comment
satra
left a comment
Collaborator
There was a problem hiding this comment.
Thanks for tracking this down — the head-structure mismatch diagnosis is solid and the new _Wav2Vec2EmotionModel reproduces the original audeering class faithfully. A few items worth addressing before merge:
CI: cpu-tests (3.12) is failing on this commit (job 75099350753). The PR description ticks "All new and existing tests passed" — please investigate; test_speech_emotion_recognition_continuous exercises this exact path now.
Correctness (see inline):
- Silent sampling-rate hazard —
Wav2Vec2Processordoes not resample. id2label[i]canKeyErrorin the raw-dict fallback path because HF JSON keys are strings.
Test coverage: the existing continuous-SER test only checks label names, which would have passed before this fix too (random head still produced correct labels with bogus scores). Worth adding:
assert set(labels) == {"arousal", "valence", "dominance"}- scores in
[0, 1] - regression guard that not all scores ≈ 0.33 (
any(abs(s - 0.33) > 0.05 for s in scores))
Minor cleanup (see inline): redundant _AC alias; Tuple import unused if switched to lowercase tuple; class factory rebuilds on every cache miss.
Otherwise looks good — once the CI failure is understood and the resample/id2label edges are handled, happy to see this land.
Generated by Claude Code
- Auto-resample audio to the processor's expected sampling rate to
avoid silent correctness bugs when callers pass non-16 kHz input.
- Coerce id2label keys to int (HF stores them as strings in JSON) and
fall back to class_{i} so the raw-config path doesn't KeyError.
- Memoize the dynamically built Wav2Vec2 emotion model class so
isinstance checks stay stable across calls.
- Drop redundant `AutoConfig as _AC` alias and unused `Tuple` import.
- Strengthen the continuous-SER test with score-range and regression
assertions that would catch the random-head bug this PR fixes.
Address Gemini review on PR #513: the previous slice-only logic could produce a labels list shorter than the model output if num_labels in the config doesn't match the actual head size. Pad with class_{i} so AudioClassificationResult always sees matching lengths.
The 'any(abs(s - 0.33) > 0.05 for s in scores)' check would fail if the test fixture happened to produce audeering outputs all in [0.28, 0.38]. Keep only the label-set and [0, 1] range checks, which are robust to any fixture content.
Three follow-on issues surfaced when running the new continuous-SER backend
end-to-end against audeering/wav2vec2-large-robust-12-ft-emotion-msp-dim:
1. AutoConfig fallback raised TypeError because config.json's own
"model_type" key collided with the positional arg in
AutoConfig.for_model("wav2vec2", ...). Pop it before splat.
2. The same checkpoints ship vocab_size: null, which trips
huggingface_hub>=1.0 strict-dataclass validators when Wav2Vec2Config is
instantiated. The regression head doesn't use vocab_size, so drop the
key and let the dataclass default apply.
3. _Wav2Vec2EmotionModel called self.init_weights(), but
transformers>=5.0's from_pretrained finalizer reads
self.all_tied_weights_keys, which is populated by post_init. Switch to
post_init() so loading succeeds.
Additionally swap Wav2Vec2Processor.from_pretrained for
Wav2Vec2FeatureExtractor.from_pretrained, since the processor path
re-invokes AutoConfig.from_pretrained for the tokenizer (re-tripping the
vocab_size validator) and regression inference doesn't need a tokenizer.
Mypy cleanups along the way: cast config to Wav2Vec2Config at the
Wav2Vec2Model construction site, and route feature_extractor.sampling_rate
through getattr to drop the attr-defined warning on Wav2Vec2Processor.
https://claude.ai/code/session_01DyXMQNsF7F5Ze8j3GGVw66
Fix Wav2Vec2 emotion model initialization and sampling rate handling
Wires the audeering Wav2Vec2ForSpeechClassification fix from PR #513 into the discrete-SER path so the ehcalabres RAVDESS checkpoint stops emitting ~1/8-uniform softmax outputs: * Generalize the custom emotion classifier head to support both final-layer attribute names actually shipped in the wild: ``classifier.out_proj.*`` (audeering MSP-Dim) and ``classifier.output.*`` (ehcalabres). Build the matching nn.Module per layout and cache by name. * Detect the head layout cheaply via the checkpoint's safetensors index or bin shard manifest (no full weight download); fall back to a small registry for single-bin checkpoints like ehcalabres. * Dispatcher now routes any Wav2Vec2 emotion checkpoint with a 2-layer dense+linear head through ``_classify_wav2vec2_speech_cls_ser``, applying softmax for discrete and leaving raw logits for continuous. * Discrete SER test asserts max(scores) > 0.2 to detect random-head regressions, plus that scores form a softmax distribution. Wraps every SpeechBrain ``from_hparams`` call in ``speechbrain_loading_cwd(savedir)`` so CWD-relative ``save_path`` lobes inside hyperparams.yaml (e.g. ``save_path: wav2vec2_checkpoints`` on ``speechbrain/emotion-recognition-wav2vec2-IEMOCAP``) resolve under ``{HF_HOME}/senselab_cache/speechbrain/<safe_key>/`` instead of dropping 700+MB into the user's working tree. Verified end-to-end against two HAPPY recordings: ehcalabres now returns ``happy: 0.49`` (vs. ~0.13 uniform before), and CWD stays clean after a full SER pipeline run. https://claude.ai/code/session_01DyXMQNsF7F5Ze8j3GGVw66
…or, KeyError) Addresses gemini-code-assist's review on PR #511 about broad-except. The actual exception observed in the audeering MSP-Dim case is huggingface_hub.errors.StrictDataclassError (parent of StrictDataclassFieldValidationError, raised on vocab_size: null with huggingface_hub>=1.0). Standard parsing failures fall under ValueError / TypeError / KeyError. OSError is intentionally excluded — that means the network or filesystem is broken, and the fallback path needs the same network hop, so it can't recover. The tuple is constructed with a try/except ImportError to stay compatible with huggingface_hub<1.0 (where StrictDataclassError doesn't exist). Verified by re-running test_speech_emotion_recognition_continuous against audeering/wav2vec2-large-robust-12-ft-emotion-msp-dim — the StrictDataclassFieldValidationError is correctly caught by the narrowed tuple and the raw-config-dict fallback runs as before. https://claude.ai/code/session_01DyXMQNsF7F5Ze8j3GGVw66
…ation plan Phase 1 of specs/20260509-204500-ser-silent-failure-guards/plan.md. The custom Wav2Vec2 emotion-head path in `_classify_wav2vec2_speech_cls_ser` trusts the dispatcher peek to have correctly matched the checkpoint's head layout. When that assumption is wrong (new encoder family, unrecognized final-layer attribute name, or single-bin checkpoint not in the registry) transformers silently random-initializes the head and inference returns ~uniform softmax — the exact bug PR #511 is fixing, just one routing branch upstream. Convert that to a loud `RuntimeError`. After `from_pretrained` returns we inspect `loading_info` for any `classifier.*` entry in `missing_keys` (weights absent, would be randomly initialized) **or** `mismatched_keys` (weights present but wrong shape — equally a silent-random failure that the original missing-keys check would miss; flagged by architect review). We also assert the final layer's `out_features` equals `config.num_labels` / `len(id2label)`, catching the case where head and config disagree on the number of classes. Adds a parametrised unit test that mocks `EmotionModel.from_pretrained` to return each broken-load shape and asserts the diagnostic raises with a matching message. Also commits the full remediation plan (specs/.../plan.md) enumerating Phases 2-7 (standard-pipeline missing-key warning, encoder-family generalization for HuBERT/WavLM, head-registry richening, problem_type based discrete/continuous detection, resampling in the standard pipeline, diagnostics CLI). Plan was self-audited and reviewed by the Plan agent before this commit; its recommendations are folded in (mismatched_keys coverage promoted into Phase 1, dynamic HeadSpec cut from Phase 4 as premature, encoder-attribute-name registry promoted from a Phase 3 footnote to a load-bearing task). https://claude.ai/code/session_01DyXMQNsF7F5Ze8j3GGVw66
4 tasks
Implements the remaining phases from
specs/20260509-204500-ser-silent-failure-guards/plan.md, layered on top
of Phase 1's loud-fail guard:
- Phase 4 (head registry): introduce ``_HeadEntry`` dataclass holding
``final_layer`` / ``activation`` / ``dropout_field``; promote
``_KNOWN_WAV2VEC2_EMOTION_HEADS`` (str→str) to ``_KNOWN_HEAD_LAYOUTS``
(str→_HeadEntry). Generalize ``_make_wav2vec2_regression_head_class``
to ``_make_regression_head_class(head: _HeadEntry)`` with a small
``_ACTIVATIONS`` table. Cache key now includes activation/dropout_field
so future entries don't collide.
- Phase 3 (encoder-family generalization): ``_BASE_REGISTRY`` maps
``config.model_type`` → ``(PreTrainedModel base name, encoder attr
name)`` for wav2vec2/hubert/wavlm/wav2vec2-bert. Resolved lazily so
imports only fire when the custom-head path runs. ``_emotion_head_kind``
now returns ``(model_type, _HeadEntry)`` tuples; ``_make_emotion_model_class``
threads the right encoder class and attribute name. The Phase-1 guard
catches mis-registered families: any wrong attr-name will surface as
"every classifier.* key missing" with a clear ``RuntimeError``.
- Phase 5 (problem_type detection): ``_resolve_apply_softmax`` prefers
``config.problem_type`` over the keyword-based label heuristic. A
regression head with non-AVD axes (e.g. ``["energy","pleasantness"]``)
no longer gets softmax misapplied. Falls back to ``_get_ser_type`` for
legacy checkpoints that omit ``problem_type``.
- Phase 2 (standard-pipeline missing-key warning):
``_get_hf_audio_classification_pipeline`` now loads the model
explicitly via ``AutoModelForAudioClassification.from_pretrained(...,
output_loading_info=True)`` and threads the loaded instance into
``pipeline(model=..., feature_extractor=...)``. ``_check_head_loaded_cleanly``
warns on suspect missing/mismatched head keys
(``classifier.``/``head.``/``score.``/``out_proj.``/``projector.``).
``SENSELAB_STRICT_HEAD_LOAD=1`` promotes the warning to a hard error.
- Phase 6 (resampling in standard pipeline): verified already implemented
at ``classify_audios_with_transformers`` lines 141-142; no code change
required. Plan updated to reflect the discovery.
- Phase 7 (docs + diagnostics CLI):
- ``doc.md`` enumerates the three dispatch paths (custom in-process
head / subprocess venv / standard pipeline) and points at the CLI.
- New ``python -m
senselab.audio.tasks.classification.speech_emotion_recognition
--probe <repo_id>`` reports model_type, architectures, dispatch
decision, head source, and emits advisories. Pure config / manifest
inspection — no weights loaded.
Backwards-compat shims kept in place:
- ``_make_wav2vec2_emotion_model_class(final_layer)`` (delegates to
``_make_emotion_model_class("wav2vec2", _HeadEntry(final_layer=...))``)
- ``_wav2vec2_emotion_head_kind`` (returns just the final-layer name from
the new ``_emotion_head_kind`` tuple).
- ``final_layer=`` kwarg on ``_classify_wav2vec2_speech_cls_ser``
short-circuits to ``_HeadEntry(final_layer=...)``.
Verified locally: all 8 classification tests pass (1 GPU-skip), audeering
happy path passes, ehcalabres CLI probe correctly identifies its head
from ``_KNOWN_HEAD_LAYOUTS`` registry, audeering CLI probe peeks the
manifest, superb correctly routes to the standard pipeline with an
advisory.
Test coverage gaps (deferred per plan):
- Phase 2 unit test for ``_check_head_loaded_cleanly`` (covered
end-to-end by audeering/ehcalabres tests for now).
- Phase 5 unit test for ``_resolve_apply_softmax`` (no fixture model
with problem_type=regression + confounding labels).
- Phase 4 numerical-regression test on existing checkpoints.
https://claude.ai/code/session_01DyXMQNsF7F5Ze8j3GGVw66
Two small fixes from the second-pass architect review:
1. Make ``_BASE_REGISTRY`` a 3-tuple ``(base_pretrained_cls_name,
encoder_model_cls_name, encoder_attr_name)`` instead of deriving the
encoder class name via ``base_name.replace("PreTrainedModel", "Model")``.
The string-mutation worked for current Wav2Vec2/HuBERT/WavLM/W2V-BERT
but would silently break the day a family ships an encoder with a
different naming convention. Explicit > clever.
Updated both indexers (``_emotion_head_kind`` in api.py and the
probe CLI's status print) to read attr name from index ``[2]``.
2. Probe CLI: when ``AutoConfig.from_pretrained`` fails (e.g. audeering's
``vocab_size: null`` tripping the strict validator), make the
diagnostic note explicit that the dispatcher's raw-config-dict
fallback handles this — the routing decision and advisories that
follow still work correctly because ``getattr(None, ...)`` short-
circuits the ``config.model_type`` checks.
Verified: existing guard tests pass; ``uv run python -m
senselab.audio.tasks.classification.speech_emotion_recognition --probe
audeering/wav2vec2-large-robust-12-ft-emotion-msp-dim`` still produces a
correct routing report end-to-end.
https://claude.ai/code/session_01DyXMQNsF7F5Ze8j3GGVw66
CI pre-commit hook runs mypy with --extra-checks which flagged the inferred type for ``model = HFModel(...)`` in __main__.py. Add the explicit annotation. cpu-tests already pass; this is the only blocker. https://claude.ai/code/session_01DyXMQNsF7F5Ze8j3GGVw66
…2 unit tests
The strict reviewer identified 5 critical defects against the user's
"zero tolerance for silent failures" criterion. All addressed in this
commit:
D1. ``wav2vec2-bert`` removed from ``_BASE_REGISTRY``.
The family consumes log-mel ``input_features`` (via
``SeamlessM4TFeatureExtractor``) rather than raw ``input_values``.
The loader hardcodes the latter and accesses ``inputs["input_values"]``
so any wav2vec2-bert checkpoint would have produced ``KeyError`` at
best, silent garbage at worst. Add a per-family preprocessing
adapter before re-introducing.
D2. Phase-1 guard now covers encoder-key misses.
Previous filter was ``classifier.*`` only, which meant a wrong
``_BASE_REGISTRY`` entry (wrong encoder attribute name) would leave
every encoder weight missing yet pass the guard — silent garbage
output. The guard now treats any missing key under either
``classifier.`` or ``{encoder_attr}.`` as suspect, with a small
whitelist of HF post_init buffers (``masked_spec_embed``).
D3. ``_resolve_apply_softmax`` peeks ``architectures`` on
``problem_type``-absent checkpoints. Any model declaring
``Wav2Vec2ForSpeechClassification`` (audeering signature) is treated
as regression. Reduces the silent-corruption window for regression
heads with English-emotion-word labels.
D4. ``_peek_head_final_layer`` raises when ``classifier.dense.weight``
is present but no recognized final-layer attribute matches.
Previously returned None and let the standard pipeline silently
random-initialize the head. The error message lists the actual
classifier keys so a maintainer can extend ``_FINAL_LAYER_NAMES``
or ``_KNOWN_HEAD_LAYOUTS``.
D5. ``SENSELAB_STRICT_HEAD_LOAD`` default inverted to ``=1`` (raise).
Warn-only meant downstream consumers like
``classify_audios`` windowed segmentation kept aggregating
corrupted scores. Set ``SENSELAB_STRICT_HEAD_LOAD=0`` to demote to
a warning if you intentionally want a partial-head load (e.g.
fine-tuning warm-start).
T203. Three unit tests for ``_check_head_loaded_cleanly``: parametrised
raise on suspect missing/mismatched-shape keys (default strict),
warn-only mode under ``SENSELAB_STRICT_HEAD_LOAD=0``, and silent
pass on clean loads (encoder-buffer misses don't trip).
Plan checkboxes T203 / T304 / T502 promoted from deferred-shortcut to
actually-shipped. Audeering happy-path test still passes; mypy clean
with ``--extra-checks``.
https://claude.ai/code/session_01DyXMQNsF7F5Ze8j3GGVw66
Reviewer round 3 verified D1-D5 + T203 fixes are functionally correct, but flagged three pieces of user-facing text that still contradict the new behavior — themselves a silent-assumption violation per the user's "no silent failures and assumptions beneath those" criterion: * ``doc.md`` listed ``wav2vec2-bert`` as a supported encoder for the custom-head path (it isn't; D1 dropped it from ``_BASE_REGISTRY``) and described ``SENSELAB_STRICT_HEAD_LOAD=1`` as a way to "promote the warning to a RuntimeError" — but ``=1`` is now the default (D5). A user who reads the doc, doesn't set the env var, and expects warn-mode would be surprised by a hard fail. * ``__main__.py`` probe CLI advisory had the same inverted message about the env var. * ``api.py:503`` dispatcher comment + ``api.py:144`` encoder-attribute comment still mentioned ``wav2vec2-bert``. This commit: 1. Removes ``wav2vec2-bert`` from the doc's supported-encoder list, adds an explicit "intentionally not supported" sentence with the ``SeamlessM4TFeatureExtractor`` reason, so users know why and what it would take to add. 2. Inverts the env-var description in both the doc and the CLI to match the strict-by-default semantics. 3. Drops stale ``wav2vec2-bert`` mentions from the api.py comments. No code-path or behavior change. Verified clean: ruff, ruff format, mypy --extra-checks all pass. https://claude.ai/code/session_01DyXMQNsF7F5Ze8j3GGVw66
When the strict reviewer's second pass forced T203 and T502 from deferred-shortcut to actually-shipped, I appended the new ``[x]`` text and kept the original ``[ ]`` lines tagged "legacy text, kept for traceability". On re-read those duplicates just made the plan look like there were unfinished items when there weren't. Drop the duplicates; the ``[x]`` entries are sufficient. Remaining ``[ ]`` items after cleanup: T305 — attention-mask shim for batched HuBERT/WavLM (don't optimize for unknowns) T403 — numerical regression test against fixed scores (Phase-1 guard + bounded-score covers) T602 — sampling-rate test (Phase 6 was a no-change phase, end-to-end coverage exists) All three are honest deferrals, not shortcuts. https://claude.ai/code/session_01DyXMQNsF7F5Ze8j3GGVw66
Implements all phases of the SER silent-failure-guards plan, with two strict reviewer passes folding in 5 critical defects + 3 stale-doc fixes. Phases: - 1: Loud-fail post-load assertion in the custom-head path (covers classifier and encoder keys). - 2: Standard-pipeline head-load guard (RuntimeError by default; SENSELAB_STRICT_HEAD_LOAD=0 to demote). - 3: Encoder-family generalization (wav2vec2/hubert/wavlm; wav2vec2-bert deferred). - 4: HeadEntry registry refactor. - 5: problem_type detection with architecture peek fallback. - 6: Sampling-rate resampling (verified pre-existing). - 7: Doc updates + diagnostic --probe CLI. See PR #514 description for the full reviewer history and explicit caveats. https://claude.ai/code/session_01DyXMQNsF7F5Ze8j3GGVw66
Collaborator
Author
|
seems reasonable. the main issues that were pointed out by claude:
|
Collaborator
|
Thanks for the careful review — all five points addressed (or, in #5's case, defended) in #515 which stacks on this branch:
Generated by Claude Code |
… lock + registry smoke tests (#515) Stacks on top of PR #511 to address @wilke0818's review. Changes: 1. **CWD race fix**: `speechbrain_loading_cwd` now serializes via a module-level `threading.Lock` so concurrent threads can't race on the process-global `os.chdir`. Docstring documents the trade-off (use processes, not threads, for parallel SpeechBrain init). 2. **AutoConfig deduplication**: introduces `_load_config_cached(model)` returning `(typed_config_or_None, raw_dict_or_None)`, memoized per-(path, revision). The four call sites (`_get_ser_type`, `_emotion_head_kind`, `_resolve_apply_softmax`, `_classify_wav2vec2_speech_cls_ser`) now share one network round-trip instead of four. Type annotated with `PretrainedConfig`. 3. **Per-key locking** (not a single global lock): mirrors the strategy used by `dependencies.ensure_hf_model` (`_HeartbeatLock` per `(repo_id, revision)`). Two threads loading the *same* model dedupe to one round-trip; two threads loading *different* models proceed in parallel. In-memory `threading.Lock` per key here, since the memo is per-process (cross-process is already handled by HF_HOME cache and `ensure_hf_model`). 4. **`_BASE_REGISTRY` smoke tests**: three parametrised unit tests verify `_resolve_base("wav2vec2"|"hubert"|"wavlm")` resolve to real transformers classes — catches typos at test-collection time without downloading model weights. 5. **Test isolation fix**: autouse fixture clears the module-level caches (`_config_memo`, `_config_memo_locks`, `_wav2vec2_emotion_models`) between tests. Previously the parametrised `test_wav2vec2_speech_cls_ser_raises_on_random_head` and `_raises_on_shape_mismatch` would silently serve a cached MagicMock from a prior test's monkeypatch instead of re-invoking their own; tests only passed because the mocks happened to agree. Verified locally: ruff, format, mypy --extra-checks clean. 18 fast tests pass including 4 new ones. CI: ✅ cpu-tests, pre-commit, build-and-deploy all green on `be77c5d`. https://claude.ai/code/session_01DyXMQNsF7F5Ze8j3GGVw66
Collaborator
|
Follow-up #515 just landed on this branch (squash-merged as 834be0c). The diff on this PR now reflects all five fixes from your review:
CI on the latest commit is green; ready when you are. Generated by Claude Code |
Collaborator
Author
|
seems good |
Collaborator
|
adding the gpu tests also to check the venv stuff (i think the image does have cuda 12.8 so this should still pass there). |
This was referenced May 15, 2026
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
3 participants
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.
Description
Fixes continuous SER, at least for a subset of models by using a custom class that matches the expected class
Related Issue(s)
None
Motivation and Context
Use of the continuous SER by collaborators
How Has This Been Tested?
Manually
Screenshots (if appropriate):
Types of changes
Checklist:
Version
Published prerelease version:
1.3.1-alpha.30Changelog
🐛 Bug Fix
Authors: 3