Training in progress - step 1000

asr_config.py

@@ -49,6 +49,13 @@ class ASRConfig(transformers.PretrainedConfig):
         mask_feature_prob: float = 0.0,  # Probability of masking frequency bins (disabled by default)
         mask_feature_length: int = 10,  # Max length of frequency mask
         mask_feature_min_masks: int = 0,  # Min number of frequency masks
+        # LoRA configuration (for Stage 2 fine-tuning)
+        use_lora: bool = False,
+        lora_rank: int = 8,  # SALMONN default
+        lora_alpha: int = 32,  # SALMONN default (scaling factor 4.0)
+        lora_dropout: float = 0.0,
+        lora_target_modules: Optional[list] = None,  # Default: ["q_proj", "v_proj"]
+        freeze_projector: bool = False,  # True for Stage 2 (LoRA-only training)
         max_new_tokens: Optional[int] = None,
         min_new_tokens: Optional[int] = None,
         repetition_penalty: Optional[float] = None,
@@ -109,6 +116,13 @@ class ASRConfig(transformers.PretrainedConfig):
         self.mask_feature_prob = mask_feature_prob
         self.mask_feature_length = mask_feature_length
         self.mask_feature_min_masks = mask_feature_min_masks
+        # LoRA configuration
+        self.use_lora = use_lora
+        self.lora_rank = lora_rank
+        self.lora_alpha = lora_alpha
+        self.lora_dropout = lora_dropout
+        self.lora_target_modules = lora_target_modules or ["q_proj", "v_proj"]
+        self.freeze_projector = freeze_projector
 
         # Generation parameters (use explicit value if provided, else use default)
         self.num_beams = num_beams if num_beams is not None else generation_defaults["num_beams"]

asr_modeling.py

@@ -190,6 +190,28 @@ class ASRModel(PreTrainedModel, GenerationMixin):
                 state_dict = load_file(model_file)
                 model.load_state_dict(state_dict, strict=False)
 
+            # Load LoRA adapters if they exist and use_lora is enabled
+            if getattr(config, "use_lora", False):
+                adapter_file = cached_file(
+                    pretrained_model_name_or_path,
+                    "adapter_model.safetensors",
+                    _raise_exceptions_for_missing_entries=False,
+                    **cache_kwargs,
+                )
+                if adapter_file is not None:
+                    from peft import PeftModel
+
+                    # Get the directory containing the adapter
+                    import os
+
+                    adapter_dir = os.path.dirname(adapter_file)
+                    # Load adapter weights into the already-wrapped PeftModel
+                    model.language_model = PeftModel.from_pretrained(
+                        model.language_model.base_model,
+                        adapter_dir,
+                        is_trainable=True,
+                    )
+
             return model
         finally:
             cls._is_loading_from_pretrained = False
@@ -233,9 +255,17 @@ class ASRModel(PreTrainedModel, GenerationMixin):
         # Feature extractor for audio preprocessing
         self.feature_extractor = self._create_feature_extractor(config)
 
-        # Audio projector (trainable)
+        # Audio projector (trainable unless freeze_projector is set)
         self.projector = self._create_projector(config, target_dtype)
 
+        # Setup LoRA if enabled (Stage 2 fine-tuning)
+        if getattr(config, "use_lora", False):
+            self._setup_lora(config)
+
+        # Freeze projector if specified (for Stage 2 LoRA-only training)
+        if getattr(config, "freeze_projector", False):
+            self.projector.requires_grad_(False)
+
         # For model parallelism
         self._no_split_modules = getattr(self.language_model, "_no_split_modules", [])
 
@@ -333,6 +363,21 @@ class ASRModel(PreTrainedModel, GenerationMixin):
         device = next(self.language_model.parameters()).device
         return projector.to(device=device, dtype=dtype)
 
+    def _setup_lora(self, config: ASRConfig):
+        """Apply LoRA adapters to the language model for Stage 2 fine-tuning."""
+        from peft import LoraConfig, get_peft_model
+
+        lora_config = LoraConfig(
+            r=config.lora_rank,
+            lora_alpha=config.lora_alpha,
+            target_modules=config.lora_target_modules,
+            lora_dropout=config.lora_dropout,
+            bias="none",
+            task_type="CAUSAL_LM",
+        )
+        self.language_model = get_peft_model(self.language_model, lora_config)
+        # LoRA params are trainable by default, base model stays frozen
+
     def _init_tokenizer(self, config: ASRConfig):
         """Initialize tokenizer with audio token."""
         self.tokenizer = AutoTokenizer.from_pretrained(config.text_model_id, trust_remote_code=True)
@@ -600,6 +645,7 @@ class ASRModel(PreTrainedModel, GenerationMixin):
                 tokenize=True,
                 add_generation_prompt=True,
                 return_tensors="pt",
+                enable_thinking=False,  # Disable Qwen3 thinking mode for ASR
             )
             input_ids = chat_result.input_ids.to(device)
 
@@ -674,6 +720,7 @@ class ASRModel(PreTrainedModel, GenerationMixin):
             tokenize=True,
             add_generation_prompt=True,
             return_tensors="pt",
+            enable_thinking=False,  # Disable Qwen3 thinking mode for ASR
         )
         input_ids = chat_result.input_ids.to(device)
 
@@ -773,6 +820,10 @@ class ASRModel(PreTrainedModel, GenerationMixin):
         self.tokenizer.save_pretrained(save_dir)
         self.feature_extractor.save_pretrained(save_dir)
 
+        # Save LoRA adapters if present (creates adapter_model.safetensors and adapter_config.json)
+        if hasattr(self.language_model, "peft_config"):
+            self.language_model.save_pretrained(save_dir)
+
         # Add processor auto_map to preprocessor_config.json
         config_path = save_dir / "preprocessor_config.json"
         if config_path.exists():

asr_processing.py

@@ -99,6 +99,7 @@ class ASRProcessor(ProcessorMixin):
             tokenize=True,
             add_generation_prompt=(text is None),
             return_tensors=return_tensors,
+            enable_thinking=False,  # Disable Qwen3 thinking mode for ASR
         )
 
         # Handle both tensor and BatchEncoding returns

chat_template.jinja

@@ -83,7 +83,7 @@
 {%- endfor %}
 {%- if add_generation_prompt %}
     {{- '<|im_start|>assistant\n' }}
-    {%- if enable_thinking is defined and enable_thinking is false %}
+    {%- if true %}
         {{- '<think>\n\n</think>\n\n' }}
     {%- endif %}
 {%- endif %}