Training in progress - step 1000

asr_config.py

@@ -66,6 +66,7 @@ class ASRConfig(transformers.PretrainedConfig):
         lora_dropout: float = 0.0,
         lora_target_modules: Optional[list] = None,  # Default: all linear layers
         freeze_projector: bool = False,  # True for Stage 2 (LoRA-only training)
+        freeze_language_model: bool = True,  # False = full decoder fine-tuning
         do_sample: bool = False,
         temperature: Optional[float] = None,
         top_p: Optional[float] = None,
@@ -152,6 +153,7 @@ class ASRConfig(transformers.PretrainedConfig):
             "down_proj",
         ]
         self.freeze_projector = freeze_projector
+        self.freeze_language_model = freeze_language_model
 
         explicit_generation_args = {
             "num_beams": num_beams,

asr_modeling.py

@@ -58,7 +58,7 @@ class ASRModel(PreTrainedModel, GenerationMixin):
     _is_loading_from_pretrained: bool = False
     _pretrained_model_path: Optional[str] = None
 
-    TRANSCRIBE_PROMPT = ""
+    TRANSCRIBE_PROMPT = "Transcribe the speech to text"
 
     @classmethod
     def from_pretrained(cls, pretrained_model_name_or_path: str, *args, **kwargs) -> "ASRModel":
@@ -265,8 +265,9 @@ class ASRModel(PreTrainedModel, GenerationMixin):
 
         decoder = AutoModelForCausalLM.from_pretrained(config.text_model_id, **decoder_kwargs)
         decoder.config.use_cache = getattr(config, "use_cache", True)
-        decoder.requires_grad_(False)
-        decoder.eval()
+        if getattr(config, "freeze_language_model", True):
+            decoder.requires_grad_(False)
+            decoder.train(False)
         return decoder
 
     def _create_projector(self, config: ASRConfig, dtype: torch.dtype) -> nn.Module:
@@ -395,8 +396,13 @@ class ASRModel(PreTrainedModel, GenerationMixin):
         )
 
     def state_dict(self, *args, **kwargs) -> dict[str, torch.Tensor]:
-        """Only save trainable projector weights."""
-        return {f"projector.{k}": v for k, v in self.projector.state_dict().items()}
+        """Save trainable weights: projector, plus the language model when fine-tuned."""
+        sd = {f"projector.{k}": v for k, v in self.projector.state_dict().items()}
+        if not getattr(self.config, "freeze_language_model", True):
+            sd.update(
+                {f"language_model.{k}": v for k, v in self.language_model.state_dict().items()}
+            )
+        return sd
 
     def _compute_encoder_output_lengths(
         self,

asr_processing.py

@@ -17,7 +17,7 @@ class ASRProcessor(ProcessorMixin):
     feature_extractor_class = "AutoFeatureExtractor"
     tokenizer_class = "AutoTokenizer"
     AUDIO_TOKEN = "<audio>"
-    TRANSCRIBE_PROMPT = ""
+    TRANSCRIBE_PROMPT = "Transcribe the speech to text"
 
     def __init__(
         self,

config.json

@@ -235,6 +235,7 @@
   ],
   "encoder_dim": 1280,
   "eos_token_id": 151645,
+  "freeze_language_model": false,
   "freeze_projector": false,
   "freq_mask_length": 27,
   "inference_warmup_tokens": 10,

model.safetensors

@@ -1,3 +1,3 @@
 version https://git-lfs.github.com/spec/v1
-oid sha256:dfbd8328dc22b130f2558cd3cf634a711688a57837e4f8d18bce72a38398dd4c
-size 25170248
+oid sha256:4121706cbf6dcdefc550c39232b722d9f8b6b180ebd04e75056da125de9ee705
+size 1216765200