Update modeling_fastslm.py

modeling_fastslm.py

@@ -289,7 +289,6 @@ class SpeechEncoder(nn.Module):
 class FastSLMPreTrainedModel(PreTrainedModel):
     config_class = FastSLMConfig
     base_model_prefix = "fastslm"
-
     def _init_weights(self, module):
         if isinstance(module, nn.Linear):
             nn.init.normal_(module.weight, std=0.02)
@@ -307,8 +306,8 @@ class FastSLMForConditionalGeneration(FastSLMPreTrainedModel, GenerationMixin):
             config.llm_config,
             trust_remote_code=True
         )
-        if self.llm._tied_weights_keys is not None:
-            self._tied_weights_keys = [f"llm.{k}" for k in self.llm._tied_weights_keys]
+        # if self.llm._tied_weights_keys is not None:
+        #     self._tied_weights_keys = [f"llm.{k}" for k in self.llm._tied_weights_keys]
 
         llm_lora_config = LoraConfig(
         r=config.lora_r,           
@@ -327,6 +326,8 @@ class FastSLMForConditionalGeneration(FastSLMPreTrainedModel, GenerationMixin):
         special_tokens = audio_token + language_token + task_token
         self.tokenizer.add_special_tokens({"additional_special_tokens": special_tokens})
 
+        self.post_init()
+        
     def get_input_embeddings(self) -> nn.Module:
         """Returns the input embedding layer of the LLM."""
         return self.llm.get_input_embeddings()
@@ -358,7 +359,7 @@ class FastSLMForConditionalGeneration(FastSLMPreTrainedModel, GenerationMixin):
     ):
         speech_query, speech_attn_mask = self.encoder(audio)
 
-        token_embedding = self.llm.get_input_embeddings()
+        token_embedding = self.get_input_embeddings()
         
         # Create speech labels (-100 to ignore in loss calculation)
         speech_label_len = int(speech_query.shape[1])
@@ -397,7 +398,7 @@ class FastSLMForConditionalGeneration(FastSLMPreTrainedModel, GenerationMixin):
         return outputs
 
     def generate(self, input_ids, audio: List[torch.Tensor] = None, **kwargs):
-        token_embedding = self.llm.get_input_embeddings()
+        token_embedding = self.get_input_embeddings()
         if audio is not None:
             speech_query, speech_attn_mask = self.encoder(audio)
             audio_token_id = self.tokenizer.convert_tokens_to_ids("<|AUDIO|>")