Training in progress - step 13000

asr_config.py

@@ -46,7 +46,6 @@ class ASRConfig(transformers.PretrainedConfig):
             "min_new_tokens": 1,
             "do_sample": False,
             "repetition_penalty": 1.05,
-            "length_penalty": 1.0,
             "no_repeat_ngram_size": 0,
             "use_cache": True,
         }

asr_modeling.py

@@ -1,5 +1,8 @@
 from pathlib import Path
-from typing import Optional, Union
+from typing import Optional, Union, Generator, NamedTuple
+
+import threading
+from concurrent import futures
 
 import torch
 import torch.nn as nn
@@ -11,6 +14,7 @@ from transformers import (
     AutoTokenizer,
     PreTrainedModel,
     Wav2Vec2FeatureExtractor,
+    TextIteratorStreamer,
 )
 from transformers.generation.utils import (
     GenerateBeamDecoderOnlyOutput,
@@ -25,6 +29,17 @@ 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__()
@@ -118,8 +133,12 @@ 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)
+        return Wav2Vec2FeatureExtractor.from_pretrained(
+            audio_model_id,
+            do_normalize=True,
+        )
 
     @classmethod
     def from_pretrained(cls, pretrained_model_name_or_path, *args, **kwargs):
@@ -206,10 +225,6 @@ class ASRModel(PreTrainedModel):
         self.decoder = self._create_decoder(config)
         self.generation_config = self.decoder.generation_config
 
-        # Set default generation parameters
-        self.generation_config.num_beams = 1
-        self.generation_config.length_penalty = 1.0
-
         self._init_tokenizer()
 
         from types import SimpleNamespace
@@ -691,6 +706,145 @@ 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
+from typing import Any, Dict, Generator, Union
 
 import torch
 import transformers
 from truecase import get_true_case
 
 try:
-    from .asr_modeling import ASRModel
+    from .asr_modeling import ASRModel, StreamChunk, StreamStats
 except ImportError:
-    from asr_modeling import ASRModel  # type: ignore[no-redef]
+    from asr_modeling import ASRModel, StreamChunk, StreamStats  # type: ignore[no-redef]
 
 
 class ASRPipeline(transformers.AutomaticSpeechRecognitionPipeline):
@@ -31,6 +31,11 @@ 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",
@@ -292,3 +297,111 @@ 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,
+                }