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,
@@ -707,6 +708,10 @@ class ASRModel(PreTrainedModel):
         print(f"DEBUG generate (non-streaming): task={task}, system_prompt={system_prompt}, user_prompt={user_prompt}", file=sys.stderr)
         print(f"DEBUG generate (non-streaming): generate_kwargs={generate_kwargs}", file=sys.stderr)
 
+        # Add streamer if provided
+        if streamer is not None:
+            generate_kwargs["streamer"] = streamer
+
         generated_ids = self.decoder.generate(
             input_ids=expanded_prompt_ids,
             inputs_embeds=inputs_embeds,
@@ -724,157 +729,11 @@ 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")
-
-        # Debug: Check audio inputs
-        import sys
-        print(f"DEBUG generate_stream: audio_inputs shape={audio_inputs.shape if audio_inputs is not None else None}", file=sys.stderr)
-        print(f"DEBUG generate_stream: audio_inputs type={type(audio_inputs)}", file=sys.stderr)
-
-        # 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 from config (same as non-streaming generate)
-        config_params = [
-            "max_new_tokens",
-            "min_new_tokens",
-            "num_beams",
-            "do_sample",
-            "temperature",
-            "top_k",
-            "top_p",
-            "repetition_penalty",
-            "length_penalty",
-            "no_repeat_ngram_size",
-            "early_stopping",
-        ]
-        for param in config_params:
-            if hasattr(self.config, param) and getattr(self.config, param) is not None:
-                generate_kwargs.setdefault(param, getattr(self.config, param))
-
-        # Override with explicit parameters if provided
-        if max_new_tokens is not None:
-            generate_kwargs["max_new_tokens"] = max_new_tokens
-
-        if temperature is not None:
-            generate_kwargs["temperature"] = temperature
-            generate_kwargs["do_sample"] = True
-
-        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)
-
-        # Debug: Check if audio embeds are in inputs_embeds
-        import sys
-        print(f"DEBUG generate_stream: task={task}, system_prompt={system_prompt}, user_prompt={user_prompt}", file=sys.stderr)
-        print(f"DEBUG generate_stream: inputs_embeds shape={inputs_embeds.shape}", file=sys.stderr)
-        print(f"DEBUG generate_stream: expanded_prompt_ids shape={expanded_prompt_ids.shape}", file=sys.stderr)
-        print(f"DEBUG generate_stream: audio_embeds shape={audio_embeds.shape}", file=sys.stderr)
-        print(f"DEBUG generate_stream: num_audio_tokens={num_audio_tokens}", file=sys.stderr)
-        print(f"DEBUG generate_stream: generate_kwargs={generate_kwargs}", file=sys.stderr)
-
-        # Debug: Check devices and values
-        print(f"DEBUG: inputs_embeds device={inputs_embeds.device}", file=sys.stderr)
-        print(f"DEBUG: expanded_prompt_ids device={expanded_prompt_ids.device}", file=sys.stderr)
-        print(f"DEBUG: attention_mask device={attention_mask.device}", file=sys.stderr)
-        print(f"DEBUG: decoder device={next(self.decoder.parameters()).device}", file=sys.stderr)
-
-        # Check if audio embeddings are non-zero
-        audio_mask = (expanded_prompt_ids == self.audio_token_id)
-        print(f"DEBUG: audio_mask sum={audio_mask.sum().item()} (should be {num_audio_tokens})", file=sys.stderr)
-
-        # Check a sample of the embeddings where audio should be
-        audio_positions = torch.where(audio_mask[0])[0]
-        if len(audio_positions) > 0:
-            sample_pos = audio_positions[0].item()
-            print(f"DEBUG: Sample audio embed at pos {sample_pos}: mean={inputs_embeds[0, sample_pos].mean().item():.4f}, std={inputs_embeds[0, sample_pos].std().item():.4f}", file=sys.stderr)
-
-        # Test: Try without threading first to see if that's the issue
-        print(f"DEBUG: Testing non-threaded generation first", file=sys.stderr)
-        print(f"DEBUG: input_token_count (prompt length) = {input_token_count}", file=sys.stderr)
-
-        test_output = self.decoder.generate(
-            input_ids=expanded_prompt_ids,
-            inputs_embeds=inputs_embeds,
-            attention_mask=attention_mask,
-            max_new_tokens=10,  # Just generate a few tokens to test
-            **{k: v for k, v in generate_kwargs.items() if k != 'max_new_tokens'}
-        )
-
-        # Debug the output
-        full_text = self.tokenizer.decode(test_output[0], skip_special_tokens=True)
-        print(f"DEBUG: Full output text: {full_text}", file=sys.stderr)
-
-        test_text = self.tokenizer.decode(test_output[0, input_token_count:], skip_special_tokens=True)
-        print(f"DEBUG: Non-threaded test output (after removing prompt): {test_text}", file=sys.stderr)
-
         # Set up the streamer
         streamer = TextIteratorStreamer(
             self.tokenizer,
@@ -882,13 +741,26 @@ class ASRModel(PreTrainedModel):
             skip_special_tokens=True
         )
 
-        # Generate in a separate thread
+        # Count prompt length for stats
+        # We need to encode just to get the prompt length
+        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")
+
+        # Simple way to get prompt length - just count audio tokens
+        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,
+                # Just call the working generate method with the streamer
+                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,
                 )
@@ -896,30 +768,29 @@ class ASRModel(PreTrainedModel):
             except Exception as e:
                 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()
 
         # 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()
 
+        # 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