Args: audio16k: shape (b, t) audio16k_length: (b,) Returns: token: shape (b, nq, l) token_length: (b,)
(
self,
audio16k: torch.Tensor,
audio16k_length: torch.Tensor = None,
batch_size: int = 96,
)
| 243 | |
| 244 | @torch.inference_mode() |
| 245 | def encode( |
| 246 | self, |
| 247 | audio16k: torch.Tensor, |
| 248 | audio16k_length: torch.Tensor = None, |
| 249 | batch_size: int = 96, |
| 250 | ): |
| 251 | """ |
| 252 | Args: |
| 253 | audio16k: shape (b, t) |
| 254 | audio16k_length: (b,) |
| 255 | Returns: |
| 256 | token: shape (b, nq, l) |
| 257 | token_length: (b,) |
| 258 | """ |
| 259 | if audio16k_length is None: |
| 260 | assert audio16k.shape[0] == 1 |
| 261 | audio16k_length = torch.tensor( |
| 262 | [audio16k.shape[1]], dtype=torch.long, device=audio16k.device |
| 263 | ) |
| 264 | |
| 265 | CHUNK_SIZE = 6 * 16000 |
| 266 | B, T = audio16k.shape |
| 267 | # Pad, chunk, and batch |
| 268 | audio16k_batch = [] |
| 269 | batch_size_list = [] |
| 270 | for i in range(B): |
| 271 | # Remove extra paddings |
| 272 | one_audio_chunks = self._pad_and_chunk( |
| 273 | audio16k[i : (i + 1), : audio16k_length[i]], CHUNK_SIZE |
| 274 | ) |
| 275 | audio16k_batch += one_audio_chunks |
| 276 | batch_size_list.append(len(one_audio_chunks)) |
| 277 | audio16k_batch = torch.cat(audio16k_batch, dim=0) |
| 278 | # Batch encode |
| 279 | token_batch = [] |
| 280 | for i in range(0, audio16k_batch.shape[0], batch_size): |
| 281 | one_audio_batch = audio16k_batch[i : (i + batch_size)] |
| 282 | one_token_batch = self._encode_one_batch(one_audio_batch) |
| 283 | token_batch.append(one_token_batch) |
| 284 | token_batch = torch.cat(token_batch, dim=0) |
| 285 | # Recover & concat |
| 286 | token_list = torch.split( |
| 287 | token_batch, batch_size_list, dim=0 |
| 288 | ) # [(B=1, nq, l), (B=3, nq, l), ...] |
| 289 | token_list = [ |
| 290 | torch.cat(token_ts.split(1, dim=0), dim=-1) # (B=1, nq, l) |
| 291 | for token_ts in token_list |
| 292 | ] |
| 293 | # Pad tokens |
| 294 | token = pad_sequence( |
| 295 | [ts.squeeze(0).transpose(1, 0) for ts in token_list], |
| 296 | batch_first=True, |
| 297 | padding_value=0, |
| 298 | ).transpose( |
| 299 | 1, 2 |
| 300 | ) # (B, nq, L) |
| 301 | token_length = (audio16k_length / 1280).ceil().long() |
| 302 | token = token[ |
no test coverage detected