(self, queries)
| 141 | self.vocab_size = final_layer.out_features |
| 142 | |
| 143 | def get_knns(self, queries): |
| 144 | if not self.knn_gpu: |
| 145 | queries = queries.cpu() |
| 146 | dists, knns = self.index.search(queries.to(torch.float32), self.k) |
| 147 | dists, knns = dists.to(self.device), knns.to(self.device) |
| 148 | return dists, knns |
| 149 | |
| 150 | def pre_forward_hook(self, input_ids=None, attention_mask=None, labels=None, **kwargs): |
| 151 | self.labels = labels |