From 8a9e5dc41f9df3bef7f849c9aab638914669a954 Mon Sep 17 00:00:00 2001 From: Tom Aarsen Date: Tue, 3 Dec 2024 12:18:34 +0100 Subject: [PATCH] Save custom module kwargs if specified This should have been included in the save all along --- sentence_transformers/SentenceTransformer.py | 6 +++++- 1 file changed, 5 insertions(+), 1 deletion(-) diff --git a/sentence_transformers/SentenceTransformer.py b/sentence_transformers/SentenceTransformer.py index 830564a57..15cf36cba 100644 --- a/sentence_transformers/SentenceTransformer.py +++ b/sentence_transformers/SentenceTransformer.py @@ -1187,7 +1187,11 @@ def save( # For other cases, we want to add the class name: elif not class_ref.startswith("sentence_transformers."): class_ref = f"{class_ref}.{type(module).__name__}" - modules_config.append({"idx": idx, "name": name, "path": os.path.basename(model_path), "type": class_ref}) + + module_config = {"idx": idx, "name": name, "path": os.path.basename(model_path), "type": class_ref} + if self.module_kwargs and name in self.module_kwargs and (module_kwargs := self.module_kwargs[name]): + module_config["kwargs"] = module_kwargs + modules_config.append(module_config) with open(os.path.join(path, "modules.json"), "w") as fOut: json.dump(modules_config, fOut, indent=2)