(self, text, context=None)
| 494 | return mask |
| 495 | |
| 496 | def forward(self, text, context=None): |
| 497 | if context is not None: |
| 498 | x_text = self.token_embedding(text) # n_clas, n_text, C |
| 499 | K, N1, C = x_text.shape |
| 500 | if len(context.shape) == 3: |
| 501 | B, N2, C = context.shape |
| 502 | |
| 503 | eos_indx = text.argmax(dim=-1) + N2 |
| 504 | eos_indx = eos_indx.reshape(1, K).expand(B, K).reshape(-1) |
| 505 | |
| 506 | x_text = x_text.reshape(1, K, N1, C).expand(B, K, N1, C) |
| 507 | context = context.reshape(B, 1, N2, C).expand(B, K, N2, C) |
| 508 | |
| 509 | elif len(context.shape) == 4: |
| 510 | B, K, N2, C = context.shape |
| 511 | |
| 512 | eos_indx = text.argmax(dim=-1) + N2 |
| 513 | eos_indx = eos_indx.reshape(1, K).expand(B, K).reshape(-1) |
| 514 | |
| 515 | x_text = x_text.reshape(1, K, N1, C).expand(B, K, N1, C) |
| 516 | x = torch.cat([x_text[:,:,0:1], context, x_text[:, :, 1:]], dim=2).reshape(B*K, N1+N2, C) |
| 517 | x = x + self.positional_embedding |
| 518 | x = x.permute(1, 0, 2) # NLD -> LND |
| 519 | x = self.transformer(x) |
| 520 | x = x.permute(1, 0, 2) # LND -> NLD |
| 521 | x = self.ln_final(x) |
| 522 | x = x[torch.arange(x.shape[0]), eos_indx] @ self.text_projection |
| 523 | x = x.reshape(B, K, self.embed_dim) # 1 19 512 |
| 524 | return x |
| 525 | |
| 526 | else: |
| 527 | x = self.token_embedding(text) # [batch_size, n_ctx, d_model] |
| 528 | x = x + self.positional_embedding |
| 529 | x = x.permute(1, 0, 2) # NLD -> LND |
| 530 | x = self.transformer(x) |
| 531 | x = x.permute(1, 0, 2) # LND -> NLD |
| 532 | x = self.ln_final(x) |
| 533 | x = x[torch.arange(x.shape[0]), text.argmax(dim=-1)] @ self.text_projection |
| 534 | # x = self.out_proj(x) |
| 535 | return x |
| 536 | |
| 537 | @BACKBONES.register_module() |
| 538 | class ContextDecoder(nn.Module): |
nothing calls this directly
no outgoing calls
no test coverage detected