(self, model)
| 371 | self.accelerator.wait_for_everyone() |
| 372 | |
| 373 | def break_into(self, model): |
| 374 | self.model = model |
| 375 | model.broken_into = True |
| 376 | self.is_encoder_decoder = model.config.is_encoder_decoder |
| 377 | |
| 378 | # Inject our activation_capturer to capture the activations at every forward pass |
| 379 | layer_to_capture_fn, capture_input = KNNWrapperMulti.model_layer_to_capture[model.config.model_type][self.knn_keytype] |
| 380 | layer_to_capture = layer_to_capture_fn(model) |
| 381 | self.activation_capturer = ActivationCapturer(layer_to_capture, capture_input=capture_input) |
| 382 | self.register_hook(layer_to_capture, self.activation_capturer) |
| 383 | |
| 384 | # Inject our pre_forward_hook to capture the labels at every forward pass |
| 385 | self.original_forward_func = model.forward |
| 386 | |
| 387 | # Create a wrapper that calls pre_forward_hook |
| 388 | def forward_wrapper(input_ids=None, attention_mask=None, labels=None, **kwargs): |
| 389 | return self.pre_forward_hook(input_ids=input_ids, attention_mask=attention_mask, labels=labels, **kwargs) |
| 390 | |
| 391 | # Override the model's forward with our wrapper |
| 392 | model.forward = forward_wrapper |
| 393 | |
| 394 | # Inject our main function after the model's final layer |
| 395 | final_layer = KNNWrapperMulti.get_model_last_layer(model.config.model_type)(model) |
| 396 | self.register_hook(final_layer, self.post_forward_hook) |
| 397 | |
| 398 | # Set up Arrow writer now that we have the model info |
| 399 | self.arrow_file_path = self._get_arrow_file_path() |
| 400 | if os.path.exists(self.arrow_file_path): |
| 401 | logger.info(f'Arrow file already exists at {self.arrow_file_path}.') |
| 402 | else: |
| 403 | logger.info(f"Creating Arrow file at {self.arrow_file_path}.") |
| 404 | self._setup_arrow_writer() |
| 405 | |
| 406 | # Create directory if it doesn't exist |
| 407 | Path(self.dstore_dir).mkdir(parents=True, exist_ok=True) |
| 408 | |
| 409 | def pre_forward_hook(self, input_ids=None, attention_mask=None, labels=None, **kwargs): |
| 410 | if labels is None: |
no test coverage detected