Update custom model files, README, and requirements

asr_pipeline.py

@@ -1,5 +1,6 @@
 from typing import Any
 
+import numpy as np
 import torch
 import transformers
 
@@ -9,6 +10,211 @@ except ImportError:
     from asr_modeling import ASRModel  # type: ignore[no-redef]
 
 
+class ForcedAligner:
+    """Lazy-loaded forced aligner for word-level timestamps."""
+
+    _instance = None
+    _model = None
+    _tokenizer = None
+
+    @classmethod
+    def get_instance(cls, device: str = "cuda"):
+        if cls._model is None:
+            from ctc_forced_aligner import load_alignment_model
+
+            dtype = torch.float16 if device == "cuda" else torch.float32
+            cls._model, cls._tokenizer = load_alignment_model(device, dtype=dtype)
+        return cls._model, cls._tokenizer
+
+    @classmethod
+    def align(
+        cls,
+        audio: np.ndarray,
+        text: str,
+        sample_rate: int = 16000,
+        language: str = "eng",
+        batch_size: int = 16,
+    ) -> list[dict]:
+        """Align transcript to audio and return word-level timestamps.
+
+        Args:
+            audio: Audio waveform as numpy array
+            text: Transcript text to align
+            sample_rate: Audio sample rate (default 16000)
+            language: ISO-639-3 language code (default "eng" for English)
+            batch_size: Batch size for alignment model
+
+        Returns:
+            List of dicts with 'word', 'start', 'end' keys
+        """
+        from ctc_forced_aligner import (
+            generate_emissions,
+            get_alignments,
+            get_spans,
+            postprocess_results,
+            preprocess_text,
+        )
+
+        device = "cuda" if torch.cuda.is_available() else "cpu"
+        model, tokenizer = cls.get_instance(device)
+
+        # Convert audio to tensor
+        if isinstance(audio, np.ndarray):
+            audio_tensor = torch.from_numpy(audio).to(model.dtype).to(model.device)
+        else:
+            audio_tensor = audio.to(model.dtype).to(model.device)
+
+        # Ensure 1D
+        if audio_tensor.dim() > 1:
+            audio_tensor = audio_tensor.squeeze()
+
+        # Generate emissions
+        emissions, stride = generate_emissions(model, audio_tensor, batch_size=batch_size)
+
+        # Preprocess text
+        tokens_starred, text_starred = preprocess_text(text, romanize=True, language=language)
+
+        # Get alignments
+        segments, scores, blank_token = get_alignments(emissions, tokens_starred, tokenizer)
+
+        # Get spans
+        spans = get_spans(tokens_starred, segments, blank_token)
+
+        # Get word timestamps
+        word_timestamps = postprocess_results(text_starred, spans, stride, scores)
+
+        # Convert to simple format
+        return [{"word": w["word"], "start": w["start"], "end": w["end"]} for w in word_timestamps]
+
+
+class SpeakerDiarizer:
+    """Lazy-loaded speaker diarization using pyannote-audio."""
+
+    _pipeline = None
+
+    @classmethod
+    def get_instance(cls, hf_token: str | None = None):
+        """Get or create the diarization pipeline.
+
+        Args:
+            hf_token: HuggingFace token with access to pyannote models.
+                     Can also be set via HF_TOKEN environment variable.
+        """
+        if cls._pipeline is None:
+            import os
+
+            from pyannote.audio import Pipeline
+
+            token = hf_token or os.environ.get("HF_TOKEN")
+            cls._pipeline = Pipeline.from_pretrained(
+                "pyannote/speaker-diarization-3.1",
+                use_auth_token=token,
+            )
+
+            # Move to GPU if available
+            if torch.cuda.is_available():
+                cls._pipeline.to(torch.device("cuda"))
+
+        return cls._pipeline
+
+    @classmethod
+    def diarize(
+        cls,
+        audio: np.ndarray | str,
+        sample_rate: int = 16000,
+        num_speakers: int | None = None,
+        min_speakers: int | None = None,
+        max_speakers: int | None = None,
+        hf_token: str | None = None,
+    ) -> list[dict]:
+        """Run speaker diarization on audio.
+
+        Args:
+            audio: Audio waveform as numpy array or path to audio file
+            sample_rate: Audio sample rate (default 16000)
+            num_speakers: Exact number of speakers (if known)
+            min_speakers: Minimum number of speakers
+            max_speakers: Maximum number of speakers
+            hf_token: HuggingFace token for pyannote models
+
+        Returns:
+            List of dicts with 'speaker', 'start', 'end' keys
+        """
+        pipeline = cls.get_instance(hf_token)
+
+        # Prepare audio input
+        if isinstance(audio, np.ndarray):
+            # pyannote expects {"waveform": tensor, "sample_rate": int}
+            waveform = torch.from_numpy(audio).unsqueeze(0)  # Add channel dim
+            if waveform.dim() == 1:
+                waveform = waveform.unsqueeze(0)
+            audio_input = {"waveform": waveform, "sample_rate": sample_rate}
+        else:
+            # File path
+            audio_input = audio
+
+        # Run diarization
+        diarization_args = {}
+        if num_speakers is not None:
+            diarization_args["num_speakers"] = num_speakers
+        if min_speakers is not None:
+            diarization_args["min_speakers"] = min_speakers
+        if max_speakers is not None:
+            diarization_args["max_speakers"] = max_speakers
+
+        diarization = pipeline(audio_input, **diarization_args)
+
+        # Convert to simple format
+        segments = []
+        for turn, _, speaker in diarization.itertracks(yield_label=True):
+            segments.append({
+                "speaker": speaker,
+                "start": turn.start,
+                "end": turn.end,
+            })
+
+        return segments
+
+    @classmethod
+    def assign_speakers_to_words(
+        cls,
+        words: list[dict],
+        speaker_segments: list[dict],
+    ) -> list[dict]:
+        """Assign speaker labels to words based on timestamp overlap.
+
+        Args:
+            words: List of word dicts with 'word', 'start', 'end' keys
+            speaker_segments: List of speaker dicts with 'speaker', 'start', 'end' keys
+
+        Returns:
+            Words list with 'speaker' key added to each word
+        """
+        for word in words:
+            word_mid = (word["start"] + word["end"]) / 2
+
+            # Find the speaker segment that contains this word's midpoint
+            best_speaker = None
+            for seg in speaker_segments:
+                if seg["start"] <= word_mid <= seg["end"]:
+                    best_speaker = seg["speaker"]
+                    break
+
+            # If no exact match, find closest segment
+            if best_speaker is None and speaker_segments:
+                min_dist = float("inf")
+                for seg in speaker_segments:
+                    seg_mid = (seg["start"] + seg["end"]) / 2
+                    dist = abs(word_mid - seg_mid)
+                    if dist < min_dist:
+                        min_dist = dist
+                        best_speaker = seg["speaker"]
+
+            word["speaker"] = best_speaker
+
+        return words
+
+
 class ASRPipeline(transformers.AutomaticSpeechRecognitionPipeline):
     """ASR Pipeline for audio-to-text transcription."""
 
