| 159 | logging.info(f"rank: {rank}, model is builded.") |
| 160 | |
| 161 | def forward( |
| 162 | self, |
| 163 | speech: torch.Tensor = None, |
| 164 | speech_lengths: torch.Tensor = None, |
| 165 | input_ids: torch.Tensor = None, |
| 166 | attention_mask: torch.Tensor = None, |
| 167 | labels_ids: torch.Tensor = None, |
| 168 | fbank_beg: torch.Tensor = None, |
| 169 | fbank_mask: torch.Tensor = None, |
| 170 | **kwargs, |
| 171 | ): |
| 172 | batch_size, token_num = input_ids.shape |
| 173 | stats = {} |
| 174 | input_ids[input_ids < 0] = 0 |
| 175 | inputs_embeds = self.llm.model.get_input_embeddings()(input_ids) |
| 176 | if speech is not None: |
| 177 | if len(speech_lengths.size()) > 1: |
| 178 | speech_lengths = speech_lengths[:, 0] |
| 179 | batch_size_speech, frames, _ = speech.shape |
| 180 | |
| 181 | # audio encoder |
| 182 | if self.audio_encoder_activation_checkpoint: |
| 183 | from torch.utils.checkpoint import checkpoint |
| 184 | |
| 185 | encoder_out, encoder_out_lens = checkpoint( |
| 186 | self.encode, speech, speech_lengths, use_reentrant=False |
| 187 | ) |
| 188 | else: |
| 189 | encoder_out, encoder_out_lens = self.encode(speech, speech_lengths) |
| 190 | |
| 191 | # audio_adaptor |
| 192 | encoder_out, encoder_out_lens = self.audio_adaptor(encoder_out, encoder_out_lens) |
| 193 | |
| 194 | batch_size, token_num, dims = inputs_embeds.shape |
| 195 | fake_token_len = kwargs.get("fake_token_len") |
| 196 | fake_token_len[fake_token_len < 0] = 0 |
| 197 | fbank_beg[fbank_beg < 0] = 0 |
| 198 | |
| 199 | speech_idx = 0 |
| 200 | for batch_idx in range(batch_size): |
| 201 | for turn_id in range(fbank_beg.shape[1]): |
| 202 | fbank_beg_idx = fbank_beg[batch_idx, turn_id].item() |
| 203 | if fbank_beg_idx > 0: |
| 204 | speech_token_len = fake_token_len[batch_idx, turn_id] |
| 205 | speech_token = encoder_out[speech_idx, :speech_token_len, :] |
| 206 | |
| 207 | try: |
| 208 | inputs_embeds[ |
| 209 | batch_idx, |
| 210 | fbank_beg_idx : fbank_beg_idx + speech_token_len, |
| 211 | :, |
| 212 | ] = speech_token |
| 213 | except Exception as e: |
| 214 | logging.error(f"{str(e)}, {traceback.format_exc()}") |
| 215 | logging.info( |
| 216 | f"batch_idx: {batch_idx}, inputs_embeds: {inputs_embeds.shape}, fbank_beg_idx: {fbank_beg_idx}, speech_token_len: {speech_token_len}, encoder_out: {encoder_out.shape}, encoder_out_lens: {encoder_out_lens}, fake_token_len: {fake_token_len}, speech_lengths: {speech_lengths}" |
| 217 | ) |
| 218 | speech_token_len = encoder_out_lens[speech_idx].item() |