Training in progress - step 15000

asr_modeling.py

@@ -1,8 +1,5 @@
 from pathlib import Path
-from typing import Optional, Union, Generator, NamedTuple
-
-import threading
-from concurrent import futures
+from typing import Optional, Union
 
 import torch
 import torch.nn as nn
@@ -14,7 +11,6 @@ from transformers import (
     AutoTokenizer,
     PreTrainedModel,
     Wav2Vec2FeatureExtractor,
-    TextIteratorStreamer,
 )
 from transformers.generation.utils import (
     GenerateBeamDecoderOnlyOutput,
@@ -29,17 +25,6 @@ except ImportError:
     from asr_config import ASRConfig  # type: ignore[no-redef]
 
 
-class StreamChunk(NamedTuple):
-    """A chunk of streaming transcription text."""
-    text: str
-
-
-class StreamStats(NamedTuple):
-    """Statistics about the streaming inference."""
-    input_tokens: int
-    output_tokens: int
-
-
 class SwiGLU(nn.Module):
     def __init__(self, in_features, hidden_features, out_features, bias=False, dropout=0.0):
         super().__init__()
@@ -133,12 +118,8 @@ class ASRModel(PreTrainedModel):
             return WhisperFeatureExtractor.from_pretrained(
                 audio_model_id,
                 feature_size=num_mel_bins,
-                do_normalize=True,
             )
-        return Wav2Vec2FeatureExtractor.from_pretrained(
-            audio_model_id,
-            do_normalize=True,
-        )
+        return Wav2Vec2FeatureExtractor.from_pretrained(audio_model_id)
 
     @classmethod
     def from_pretrained(cls, pretrained_model_name_or_path, *args, **kwargs):
@@ -706,145 +687,6 @@ class ASRModel(PreTrainedModel):
 
         return generated_ids[:, prompt_length:]
 
-    @torch.no_grad()
-    def generate_stream(
-        self,
-        input_values: Optional[torch.Tensor] = None,
-        input_features: Optional[torch.Tensor] = None,
-        system_prompt: Optional[str] = None,
-        user_prompt: Optional[str] = None,
-        task: Optional[str] = None,
-        max_new_tokens: Optional[int] = None,
-        temperature: Optional[float] = None,
-        **generate_kwargs,
-    ) -> Generator[Union[StreamChunk, StreamStats], None, None]:
-        """
-        Generate transcription in streaming mode, yielding text chunks as they're generated.
-
-        Args:
-            input_values: Audio input tensor for non-Whisper models
-            input_features: Audio input tensor for Whisper models
-            system_prompt: System prompt override
-            user_prompt: User prompt override
-            task: Task type (transcribe, describe, emotion, continue)
-            max_new_tokens: Maximum tokens to generate
-            temperature: Sampling temperature
-            **generate_kwargs: Additional generation parameters
-
-        Yields:
-            StreamChunk: Text chunks as they're generated
-            StreamStats: Final statistics (input_tokens, output_tokens)
-        """
-        audio_inputs = input_values if input_values is not None else input_features
-        if audio_inputs is None:
-            raise ValueError("input_values or input_features must be provided for generation")
-
-        # Encode audio once and prepare prompt
-        audio_embeds = self._encode_audio(audio_inputs)
-        batch_size = audio_embeds.shape[0]
-        device = audio_embeds.device
-
-        if batch_size > 1:
-            raise ValueError("Streaming generation only supports batch_size=1")
-
-        if system_prompt is None:
-            system_prompt = self.system_prompt
-
-        if user_prompt is None:
-            user_prompt = (
-                self.TASK_PROMPTS.get(task, self.config.user_prompt or "Transcribe: <audio>")
-                or "Transcribe: <audio>"
-            )
-
-        messages = []
-        if system_prompt:
-            messages.append({"role": "system", "content": system_prompt})
-        messages.append({"role": "user", "content": user_prompt})
-
-        prompt_ids = self.tokenizer.apply_chat_template(
-            messages,
-            tokenize=True,
-            add_generation_prompt=True,
-            return_tensors="pt",
-            enable_thinking=False,
-        ).to(device)
-
-        if len(prompt_ids.shape) == 1:
-            prompt_ids = prompt_ids.unsqueeze(0)
-
-        if not (prompt_ids == self.audio_token_id).any():
-            raise ValueError("Audio token <audio> not found in prompt")
-
-        num_audio_tokens = audio_embeds.shape[1]
-        expanded_prompt_ids = self._expand_audio_tokens(prompt_ids, num_audio_tokens)
-        inputs_embeds = self._prepare_audio_inputs_embeds(expanded_prompt_ids, audio_embeds)
-        input_token_count = expanded_prompt_ids.shape[1]
-
-        attention_mask = torch.ones(
-            batch_size, input_token_count, dtype=torch.long, device=device
-        )
-
-        # Set up generation parameters
-        if max_new_tokens is None:
-            max_new_tokens = getattr(self.config, "max_new_tokens", 256)
-
-        generate_kwargs.setdefault("max_new_tokens", max_new_tokens)
-        generate_kwargs.setdefault("use_cache", True)
-        generate_kwargs.setdefault(
-            "eos_token_id", self.tokenizer.convert_tokens_to_ids("<|im_end|>")
-        )
-        generate_kwargs.setdefault("pad_token_id", self.tokenizer.pad_token_id)
-
-        if temperature is not None:
-            generate_kwargs["temperature"] = temperature
-            generate_kwargs.setdefault("do_sample", True)
-
-        # Set up the streamer
-        streamer = TextIteratorStreamer(
-            self.tokenizer,
-            skip_prompt=True,
-            skip_special_tokens=True
-        )
-
-        # Generate in a separate thread
-        def generation_thread(future: futures.Future):
-            try:
-                result = self.decoder.generate(
-                    input_ids=expanded_prompt_ids,
-                    inputs_embeds=inputs_embeds,
-                    attention_mask=attention_mask,
-                    streamer=streamer,
-                    **generate_kwargs,
-                )
-                future.set_result(result)
-            except Exception as e:
-                future.set_exception(e)
-
-        future: futures.Future[torch.Tensor] = futures.Future()
-        thread = threading.Thread(target=generation_thread, args=(future,))
-        thread.start()
-
-        # Stream the output
-        output_text = ""
-        output_token_count = 0
-
-        try:
-            for chunk in streamer:
-                if chunk:
-                    output_text += chunk
-                    output_token_count += 1
-                    yield StreamChunk(chunk)
-        finally:
-            # Wait for generation to complete
-            thread.join()
-
-            # Check if there was an exception
-            if future.exception():
-                raise future.exception()
-
-        # Yield final statistics
-        yield StreamStats(input_token_count, output_token_count)
-
     def save_pretrained(self, save_directory: Union[str, Path], **kwargs):
         import shutil
         from pathlib import Path as PathlibPath

asr_pipeline.py

@@ -1,13 +1,13 @@
-from typing import Any, Dict, Generator, Union
+from typing import Any, Dict
 
 import torch
 import transformers
 from truecase import get_true_case
 
 try:
-    from .asr_modeling import ASRModel, StreamChunk, StreamStats
+    from .asr_modeling import ASRModel
 except ImportError:
-    from asr_modeling import ASRModel, StreamChunk, StreamStats  # type: ignore[no-redef]
+    from asr_modeling import ASRModel  # type: ignore[no-redef]
 
 
 class ASRPipeline(transformers.AutomaticSpeechRecognitionPipeline):
@@ -31,11 +31,6 @@ class ASRPipeline(transformers.AutomaticSpeechRecognitionPipeline):
             self.text_normalizer = WhisperTokenizer.from_pretrained("openai/whisper-tiny")
 
     def __call__(self, inputs, **kwargs):
-        # Check if streaming is requested
-        stream = kwargs.pop("stream", False)
-        if stream:
-            return self._stream_inference(inputs, **kwargs)
-
         generate_kwargs = {}
         for key in [
             "max_new_tokens",
@@ -297,111 +292,3 @@ class ASRPipeline(transformers.AutomaticSpeechRecognitionPipeline):
         text = get_true_case(text)
 
         return {"text": text}
-
-    def _stream_inference(
-        self, inputs, **kwargs
-    ) -> Generator[Union[Dict[str, str], Dict[str, int]], None, None]:
-        """
-        Perform streaming inference on audio input.
-
-        Args:
-            inputs: Audio input (same format as __call__)
-            **kwargs: Generation parameters
-
-        Yields:
-            Dict with "text" key containing text chunks as they're generated,
-            followed by a final dict with "input_tokens" and "output_tokens" statistics
-        """
-        # Extract generation kwargs
-        generate_kwargs = {}
-        for key in [
-            "max_new_tokens",
-            "temperature",
-            "do_sample",
-            "top_k",
-            "top_p",
-            "user_prompt",
-            "task",
-            "system_prompt",
-        ]:
-            if key in kwargs:
-                generate_kwargs[key] = kwargs.pop(key)
-
-        # Disable chunking for streaming - we want the whole audio at once
-        kwargs.pop("chunk_length_s", None)
-        kwargs.pop("stride_length_s", None)
-
-        # Preprocess audio to get model inputs
-        model_inputs = self.preprocess(inputs, chunk_length_s=0, **kwargs)
-
-        # Handle different input formats
-        audio_inputs = None
-        is_whisper = False
-
-        # Check if preprocess returned an iterator (shouldn't with chunk_length_s=0)
-        from collections.abc import Iterator
-        if isinstance(model_inputs, Iterator):
-            # Get the first (and should be only) chunk
-            try:
-                model_inputs = next(model_inputs)
-            except StopIteration:
-                raise ValueError("Preprocess returned empty iterator")
-
-        if isinstance(model_inputs, torch.Tensor):
-            audio_inputs = model_inputs
-        elif isinstance(model_inputs, dict):
-            # Remove metadata fields
-            model_inputs.pop("is_last", None)
-            model_inputs.pop("stride", None)
-
-            # Get audio input (Whisper uses input_features, others use input_values)
-            if "input_features" in model_inputs:
-                audio_inputs = model_inputs["input_features"]
-                is_whisper = True
-            else:
-                audio_inputs = model_inputs.get("input_values")
-
-        if audio_inputs is None:
-            # Debug info
-            import sys
-            print(f"DEBUG: model_inputs type: {type(model_inputs)}", file=sys.stderr)
-            if isinstance(model_inputs, dict):
-                print(f"DEBUG: model_inputs keys: {model_inputs.keys()}", file=sys.stderr)
-            raise ValueError(f"Could not extract audio inputs from preprocessing. Got type: {type(model_inputs)}")
-
-        if isinstance(audio_inputs, torch.Tensor):
-            audio_inputs = audio_inputs.to(self.model.device)
-        else:
-            raise ValueError(f"audio inputs must be a tensor, got {type(audio_inputs)}")
-
-        # Call the streaming generate method
-        if is_whisper:
-            stream_generator = self.model.generate_stream(
-                input_features=audio_inputs,
-                **generate_kwargs,
-            )
-        else:
-            stream_generator = self.model.generate_stream(
-                input_values=audio_inputs,
-                **generate_kwargs,
-            )
-
-        # Track full text for post-processing
-        full_text = ""
-
-        # Stream the chunks
-        for item in stream_generator:
-            if isinstance(item, StreamChunk):
-                full_text += item.text
-                yield {"text": item.text}
-            elif isinstance(item, StreamStats):
-                # Apply post-processing to the full text
-                processed_text = self.text_normalizer.normalize(full_text)
-                processed_text = get_true_case(processed_text)
-
-                # Yield final statistics with processed text
-                yield {
-                    "input_tokens": item.input_tokens,
-                    "output_tokens": item.output_tokens,
-                    "full_text": processed_text,
-                }