| 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: |
| 411 | raise ValueError('labels must be provided when saving a datastore. Are you using --predict_with_generate by mistake? If so, disable it') |
| 412 | self.labels = labels |
| 413 | return self.original_forward_func(input_ids=input_ids, labels=labels, attention_mask=attention_mask, **kwargs) |
| 414 | |
| 415 | def post_forward_hook(self, module, input, output): |
| 416 | captured_keys = self.activation_capturer.captured |