| 59 | |
| 60 | @staticmethod |
| 61 | def encode(model, input_ids, attention_mask, gck_segment=32): |
| 62 | dummy_tensor = torch.ones(1, dtype=torch.float32, requires_grad=True) |
| 63 | pooled_output = [] |
| 64 | for mini_batch in range(0, input_ids.shape[0], gck_segment): |
| 65 | mini_batch_input_ids = input_ids[mini_batch:mini_batch |
| 66 | + gck_segment] |
| 67 | mini_batch_attention_mask = attention_mask[mini_batch:mini_batch |
| 68 | + gck_segment] |
| 69 | mini_batch_pooled_output = checkpoint(model, mini_batch_input_ids, |
| 70 | mini_batch_attention_mask, |
| 71 | dummy_tensor) |
| 72 | pooled_output.append(mini_batch_pooled_output) |
| 73 | return torch.cat(pooled_output, dim=0) |
| 74 | |
| 75 | def forward(self, |
| 76 | query_input_ids, |