Training in progress - step 30500

asr_config.py

@@ -18,6 +18,7 @@ class ASRConfig(transformers.PretrainedConfig):
         user_prompt: str = "Please transcribe this English audio into text: <audio>",
         encoder_dim: Optional[int] = None,
         llm_dim: Optional[int] = None,
+        encoder_stride: int = 2,  # Temporal downsampling factor of audio encoder (legacy, use encoder_conv_layers)
         # Encoder conv layers: list of (padding, kernel_size, stride) tuples
         # Default is Whisper/GLM-ASR structure: conv1(k=3,s=1,p=1) + conv2(k=3,s=2,p=1)
         encoder_conv_layers: Optional[list] = None,
@@ -51,7 +52,7 @@ class ASRConfig(transformers.PretrainedConfig):
         # Set default generation parameters (greedy decoding only)
         generation_defaults = {
             "num_beams": 1,
-            "max_new_tokens": 256,
+            "max_new_tokens": 96,
             "repetition_penalty": 1.0,
             "length_penalty": 1.0,
             "no_repeat_ngram_size": 0,
@@ -69,6 +70,7 @@ class ASRConfig(transformers.PretrainedConfig):
         self.user_prompt = user_prompt
         self.encoder_dim = encoder_dim
         self.llm_dim = llm_dim
+        self.encoder_stride = encoder_stride
         # Default conv layers for Whisper/GLM-ASR: [(pad, kernel, stride), ...]
         self.encoder_conv_layers = encoder_conv_layers or [(1, 3, 1), (1, 3, 2)]
         self.audio_sample_rate = audio_sample_rate

asr_modeling.py

@@ -96,6 +96,7 @@ class ASRModel(PreTrainedModel, GenerationMixin):
         super().__init__(config)
 
         self.system_prompt = config.system_prompt
+        self.encoder_stride = config.encoder_stride
         target_dtype = getattr(torch, config.model_dtype)
 
         # Audio encoder (frozen)
@@ -120,10 +121,7 @@ class ASRModel(PreTrainedModel, GenerationMixin):
         self.generation_config.length_penalty = config.length_penalty
         self.generation_config.repetition_penalty = config.repetition_penalty
         self.generation_config.no_repeat_ngram_size = config.no_repeat_ngram_size
-        self.generation_config.eos_token_id = [
-            self.tokenizer.convert_tokens_to_ids("<|im_end|>"),
-            self.tokenizer.convert_tokens_to_ids("<|endoftext|>"),
-        ]
+        self.generation_config.eos_token_id = self.tokenizer.convert_tokens_to_ids("<|im_end|>")
         self.generation_config.pad_token_id = self.tokenizer.pad_token_id
 
         # Feature extractor for audio preprocessing
@@ -147,7 +145,7 @@ class ASRModel(PreTrainedModel, GenerationMixin):
         encoder_kwargs = {
             "attn_implementation": config.attn_implementation,
             "low_cpu_mem_usage": True,
-            "dtype": dtype,
+            "torch_dtype": dtype,
         }
 
         if "whisper" in config.audio_model_id.lower():
@@ -298,6 +296,7 @@ class ASRModel(PreTrainedModel, GenerationMixin):
             feature_extractor=self.feature_extractor,
             tokenizer=self.tokenizer,
             projector=self.projector,
+            encoder_stride=self.encoder_stride,
             encoder_conv_layers=self.config.encoder_conv_layers,
         )
 

asr_pipeline.py

@@ -476,32 +476,4 @@ class ASRPipeline(transformers.AutomaticSpeechRecognitionPipeline):
         text = self.tokenizer.decode(tokens, skip_special_tokens=True).strip()
         # Strip <think>...</think> tags (Qwen3 doesn't respect /no_think prompt)
         text = re.sub(r"<think>.*?</think>\s*", "", text, flags=re.DOTALL).strip()
-        # Truncate if a word repeats more than 3 times consecutively
-        text = self._truncate_repetitions(text, max_repeats=3)
         return {"text": text}
-
-    def _truncate_repetitions(self, text: str, max_repeats: int = 3) -> str:
-        """Truncate text when a word repeats more than max_repeats times consecutively.
-
-        Args:
-            text: Input text to check for repetitions
-            max_repeats: Maximum allowed consecutive repetitions (default 3)
-
-        Returns:
-            Truncated text if repetition detected, otherwise original text
-        """
-        words = text.split()
-        if len(words) <= max_repeats:
-            return text
-
-        repeat_count = 1
-        for i in range(1, len(words)):
-            if words[i].lower() == words[i - 1].lower():
-                repeat_count += 1
-                if repeat_count > max_repeats:
-                    # Keep up to max_repeats of the repeated word
-                    return " ".join(words[:i])
-            else:
-                repeat_count = 1
-
-        return text

asr_processing.py

@@ -26,12 +26,14 @@ class ASRProcessor(ProcessorMixin):
         feature_extractor,
         tokenizer,
         projector=None,
+        encoder_stride: int = 2,
         encoder_conv_layers: Optional[list] = None,
     ):
         self.feature_extractor = feature_extractor
         self.tokenizer = tokenizer
         self.audio_token_id = tokenizer.convert_tokens_to_ids(self.AUDIO_TOKEN)
         self.projector = projector
+        self.encoder_stride = encoder_stride  # Legacy, kept for compatibility
         self.encoder_conv_layers = encoder_conv_layers or self.DEFAULT_ENCODER_CONV_LAYERS
 
     def _compute_encoder_output_length(self, mel_length: int) -> int:

model.safetensors

@@ -1,3 +1,3 @@
 version https://git-lfs.github.com/spec/v1
-oid sha256:4ba6c7ffd625764146a13f4678f459fe084bec15e140db28016239aac516f158
+oid sha256:c5ce422bc8492d610e980da968fec9e97cc628c00a928b8d9cdc24197ab5910c
 size 58732960