Training in progress - step 15000

asr_pipeline.py

@@ -1,6 +1,5 @@
 from typing import Any
 
-import numpy as np
 import torch
 import transformers
 
@@ -10,211 +9,6 @@ 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."""
 
@@ -230,125 +24,6 @@ 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
-        """
-        # Store audio for timestamp alignment and diarization
-        if self._return_timestamps or self._return_speakers:
-            self._current_audio = self._extract_audio(inputs)
-
-        # Run standard transcription
-        result = super().__call__(inputs, **kwargs)
-
-        # Add timestamps if requested
-        if self._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 self._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 self._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
-        if self._return_timestamps or self._return_speakers:
-            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)