Training in progress - step 8000

asr_modeling.py

@@ -616,7 +616,6 @@ class ASRModel(PreTrainedModel):
         system_prompt: Optional[str] = None,
         user_prompt: Optional[str] = None,
         task: Optional[str] = None,
-        streamer: Optional[TextIteratorStreamer] = None,
         **generate_kwargs,
     ) -> Union[
         torch.Tensor,
@@ -698,27 +697,14 @@ class ASRModel(PreTrainedModel):
         generate_kwargs.setdefault("pad_token_id", self.tokenizer.pad_token_id)
         prompt_length = expanded_prompt_ids.shape[1]
 
-        # Generate with or without streamer
-        if streamer is not None:
-            generated_ids = self.decoder.generate(
-                input_ids=expanded_prompt_ids,
-                inputs_embeds=inputs_embeds,
-                attention_mask=attention_mask,
-                streamer=streamer,
-                **generate_kwargs,
-            )
-            # When using a streamer, return the full output (streamer will handle skipping prompt)
-            # The streamer needs the full sequence to properly identify what to skip
-            return generated_ids
-        else:
-            generated_ids = self.decoder.generate(
-                input_ids=expanded_prompt_ids,
-                inputs_embeds=inputs_embeds,
-                attention_mask=attention_mask,
-                **generate_kwargs,
-            )
-            # When not streaming, return only the new tokens (without prompt)
-            return generated_ids[:, prompt_length:]
+        generated_ids = self.decoder.generate(
+            input_ids=expanded_prompt_ids,
+            inputs_embeds=inputs_embeds,
+            attention_mask=attention_mask,
+            **generate_kwargs,
+        )
+
+        return generated_ids[:, prompt_length:]
 
     @torch.no_grad()
     def generate_stream(
@@ -728,39 +714,105 @@ class ASRModel(PreTrainedModel):
         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]:
         """
-        Stream generation by using the working generate() method with a TextIteratorStreamer.
+        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)
         """
-        # Set up the streamer - use skip_prompt=True like Ultravox
-        # The key is that when we return the full sequence from generate(),
-        # the streamer can properly identify and skip the prompt
-        streamer = TextIteratorStreamer(
-            self.tokenizer,
-            skip_prompt=True,  # Skip the prompt tokens
-            skip_special_tokens=True,
-            timeout=30.0
-        )
-
         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")
+            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
 
-        import threading
-        from concurrent import futures
+        if batch_size > 1:
+            raise ValueError("Streaming generation only supports batch_size=1")
 
-        # Run generation in a thread with streamer
+        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:
-                # Call generate with the streamer
-                # Important: This now returns the FULL sequence when streaming
-                result = self.generate(
-                    input_values=input_values,
-                    input_features=input_features,
-                    system_prompt=system_prompt,
-                    user_prompt=user_prompt,
-                    task=task,
+                result = self.decoder.generate(
+                    input_ids=expanded_prompt_ids,
+                    inputs_embeds=inputs_embeds,
+                    attention_mask=attention_mask,
                     streamer=streamer,
                     **generate_kwargs,
                 )
@@ -768,47 +820,30 @@ class ASRModel(PreTrainedModel):
             except Exception as e:
                 future.set_exception(e)
 
-        future: futures.Future = futures.Future()
+        future: futures.Future[torch.Tensor] = futures.Future()
         thread = threading.Thread(target=generation_thread, args=(future,))
         thread.start()
 
-        # Stream the output - like Ultravox, just yield chunks as they come
+        # Stream the output
+        output_text = ""
         output_token_count = 0
+
         try:
             for chunk in streamer:
-                if chunk:  # Only yield non-empty chunks
+                if chunk:
+                    output_text += chunk
                     output_token_count += 1
                     yield StreamChunk(chunk)
-        except Exception as e:
-            # Check if it's the Empty exception from queue
-            if e.__class__.__name__ == "Empty":
-                # This happens when generation completes before we start iterating
-                pass
-            else:
-                # Re-raise other exceptions
-                raise
         finally:
             # Wait for generation to complete
             thread.join()
+
+            # Check if there was an exception
             if future.exception():
                 raise future.exception()
 
-            # Debug: If no chunks were yielded, check what was generated
-            if output_token_count == 0:
-                import sys
-                result = future.result()
-                if result is not None:
-                    # Note: result now includes the full sequence (including prompt)
-                    # when streaming, so decode the full thing
-                    decoded = self.tokenizer.decode(result[0], skip_special_tokens=True)
-                    print(f"DEBUG: No chunks yielded but generated: {decoded}", file=sys.stderr)
-
-        # For stats, estimate input tokens (we can't easily get exact count without duplicating work)
-        # Rough estimate: prompt is about 20 tokens + 750 audio tokens
-        estimated_input_tokens = 770
-
         # Yield final statistics
-        yield StreamStats(estimated_input_tokens, output_token_count)
+        yield StreamStats(input_token_count, output_token_count)
 
     def save_pretrained(self, save_directory: Union[str, Path], **kwargs):
         import shutil