Skip to content

Commit

Permalink
Save custom module kwargs if specified
Browse files Browse the repository at this point in the history
This should have been included in the save all along
  • Loading branch information
tomaarsen committed Dec 3, 2024
1 parent ba8cb2e commit 8a9e5dc
Showing 1 changed file with 5 additions and 1 deletion.
6 changes: 5 additions & 1 deletion sentence_transformers/SentenceTransformer.py
Original file line number Diff line number Diff line change
Expand Up @@ -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)
Expand Down

0 comments on commit 8a9e5dc

Please sign in to comment.