@@ -24,6 +230,137 @@ class ASRPipeline(transformers.AutomaticSpeechRecognitionPipeline):
         super().__init__(
             model=model, feature_extractor=feature_extractor, tokenizer=tokenizer, **kwargs
         )
+        self._current_audio = None
+        self._return_timestamps = False
+        self._return_speakers = False
+        self._diarization_params = {}
+
+    def _sanitize_parameters(self, **kwargs):
+        """Intercept our custom parameters before parent class validates them."""
+        # Extract our custom parameters before parent sees them
+        self._return_timestamps = kwargs.pop("return_timestamps", False)
+        self._return_speakers = kwargs.pop("return_speakers", False)
+        self._diarization_params = {
+            "num_speakers": kwargs.pop("num_speakers", None),
+            "min_speakers": kwargs.pop("min_speakers", None),
+            "max_speakers": kwargs.pop("max_speakers", None),
+            "hf_token": kwargs.pop("hf_token", None),
+        }
+
+        # return_speakers requires return_timestamps
+        if self._return_speakers:
+            self._return_timestamps = True
+
+        # Now let parent sanitize remaining params
+        return super()._sanitize_parameters(**kwargs)
+
+    def __call__(
+        self,
+        inputs,
+        **kwargs,
+    ):
+        """Transcribe audio with optional word-level timestamps and speaker diarization.
+
+        Args:
+            inputs: Audio input (file path, dict with array/sampling_rate, etc.)
+            return_timestamps: If True, return word-level timestamps using forced alignment
+            return_speakers: If True, return speaker labels for each word
+            num_speakers: Exact number of speakers (if known, for diarization)
+            min_speakers: Minimum number of speakers (for diarization)
+            max_speakers: Maximum number of speakers (for diarization)
+            hf_token: HuggingFace token for pyannote models (or set HF_TOKEN env var)
+            **kwargs: Additional arguments passed to the pipeline
+
+        Returns:
+            Dict with 'text' key, 'words' key if return_timestamps=True,
+            and speaker labels on words if return_speakers=True
+        """
+        # Extract our params before super().__call__ (which will also call _sanitize_parameters)
+        return_timestamps = kwargs.pop("return_timestamps", False)
+        return_speakers = kwargs.pop("return_speakers", False)
+        diarization_params = {
+            "num_speakers": kwargs.pop("num_speakers", None),
+            "min_speakers": kwargs.pop("min_speakers", None),
+            "max_speakers": kwargs.pop("max_speakers", None),
+            "hf_token": kwargs.pop("hf_token", None),
+        }
+
+        if return_speakers:
+            return_timestamps = True
+
+        # Store audio for timestamp alignment and diarization
+        if return_timestamps or return_speakers:
+            self._current_audio = self._extract_audio(inputs)
+
+        # Run standard transcription
+        result = super().__call__(inputs, **kwargs)
+
+        # Add timestamps if requested
+        if return_timestamps and self._current_audio is not None:
+            text = result.get("text", "")
+            if text:
+                try:
+                    words = ForcedAligner.align(
+                        self._current_audio["array"],
+                        text,
+                        sample_rate=self._current_audio.get("sampling_rate", 16000),
+                    )
+                    result["words"] = words
+                except Exception as e:
+                    result["words"] = []
+                    result["timestamp_error"] = str(e)
+            else:
+                result["words"] = []
+
+        # Add speaker diarization if requested
+        if return_speakers and self._current_audio is not None:
+            try:
+                # Run diarization
+                speaker_segments = SpeakerDiarizer.diarize(
+                    self._current_audio["array"],
+                    sample_rate=self._current_audio.get("sampling_rate", 16000),
+                    **{k: v for k, v in diarization_params.items() if v is not None},
+                )
+                result["speaker_segments"] = speaker_segments
+
+                # Assign speakers to words
+                if result.get("words"):
+                    result["words"] = SpeakerDiarizer.assign_speakers_to_words(
+                        result["words"],
+                        speaker_segments,
+                    )
+            except Exception as e:
+                result["speaker_segments"] = []
+                result["diarization_error"] = str(e)
+
+        # Clean up
+        self._current_audio = None
+
+        return result
+
+    def _extract_audio(self, inputs) -> dict | None:
+        """Extract audio array from various input formats."""
+        import librosa
+
+        if isinstance(inputs, dict):
+            if "array" in inputs:
+                return {
+                    "array": inputs["array"],
+                    "sampling_rate": inputs.get("sampling_rate", 16000),
+                }
+            if "raw" in inputs:
+                return {
+                    "array": inputs["raw"],
+                    "sampling_rate": inputs.get("sampling_rate", 16000),
+                }
+        elif isinstance(inputs, str):
+            # File path - load audio
+            audio, sr = librosa.load(inputs, sr=16000)
+            return {"array": audio, "sampling_rate": sr}
+        elif isinstance(inputs, np.ndarray):
+            return {"array": inputs, "sampling_rate": 16000}
+
+        return None
 
     def preprocess(self, inputs, **preprocess_params):
         # Handle dict with "array" key (from datasets)

requirements.txt

@@ -1,14 +1,6 @@
-# Use latest compatible versions
-gradio
-transformers>=4.57.1
-torch
-soundfile
-librosa
-peft
-truecase
+# Core dependencies for tiny-audio model inference
+# This file is pushed to HuggingFace for model repository
 
-# Forced alignment for word-level timestamps
-ctc-forced-aligner @ git+https://github.com/MahmoudAshraf97/ctc-forced-aligner.git
-
-# Speaker diarization
-pyannote-audio>=3.1.0
+# Transformers - main library for model loading and inference
+transformers>=4.57.0
+truecase