Skip to content

Commit

Permalink
Merge branch 'bump-version' of https://github.com/MinishLab/model2vec
Browse files Browse the repository at this point in the history
…into bump-version
  • Loading branch information
Pringled committed Jan 21, 2025
2 parents 6b63c6a + 14176b2 commit d003fa8
Show file tree
Hide file tree
Showing 3 changed files with 25 additions and 17 deletions.
11 changes: 9 additions & 2 deletions model2vec/hf_utils.py
Original file line number Diff line number Diff line change
Expand Up @@ -39,7 +39,14 @@ def save_pretrained(
folder_path.mkdir(exist_ok=True, parents=True)
save_file({"embeddings": embeddings}, folder_path / "model.safetensors")
tokenizer.save(str(folder_path / "tokenizer.json"))
json.dump(config, open(folder_path / "config.json", "w"))
json.dump(config, open(folder_path / "config.json", "w"), indent=4)

# Create modules.json
modules = [{"idx": 0, "name": "0", "path": ".", "type": "sentence_transformers.models.StaticEmbedding"}]
if config.get("normalize"):
# If normalize=True, add sentence_transformers.models.Normalize
modules.append({"idx": 1, "name": "1", "path": "1_Normalize", "type": "sentence_transformers.models.Normalize"})
json.dump(modules, open(folder_path / "modules.json", "w"), indent=4)

logger.info(f"Saved model to {folder_path}")

Expand Down Expand Up @@ -75,7 +82,7 @@ def _create_model_card(
base_model=base_model_name,
license=license,
language=language,
tags=["embeddings", "static-embeddings"],
tags=["embeddings", "static-embeddings", "sentence-transformers"],
library_name="model2vec",
**kwargs,
)
Expand Down
1 change: 1 addition & 0 deletions tests/test_model.py
Original file line number Diff line number Diff line change
Expand Up @@ -161,6 +161,7 @@ def test_save_pretrained(
assert (save_path / "model.safetensors").exists()
assert (save_path / "tokenizer.json").exists()
assert (save_path / "config.json").exists()
assert (save_path / "modules.json").exists()


def test_load_pretrained(
Expand Down
30 changes: 15 additions & 15 deletions uv.lock

Some generated files are not rendered by default. Learn more about how customized files appear on GitHub.

0 comments on commit d003fa8

Please sign in to comment.