(self, text, normalized=False,
hidden_z: Optional[torch.Tensor] = None,
heads_z: Optional[torch.Tensor] = None,
mha_z: Optional[torch.Tensor] = None,
intermediate_z: Optional[torch.Tensor] = None,
ffn_z: Optional[torch.Tensor] = None,
embed_dim_z: Optional[torch.Tensor] = None,
)
| 762 | return mask |
| 763 | |
| 764 | def encode_text(self, text, normalized=False, |
| 765 | hidden_z: Optional[torch.Tensor] = None, |
| 766 | heads_z: Optional[torch.Tensor] = None, |
| 767 | mha_z: Optional[torch.Tensor] = None, |
| 768 | intermediate_z: Optional[torch.Tensor] = None, |
| 769 | ffn_z: Optional[torch.Tensor] = None, |
| 770 | embed_dim_z: Optional[torch.Tensor] = None, |
| 771 | ): |
| 772 | self.hidden_z = hidden_z |
| 773 | self.embed_dim_z = embed_dim_z |
| 774 | |
| 775 | text = text.to(self.token_embedding.weight.device) |
| 776 | x = self.token_embedding(text) # [batch_size, n_ctx, d_model] |
| 777 | |
| 778 | x = x + self.positional_embedding |
| 779 | if hidden_z is not None: |
| 780 | x = torch.mul(x, hidden_z) |
| 781 | |
| 782 | x = x.permute(1, 0, 2) # NLD -> LND |
| 783 | x = self.transformer(x, attn_mask=self.attn_mask, |
| 784 | hidden_z=hidden_z, |
| 785 | heads_z=heads_z, |
| 786 | mha_z=mha_z, |
| 787 | intermediate_z=intermediate_z, |
| 788 | ffn_z=ffn_z) |
| 789 | x = x.permute(1, 0, 2) # LND -> NLD |
| 790 | x = self.ln_final(x, hidden_z) |
| 791 | |
| 792 | # if hidden_z is not None: |
| 793 | # x = torch.mul(x, hidden_z) |
| 794 | |
| 795 | x = x[torch.arange(x.shape[0]), text.argmax(dim=-1)] |
| 796 | |
| 797 | # x.shape = [batch_size, n_ctx, transformer.width] |
| 798 | # take features from the eot embedding (eot_token is the highest number in each sequence) |
| 799 | x = self.get_proj_feature(x) |
| 800 | if embed_dim_z is not None: |
| 801 | x = x.mul(embed_dim_z) |
| 802 | if normalized: |
| 803 | x = F.normalize(x, dim=-1) |
| 804 | |
| 805 | return x |
| 806 | |
| 807 | def get_proj_feature(self, x): |
| 808 | return x @ self.text_projection |
no test coverage detected