Update custom model files, README, and requirements

asr_modeling.py

@@ -616,6 +616,7 @@ 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,
@@ -697,14 +698,27 @@ class ASRModel(PreTrainedModel):
         generate_kwargs.setdefault("pad_token_id", self.tokenizer.pad_token_id)
         prompt_length = expanded_prompt_ids.shape[1]
 
-        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:]
+        # 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:]
 
     @torch.no_grad()
     def generate_stream(
@@ -714,136 +728,95 @@ 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]:
         """
-        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)
+        Stream generation by using the working generate() method with a TextIteratorStreamer.
         """
-        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
+        # 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_special_tokens=True
+            skip_prompt=True,  # Skip the prompt tokens
+            skip_special_tokens=True,
+            timeout=30.0
         )
 
-        # Generate in a separate thread
+        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")
+
+        import threading
+        from concurrent import futures
+
+        # Run generation in a thread with streamer
         def generation_thread(future: futures.Future):
             try:
-                result = self.decoder.generate(
-                    input_ids=expanded_prompt_ids,
-                    inputs_embeds=inputs_embeds,
-                    attention_mask=attention_mask,
+                import sys
+                print("DEBUG: Starting generation thread", file=sys.stderr)
+                # 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,
                     streamer=streamer,
                     **generate_kwargs,
                 )
+                print("DEBUG: Generation complete", file=sys.stderr)
                 future.set_result(result)
             except Exception as e:
+                print(f"DEBUG: Generation error: {e}", file=sys.stderr)
                 future.set_exception(e)
 
-        future: futures.Future[torch.Tensor] = futures.Future()
+        future: futures.Future = futures.Future()
         thread = threading.Thread(target=generation_thread, args=(future,))
         thread.start()
+        print("DEBUG: Thread started", file=sys.stderr)
 
-        # Stream the output
-        output_text = ""
+        # Stream the output - like Ultravox, just yield chunks as they come
         output_token_count = 0
-
+        import sys
+        print("DEBUG: Starting streaming iteration", file=sys.stderr)
         try:
             for chunk in streamer:
-                if chunk:
-                    output_text += chunk
+                print(f"DEBUG: Got chunk: {repr(chunk)}", file=sys.stderr)
+                if chunk:  # Only yield non-empty chunks
                     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(input_token_count, output_token_count)
+        yield StreamStats(estimated_input_tokens, output_token_count)
 
     def save_pretrained(self, save_directory: Union[str, Path], **kwargs):
         import shutil