| 185 | pass |
| 186 | |
| 187 | class EncoderWrapper(torch.nn.Module): |
| 188 | def __init__(self, encoder,use_checkpoint=False): |
| 189 | super().__init__() |
| 190 | self.encoder = encoder |
| 191 | |
| 192 | try: |
| 193 | self.main_input_name = encoder.main_input_name |
| 194 | except: |
| 195 | pass |
| 196 | apply_checkpoint_wrapper(self.encoder, use_checkpoint) |
| 197 | |
| 198 | def forward(self, input_ids=None, attention_mask=None,**kwargs): |
| 199 | bsz, total_length = input_ids.shape |
| 200 | passage_length = total_length // self.n_passages |
| 201 | # total_input |
| 202 | input_ids = input_ids.view(bsz*self.n_passages, passage_length) |
| 203 | attention_mask = attention_mask.view(bsz*self.n_passages, passage_length) |
| 204 | outputs = self.encoder(input_ids=input_ids, attention_mask=attention_mask, **kwargs) |
| 205 | outputs.last_hidden_state = outputs.last_hidden_state.view(bsz, self.n_passages*passage_length, -1) |
| 206 | return outputs |
no outgoing calls
no test coverage detected