| 424 | self.hook_handles.append(handle) |
| 425 | |
| 426 | def break_out(self): |
| 427 | # Close the Arrow writer (main process only) |
| 428 | if self.process_index == 0 and hasattr(self, 'arrow_writer'): |
| 429 | self.arrow_writer.close() |
| 430 | self.arrow_file.close() |
| 431 | |
| 432 | # Remove hooks |
| 433 | for h in self.hook_handles: |
| 434 | h.remove() |
| 435 | |
| 436 | if self.model is not None and hasattr(self.model, 'broken_into') and self.model.broken_into is not None: |
| 437 | self.model.forward = self.original_forward_func |
| 438 | self.model.broken_into = None |
| 439 | |
| 440 | def build_index(self, num_keys_to_add_at_a_time=1_000_000, ncentroids=4096, seed=42, code_size=32, probe=8): |
| 441 | logger.info('Loading Dataset...') |