Upload modeling_prosody.py with huggingface_hub

modeling_prosody.py

@@ -0,0 +1,341 @@
+"""Prosody to text: f0 and intensity in, words out.
+
+Two signals -- interpolated f0 and intensity at 200 Hz -- through a conformer
+to a mel spectrogram, and the mel through a fine-tuned Whisper. No phonetic
+content at any point: the model sees only how the pitch moved and how loud
+it was.
+
+Loaded the way any custom architecture on the Hub is:
+
+    from transformers import AutoModel
+    model = AutoModel.from_pretrained(
+        "Daporte/prosody2text-weights", trust_remote_code=True
+    )
+    print(model.transcribe("utterance.wav").text)
+
+A note on the shape of this class. It subclasses `PreTrainedModel` so that
+`from_pretrained`, the `AutoModel` API and the Hub's download metrics all
+work, but it is a pipeline rather than a single network: a conformer and a
+`WhisperForConditionalGeneration` held alongside each other, loaded from
+their own checkpoints rather than from one state dict. `forward` runs them
+in sequence under `no_grad`, so this is not trainable as it stands -- it
+exists to run the published pipeline, which is what the paper reports.
+"""
+
+from __future__ import annotations
+
+from pathlib import Path
+
+import numpy as np
+import torch
+
+from transformers import PreTrainedModel
+
+from .configuration_prosody import ProsodyToTextConfig
+from .conformer import Conformer
+from .features import ProsodyFeatures
+from .speech import extract_features
+from .whisper_decoder import Whisper, looks_like_fallback
+
+
+class ProsodyResult:
+    """What one utterance produces.
+
+    Carries the intermediate stages as well as the words, because the whole
+    point is being able to see them: a transcription that looks wrong is
+    usually a contour that looks wrong first.
+    """
+
+    def __init__(
+        self,
+        features: ProsodyFeatures,
+        mel: np.ndarray,
+        ranked: list[tuple[str, float]],
+    ) -> None:
+        self.features = features
+        self.mel = mel
+        self.ranked = ranked
+
+    @property
+    def text(self) -> str:
+        """The chosen hypothesis, empty if nothing was transcribed."""
+        return self.ranked[0][0] if self.ranked else ""
+
+    @property
+    def candidates(self) -> list[str]:
+        """The whole beam. The top hypothesis is often a narrow winner on an
+        input this uncertain, so the runners-up say more about what the model
+        recovered than the single best string does."""
+        return [text for text, _ in self.ranked]
+
+    @property
+    def fallback(self) -> bool:
+        """Whether the result is Whisper's no-speech boilerplate rather than
+        a transcription. Trained on captions, it answers unintelligible
+        input with "Thank you." and shouted auto-caption filler, which reads
+        as a result but is not one."""
+        return bool(self.text) and looks_like_fallback(self.text)
+
+    def __repr__(self) -> str:
+        return (
+            f"ProsodyResult(text={self.text!r}, "
+            f"frames={self.features.num_frames}, mel={self.mel.shape})"
+        )
+
+
+class ProsodyToTextModel(PreTrainedModel):
+    """The published pipeline, as a Hub model."""
+
+    config_class = ProsodyToTextConfig
+    # Nothing here is trained by this class, and the weights are loaded by
+    # the two components rather than from a single state dict.
+    base_model_prefix = "prosody"
+    supports_gradient_checkpointing = False
+    _supports_sdpa = False
+
+    def __init__(self, config: ProsodyToTextConfig) -> None:
+        super().__init__(config)
+        self.config = config
+        self._conformer: Conformer | None = None
+        self._whisper: Whisper | None = None
+        self._root: Path | None = None
+
+    # ---- loading ---------------------------------------------------------
+
+    @classmethod
+    def from_pretrained(cls, model_id, *args, **kwargs):  # noqa: ANN001
+        """Build the pipeline from a repo or a local directory.
+
+        Overridden rather than relying on the base implementation: there is
+        no single state dict to load. Each half is its own checkpoint in its
+        own subfolder, so both are resolved from the snapshot and loaded by
+        their own code.
+        """
+        from transformers import AutoConfig
+
+        half = kwargs.pop("half", False)
+        cpu = kwargs.pop("cpu", False)
+        load_whisper = kwargs.pop("whisper", True)
+        config = kwargs.pop("config", None)
+        # Consumed so a caller can pass the usual from_pretrained arguments
+        # without them reaching code that has no use for them.
+        kwargs.pop("torch_dtype", None)
+        kwargs.pop("device_map", None)
+
+        root = Path(model_id)
+        if not root.is_dir():
+            from huggingface_hub import snapshot_download
+
+            patterns = ["*.py", "*.json", config_file(model_id)]
+            patterns += ["conformer/*", "tokenizer/*"]
+            if load_whisper:
+                patterns.append("whisper/*")
+            root = Path(snapshot_download(
+                repo_id=str(model_id),
+                allow_patterns=[p for p in patterns if p],
+                revision=kwargs.pop("revision", None),
+                token=kwargs.pop("token", None),
+            ))
+
+        if config is None:
+            config = AutoConfig.from_pretrained(
+                str(root), trust_remote_code=True
+            )
+
+        model = cls(config)
+        model._root = root
+        model._load(root, half=half, cpu=cpu, whisper=load_whisper)
+        return model
+
+    def _load(
+        self, root: Path, half: bool, cpu: bool, whisper: bool
+    ) -> None:
+        warn_if_driver_hidden(whisper)
+
+        self._conformer = Conformer(
+            root / self.config.conformer_subfolder, cpu=cpu
+        )
+        if not whisper:
+            return
+
+        tokenizer = root / self.config.tokenizer_subfolder
+        self._whisper = Whisper(
+            root / self.config.whisper_subfolder,
+            half=half,
+            num_beams=self.config.num_beams,
+            tokenizer_dir=(
+                tokenizer if (tokenizer / "tokenizer.json").exists() else None
+            ),
+        )
+
+    @property
+    def device_summary(self) -> str:
+        """Where each half is running, which is worth printing: the whole
+        difference between a minute and an hour is whether Whisper found the
+        GPU."""
+        parts = []
+        if self._conformer is not None:
+            parts.append(f"conformer: {self._conformer.provider}")
+        if self._whisper is not None:
+            parts.append(f"whisper: {self._whisper.provider}")
+        return ", ".join(parts) or "nothing loaded"
+
+    # ---- running ---------------------------------------------------------
+
+    def forward(
+        self,
+        f0_hz: np.ndarray | torch.Tensor | None = None,
+        intensity_db: np.ndarray | torch.Tensor | None = None,
+        transcribe: bool = True,
+        **kwargs,
+    ) -> ProsodyResult:
+        """The two contours in, the mel and the ranked beam out.
+
+        Physical units -- Hz and dB on the 200 Hz grid -- not normalised. The
+        normalisation to the model's own scale happens here so a caller
+        cannot apply it twice, which produces a mel that looks plausible and
+        is wrong.
+        """
+        if f0_hz is None or intensity_db is None:
+            raise ValueError("f0_hz and intensity_db are both required")
+
+        features = ProsodyFeatures(
+            f0_hz=as_array(f0_hz), intensity_db=as_array(intensity_db)
+        )
+        return self.run(features, transcribe=transcribe)
+
+    def run(
+        self, features: ProsodyFeatures, transcribe: bool = True
+    ) -> ProsodyResult:
+        """Features to words, with the mel kept on the way through."""
+        if self._conformer is None:
+            raise RuntimeError(
+                "the conformer is not loaded; build this with "
+                "from_pretrained()"
+            )
+        if features.num_frames == 0:
+            return ProsodyResult(
+                features, np.zeros((0, self.config.mel_bins), np.float32), []
+            )
+
+        mel = self._conformer.predict(features.normalised())
+        ranked: list[tuple[str, float]] = []
+        if transcribe and self._whisper is not None and len(mel):
+            ranked = self._whisper.hypotheses(
+                mel,
+                language=self.config.language,
+                # None under `constrain`, so the budget comes from the
+                # utterance's duration rather than the flat ceiling.
+                max_tokens=(
+                    None if self.config.constrain
+                    else self.config.max_new_tokens
+                ),
+                no_repeat_ngram_size=self.config.no_repeat_ngram_size,
+                constrain=self.config.constrain,
+            )
+        return ProsodyResult(features, mel, ranked)
+
+    def transcribe(
+        self,
+        audio,
+        sampling_rate: int | None = None,
+        trim: bool = True,
+    ) -> ProsodyResult:
+        """Audio to words, for the common case.
+
+        `audio` is a path, or samples with `sampling_rate` given. `trim` cuts
+        the quiet lead-in and tail a live recording carries -- yaapt finds a
+        pitch in room noise, so without it several hundred milliseconds of
+        nothing reach the model. Dataset clips usually arrive trimmed
+        already; pass False to match the published runs exactly.
+        """
+        if isinstance(audio, (str, Path)):
+            # In a try block on purpose. trust_remote_code scans a modeling
+            # file's imports statically and refuses to load unless every one
+            # is installed -- and it skips Try nodes precisely so a module
+            # can mark a dependency optional. Reading a path is the only
+            # thing here that needs soundfile; a caller passing arrays, as
+            # the drawing demo does, should not be made to install it.
+            try:
+                import soundfile
+            except ImportError as error:  # pragma: no cover
+                raise ImportError(
+                    "reading audio from a path needs soundfile "
+                    "(pip install soundfile); pass samples and "
+                    "sampling_rate instead to avoid it"
+                ) from error
+
+            samples, sampling_rate = soundfile.read(
+                str(audio), dtype="float64", always_2d=False
+            )
+            if getattr(samples, "ndim", 1) > 1:
+                samples = samples.mean(axis=1)
+        else:
+            if sampling_rate is None:
+                raise ValueError(
+                    "sampling_rate is required when passing samples"
+                )
+            samples = as_array(audio)
+
+        return self.run(extract_features(samples, sampling_rate, trim=trim))
+
+
+# ---- helpers -------------------------------------------------------------
+
+
+def as_array(value) -> np.ndarray:  # noqa: ANN001
+    if isinstance(value, torch.Tensor):
+        value = value.detach().cpu().numpy()
+    array = np.asarray(value, dtype=np.float64)
+    return array.squeeze() if array.ndim > 1 else array
+
+
+def config_file(model_id) -> str:  # noqa: ANN001
+    """`config.json`, named separately so the download pattern reads clearly."""
+    return "config.json"
+
+
+# Directories where a GPU driver may sit without being on the loader path.
+#
+# A pip-installed torch finds libcuda.so.1 through the dynamic loader, and on
+# most systems the driver package puts it where the loader already looks. Not
+# everywhere: NixOS and some container images keep it outside that path, so
+# torch reports "Found no NVIDIA driver on your system" on a machine whose
+# nvidia-smi works perfectly.
+DRIVER_DIRS = (
+    "/run/opengl-driver/lib",          # NixOS
+    "/usr/lib/x86_64-linux-gnu",       # Debian, Ubuntu
+    "/usr/lib64",                      # Fedora, RHEL
+    "/usr/local/nvidia/lib64",         # some container images
+)
+
+
+def warn_if_driver_hidden(whisper: bool) -> None:
+    """Say when a GPU is present but unreachable.
+
+    Falling back to CPU silently is the wrong default for something twenty
+    times slower: across a few thousand utterances that is the difference
+    between an hour and a day, and the cause is one environment variable.
+    """
+    if not whisper:
+        return
+    if torch.cuda.is_available():
+        return
+
+    import ctypes
+
+    try:
+        ctypes.CDLL("libcuda.so.1")
+        return          # loadable, so CUDA is unavailable for another reason
+    except OSError:
+        pass
+
+    for directory in DRIVER_DIRS:
+        if (Path(directory) / "libcuda.so.1").exists():
+            print(
+                f"warning: a GPU driver is installed but not on the library "
+                f"path, so this will run on CPU. Re-run with:\n"
+                f"    LD_LIBRARY_PATH={directory}:$LD_LIBRARY_PATH",
+                flush=True,
+            )
+            return