Spaces:
Runtime error
Runtime error
| """Frozen HuBERT features and a four-label sound-event head. | |
| This path deliberately has no PANNs runtime dependency. It shares the pinned | |
| HuBERT encoder used for audio/context decisions, while keeping sound labels in | |
| the fixed ``sound-v1`` taxonomy. | |
| """ | |
| from __future__ import annotations | |
| from dataclasses import dataclass | |
| import math | |
| from pathlib import Path | |
| import numpy as np | |
| import torch | |
| from torch import nn | |
| from transformers import AutoFeatureExtractor, HubertModel | |
| from .schema import SOUND_LABELS | |
| HUBERT_MODEL = "facebook/hubert-base-ls960" | |
| HUBERT_REVISION = "dba3bb02fda4248b6e082697eee756de8fe8aa8a" | |
| SAMPLE_RATE = 16_000 | |
| WINDOW_SECONDS = 5.0 | |
| HOP_SECONDS = 2.5 | |
| TAIL_SECONDS = 0.5 | |
| FRAME_SECONDS = 0.020 | |
| FEATURE_DIM = 1_536 | |
| class SoundWindows: | |
| waveform: torch.Tensor | |
| intervals: tuple[tuple[float, float], ...] | |
| observed_samples: tuple[int, ...] | |
| def make_sound_windows(waveform: np.ndarray | torch.Tensor) -> SoundWindows: | |
| """Create observed 5-s windows with 2.5-s hop, padding only after audio.""" | |
| wave = torch.as_tensor(np.asarray(waveform, dtype=np.float32)) | |
| if wave.ndim != 1 or wave.numel() == 0: | |
| raise ValueError("waveform must be finite nonempty mono 16-kHz audio") | |
| window, hop = round(WINDOW_SECONDS * SAMPLE_RATE), round(HOP_SECONDS * SAMPLE_RATE) | |
| last_start = max(0, wave.numel() - window) | |
| starts = list(range(0, last_start + 1, hop)) | |
| if not starts or starts[-1] != last_start: | |
| starts.append(last_start) | |
| frames, intervals, observed = [], [], [] | |
| for start in starts: | |
| end = min(start + window, wave.numel()) | |
| frames.append(torch.nn.functional.pad(wave[start:end], (0, window - (end - start)))) | |
| intervals.append((start / SAMPLE_RATE, end / SAMPLE_RATE)) | |
| observed.append(end - start) | |
| return SoundWindows(torch.stack(frames), tuple(intervals), tuple(observed)) | |
| def aggregate_window_scores(scores: torch.Tensor) -> torch.Tensor: | |
| if scores.ndim != 2 or scores.shape[0] == 0 or scores.shape[1] != len(SOUND_LABELS): | |
| raise ValueError("expected [nonempty windows, four-label] scores") | |
| return scores.max(dim=0).values | |
| def aggregate_hubert_hidden(hidden: torch.Tensor, output_lengths: torch.Tensor) -> torch.Tensor: | |
| """Observed-frame mean plus the observed final 0.5-s mean for every window.""" | |
| if hidden.ndim != 3 or output_lengths.ndim != 1 or hidden.shape[0] != output_lengths.numel(): | |
| raise ValueError("hidden state and HuBERT output lengths are inconsistent") | |
| tail_frames = math.ceil(TAIL_SECONDS / FRAME_SECONDS) | |
| values = [] | |
| for row, length in zip(hidden, output_lengths.tolist()): | |
| observed = min(row.shape[0], max(1, int(length))) | |
| values.append(torch.cat((row[:observed].mean(0), row[max(0, observed - tail_frames):observed].mean(0)))) | |
| return torch.stack(values) | |
| def encode_hubert_windows(encoder: HubertModel, feature_extractor: AutoFeatureExtractor, | |
| windows: torch.Tensor, observed_samples: tuple[int, ...] | None = None, | |
| batch_size: int = 32) -> torch.Tensor: | |
| """Return frozen [windows, 1536] features; no padded frames enter pooling.""" | |
| if windows.ndim != 2 or windows.shape[1] != round(WINDOW_SECONDS * SAMPLE_RATE) or batch_size < 1: | |
| raise ValueError("expected 5-s 16-kHz windows and a positive batch size") | |
| if observed_samples is None: | |
| observed_samples = tuple([windows.shape[1]] * windows.shape[0]) | |
| if len(observed_samples) != windows.shape[0] or any(length < 1 or length > windows.shape[1] for length in observed_samples): | |
| raise ValueError("each window needs one valid observed-sample length") | |
| device = next(encoder.parameters()).device | |
| encoded = [] | |
| for start in range(0, windows.shape[0], batch_size): | |
| batch = windows[start:start + batch_size] | |
| lengths = observed_samples[start:start + batch_size] | |
| values = [item[:length].numpy() for item, length in zip(batch.cpu(), lengths)] | |
| inputs = feature_extractor(values, sampling_rate=SAMPLE_RATE, padding=True, | |
| return_attention_mask=True, return_tensors="pt") | |
| mask = inputs["attention_mask"] | |
| hidden = encoder(input_values=inputs["input_values"].to(device), attention_mask=mask.to(device)).last_hidden_state | |
| lengths = encoder._get_feat_extract_output_lengths(mask.sum(dim=1)).to(device) | |
| encoded.append(aggregate_hubert_hidden(hidden, lengths).float().cpu()) | |
| return torch.cat(encoded) | |
| class WindowMILHead(nn.Module): | |
| """Small head shared by cache adapters with any fixed window feature size.""" | |
| def __init__(self, feature_dim: int): | |
| super().__init__() | |
| if feature_dim < 1: | |
| raise ValueError("feature_dim must be positive") | |
| self.feature_dim = feature_dim | |
| self.layers = nn.Sequential( | |
| nn.LayerNorm(feature_dim), nn.Linear(feature_dim, 128), nn.GELU(), nn.Dropout(0.1), | |
| nn.Linear(128, len(SOUND_LABELS)), | |
| ) | |
| def forward(self, window_features: torch.Tensor) -> torch.Tensor: | |
| if window_features.shape[-1] != self.feature_dim: | |
| raise ValueError(f"expected features ending in {self.feature_dim} dimensions") | |
| return self.layers(window_features) | |
| class HubertSoundHead(WindowMILHead): | |
| """The 1536-d shared-HuBERT specialization of ``WindowMILHead``.""" | |
| def __init__(self): | |
| super().__init__(FEATURE_DIM) | |
| def score_window_features(head: WindowMILHead, window_features: torch.Tensor) -> dict[str, torch.Tensor]: | |
| """Score cached/runtime HuBERT window features and apply documented clip max MIL.""" | |
| if window_features.ndim != 2: | |
| raise ValueError("window_features must have shape [windows, 1536]") | |
| device = next(head.parameters()).device | |
| window_scores = torch.sigmoid(head(window_features.to(device))).cpu() | |
| return {"window_scores": window_scores, "clip_scores": aggregate_window_scores(window_scores)} | |
| def padded_mil_logits(head: WindowMILHead, entries: list[dict], device: str | torch.device) -> tuple[torch.Tensor, torch.Tensor]: | |
| """Apply a window head then clipwise max MIL to a cache-adapter batch. | |
| The stable adapter contract is ``entry['features'] -> float [W, D]`` and | |
| ``entry['target'] -> four-label float vector``. Extra provenance such as | |
| intervals or duration is intentionally preserved but not interpreted by | |
| this fixed-label trainer. | |
| """ | |
| if not entries: | |
| raise ValueError("at least one cache entry is required") | |
| maximum = max(entry["features"].shape[0] for entry in entries) | |
| features = torch.zeros((len(entries), maximum, head.feature_dim), device=device) | |
| valid = torch.zeros((len(entries), maximum), dtype=torch.bool, device=device) | |
| targets = torch.stack([torch.as_tensor(entry["target"], dtype=torch.float32) for entry in entries]).to(device) | |
| for index, entry in enumerate(entries): | |
| value = torch.as_tensor(entry["features"], dtype=torch.float32) | |
| if value.ndim != 2 or value.shape[0] < 1 or value.shape[1] != head.feature_dim: | |
| raise ValueError("cache entry features must have shape [nonempty windows, configured feature_dim]") | |
| features[index, :value.shape[0]] = value.to(device) | |
| valid[index, :value.shape[0]] = True | |
| return head(features).masked_fill(~valid.unsqueeze(-1), -torch.inf).amax(dim=1), targets | |
| def load_hubert_sound_head(checkpoint: str | Path, device: str | torch.device = "cpu") -> tuple[HubertSoundHead, torch.Tensor, dict]: | |
| payload = torch.load(checkpoint, map_location="cpu", weights_only=False) | |
| if not isinstance(payload, dict) or payload.get("format") != "whisper-jev-unified-hubert-sound-v1": | |
| raise ValueError("expected a unified HuBERT sound checkpoint") | |
| if tuple(payload.get("taxonomy", ())) != tuple(SOUND_LABELS): | |
| raise ValueError("checkpoint taxonomy differs from the fixed sound taxonomy") | |
| head = HubertSoundHead().to(device).eval() | |
| head.load_state_dict(payload["head"], strict=True) | |
| thresholds = torch.as_tensor(payload["thresholds"], dtype=torch.float32) | |
| return head, thresholds, payload | |