From 3a49da8050d81b80d065d6dc8ec33e2231df9133 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?=E7=BD=97=E5=BB=B6=E6=A0=B9?= Date: Fri, 20 Dec 2024 08:54:03 +0800 Subject: [PATCH] metadata save and load --- retriv/experimental/advanced_retriever.py | 3 +++ 1 file changed, 3 insertions(+) diff --git a/retriv/experimental/advanced_retriever.py b/retriv/experimental/advanced_retriever.py index a0b17fd..a6183db 100755 --- a/retriv/experimental/advanced_retriever.py +++ b/retriv/experimental/advanced_retriever.py @@ -117,6 +117,7 @@ def save(self) -> None: "doc_lens": self.doc_lens, "relative_doc_lens": self.relative_doc_lens, "hyperparams": self.hyperparams, + "metadata": self.metadata, } np.savez_compressed(fr_state_path(self.index_name), state=state) @@ -145,6 +146,7 @@ def load(index_name: str = "new-index"): se.doc_lens = state["doc_lens"] se.relative_doc_lens = state["relative_doc_lens"] se.hyperparams = state["hyperparams"] + se.metadata = state["metadata"] state = { "init_args": se.init_args, @@ -155,6 +157,7 @@ def load(index_name: str = "new-index"): "doc_lens": se.doc_lens, "relative_doc_lens": se.relative_doc_lens, "hyperparams": se.hyperparams, + "metadata": se.metadata, } return se