| 9 | super().__init__(vocab_size=vocab_size) |
| 10 | |
| 11 | def forward(self, input_ids, clip_skip=2, extra_mask=None): |
| 12 | embeds = self.token_embedding(input_ids) |
| 13 | embeds = embeds + self.position_embeds.to(dtype=embeds.dtype, device=input_ids.device) |
| 14 | attn_mask = self.attn_mask.to(device=embeds.device, dtype=embeds.dtype) |
| 15 | if extra_mask is not None: |
| 16 | attn_mask[:, extra_mask[0]==0] = float("-inf") |
| 17 | for encoder_id, encoder in enumerate(self.encoders): |
| 18 | embeds = encoder(embeds, attn_mask=attn_mask) |
| 19 | if encoder_id + clip_skip == len(self.encoders): |
| 20 | hidden_states = embeds |
| 21 | embeds = self.final_layer_norm(embeds) |
| 22 | pooled_embeds = embeds[torch.arange(embeds.shape[0]), input_ids.to(dtype=torch.int).argmax(dim=-1)] |
| 23 | return pooled_embeds, hidden_states |
| 24 | |
| 25 | @staticmethod |
| 26 | def state_dict_converter(): |