Fixed dims bug and bumped version (#24)

MinishLab · Sep 23, 2024 · 5520cdc · 5520cdc
1 parent 9a0d0a9
commit 5520cdc
Show file tree

Hide file tree

Showing 2 changed files with 2 additions and 2 deletions.
diff --git a/model2vec/distill/inference.py b/model2vec/distill/inference.py
@@ -81,7 +81,7 @@ def create_output_embeddings_from_model_name_and_tokens(
 
 def _get_embedder_output_dim(output_value: OutputValue, embedder: SentenceTransformer) -> int:
     """Get the embeddings dimension of a sentence transformer, given an output value."""
-    embedder_output_dim = embedder.encode(["a"], show_progress_bar=False, output_value=output_value)[0].shape[0]
+    embedder_output_dim = embedder.encode(["a"], show_progress_bar=False, output_value=output_value)[0].shape[1]
 
     return embedder_output_dim
 

diff --git a/pyproject.toml b/pyproject.toml
@@ -3,7 +3,7 @@ name = "model2vec"
 description = "Distill a Small Fast Model from any Sentence Transformer"
 readme = { file = "README.md", content-type = "text/markdown" }
 license = { file = "LICENSE" }
-version = "0.1.1"
+version = "0.1.2"
 requires-python = ">=3.10"
 authors = [{ name = "Stéphan Tulkens", email = "[email protected]"}, {name = "Thomas van Dongen", email = "[email protected]"}]