Load retriever state from disk
(self, retriever_cache_file: str, retriever_cache_embeddings_file: str)
| 623 | pickle.dump(state, f) |
| 624 | |
| 625 | def load(self, retriever_cache_file: str, retriever_cache_embeddings_file: str): |
| 626 | """Load retriever state from disk""" |
| 627 | print(f"Loading retriever from {retriever_cache_file} and {retriever_cache_embeddings_file}") |
| 628 | |
| 629 | # Load embeddings |
| 630 | if os.path.exists(retriever_cache_embeddings_file): |
| 631 | print(f"Loading embeddings from {retriever_cache_embeddings_file}") |
| 632 | self.embeddings = np.load(retriever_cache_embeddings_file) |
| 633 | print(f"Embeddings shape: {self.embeddings.shape}") |
| 634 | else: |
| 635 | print(f"Embeddings file not found: {retriever_cache_embeddings_file}") |
| 636 | |
| 637 | # Load other attributes |
| 638 | if os.path.exists(retriever_cache_file): |
| 639 | print(f"Loading corpus from {retriever_cache_file}") |
| 640 | with open(retriever_cache_file, 'rb') as f: |
| 641 | state = pickle.load(f) |
| 642 | self.corpus = state['corpus'] |
| 643 | self.document_ids = state['document_ids'] |
| 644 | print(f"Loaded corpus with {len(self.corpus)} documents") |
| 645 | else: |
| 646 | print(f"Corpus file not found: {retriever_cache_file}") |
| 647 | |
| 648 | return self |
| 649 | |
| 650 | @classmethod |
| 651 | def load_from_local_memory(cls, memories: Dict, model_name: str) -> 'SimpleEmbeddingRetriever': |