| 476 | |
| 477 | |
| 478 | class IndividualTokenRefiner(torch.nn.Module): |
| 479 | def __init__( |
| 480 | self, |
| 481 | hidden_size, |
| 482 | heads_num, |
| 483 | depth, |
| 484 | mlp_width_ratio: float = 4.0, |
| 485 | mlp_drop_rate: float = 0.0, |
| 486 | act_type: str = "silu", |
| 487 | qk_norm: bool = False, |
| 488 | qk_norm_type: str = "layer", |
| 489 | qkv_bias: bool = True, |
| 490 | need_CA:bool=False, |
| 491 | dtype: Optional[torch.dtype] = None, |
| 492 | device: Optional[torch.device] = None, |
| 493 | ): |
| 494 | |
| 495 | factory_kwargs = {"device": device, "dtype": dtype} |
| 496 | super().__init__() |
| 497 | self.need_CA = need_CA |
| 498 | self.blocks = nn.ModuleList( |
| 499 | [ |
| 500 | IndividualTokenRefinerBlock( |
| 501 | hidden_size=hidden_size, |
| 502 | heads_num=heads_num, |
| 503 | mlp_width_ratio=mlp_width_ratio, |
| 504 | mlp_drop_rate=mlp_drop_rate, |
| 505 | act_type=act_type, |
| 506 | qk_norm=qk_norm, |
| 507 | qk_norm_type=qk_norm_type, |
| 508 | qkv_bias=qkv_bias, |
| 509 | need_CA=self.need_CA, |
| 510 | **factory_kwargs, |
| 511 | ) |
| 512 | for _ in range(depth) |
| 513 | ] |
| 514 | ) |
| 515 | |
| 516 | |
| 517 | def forward( |
| 518 | self, |
| 519 | x: torch.Tensor, |
| 520 | c: torch.LongTensor, |
| 521 | mask: Optional[torch.Tensor] = None, |
| 522 | y:torch.Tensor=None, |
| 523 | ): |
| 524 | self_attn_mask = None |
| 525 | if mask is not None: |
| 526 | batch_size = mask.shape[0] |
| 527 | seq_len = mask.shape[1] |
| 528 | mask = mask.to(x.device) |
| 529 | # batch_size x 1 x seq_len x seq_len |
| 530 | self_attn_mask_1 = mask.view(batch_size, 1, 1, seq_len).repeat( |
| 531 | 1, 1, seq_len, 1 |
| 532 | ) |
| 533 | # batch_size x 1 x seq_len x seq_len |
| 534 | self_attn_mask_2 = self_attn_mask_1.transpose(2, 3) |
| 535 | # batch_size x 1 x seq_len x seq_len, 1 for broadcasting of heads_num |