(self, model)
| 114 | return cpu_index, gpu_index |
| 115 | |
| 116 | def break_into(self, model): |
| 117 | self.model = model |
| 118 | model.broken_into = True |
| 119 | self.reconstruct_index, self.index = self.setup_faiss() |
| 120 | self.is_encoder_decoder = model.config.is_encoder_decoder |
| 121 | |
| 122 | # Inject our pre_forward_hook to capture the labels at every forward pass |
| 123 | self.original_forward_func = model.forward |
| 124 | |
| 125 | # Create a wrapper that calls pre_forward_hook, ensuring that self still refers to KNNSaverMulti |
| 126 | def forward_wrapper(input_ids=None, attention_mask=None, labels=None, **kwargs): |
| 127 | return self.pre_forward_hook(input_ids=input_ids, attention_mask=attention_mask, labels=labels, **kwargs) |
| 128 | |
| 129 | # Override the model's forward with our wrapper |
| 130 | model.forward = forward_wrapper |
| 131 | |
| 132 | # Inject our activation_capturer to capture the activations at every forward pass |
| 133 | layer_to_capture_fn, capture_input = KNNWrapperMulti.model_layer_to_capture[model.config.model_type][self.knn_keytype] |
| 134 | layer_to_capture = layer_to_capture_fn(model) |
| 135 | self.activation_capturer = ActivationCapturer(layer_to_capture, capture_input=capture_input) |
| 136 | self.register_hook(layer_to_capture, self.activation_capturer) |
| 137 | |
| 138 | # Inject our main function after the model's final layer |
| 139 | final_layer = KNNWrapperMulti.get_model_last_layer(model.config.model_type)(model) |
| 140 | self.register_hook(final_layer, self.post_forward_hook) |
| 141 | self.vocab_size = final_layer.out_features |
| 142 | |
| 143 | def get_knns(self, queries): |
| 144 | if not self.knn_gpu: |
nothing calls this directly
no test coverage detected