| 81 | self.hf_plugin_mode = is_plugin |
| 82 | |
| 83 | def forward(self, *args, **kwargs): |
| 84 | encoder_in = args[0] |
| 85 | |
| 86 | if self.hf_plugin_mode: |
| 87 | # In this mode, assume that HuggingFace's LongformerModel.encoder has been |
| 88 | # substituted to this class's instance |
| 89 | extended_attention_mask = kwargs['attention_mask'] |
| 90 | local_attn_mask = torch.zeros_like(extended_attention_mask) |
| 91 | local_attn_mask[extended_attention_mask > -10000.] = 1.0 |
| 92 | global_attn_mask = torch.zeros_like(extended_attention_mask) |
| 93 | global_attn_mask[extended_attention_mask > 0.] = 1.0 |
| 94 | output = self.ft_encoder.forward(encoder_in, local_attn_mask, global_attn_mask, self.all_weight, 0) |
| 95 | return LongformerBaseModelOutput( |
| 96 | last_hidden_state=output, |
| 97 | hidden_states=None, |
| 98 | attentions=None, |
| 99 | global_attentions=None, |
| 100 | ) |
| 101 | else: |
| 102 | local_attn_mask = args[1] |
| 103 | global_attn_mask = args[2] |
| 104 | return self.ft_encoder.forward(encoder_in, local_attn_mask, global_attn_mask, self.all_weight, 0) |