Update normalization

Author: stephantul

model2vec/model.py

@@ -22,7 +22,7 @@
 
 class StaticModel(nn.Module):
     def __init__(
-        self, vectors: np.ndarray, tokenizer: Tokenizer, config: dict[str, Any], normalize: bool = False
+        self, vectors: np.ndarray, tokenizer: Tokenizer, config: dict[str, Any], normalize: bool | None = None
     ) -> None:
         """
         Initialize the StaticModel.
@@ -49,7 +49,27 @@ def __init__(
             self.unk_token_id = None
 
         self.config = config
-        self.normalize = normalize
+        self.normalize = config.get("normalize", normalize if normalize is not None else False)
+
+    @property
+    def normalize(self) -> bool:
+        """
+        Get the normalize value.
+
+        :return: The normalize value.
+        """
+        return self._normalize
+
+    @normalize.setter
+    def normalize(self, value: bool) -> None:
+        """Update the config if the value of normalize changes."""
+        config_normalize = self.config.get("normalize", False)
+        self._normalize = value
+        if config_normalize is not None and value != config_normalize:
+            logger.warning(
+                f"Set normalization to `{value}`, which does not match config value `{config_normalize}`. Updating config."
+            )
+        self.config["normalize"] = value
 
     def save_pretrained(self, path: PathLike) -> None:
         """

tests/test_model.py

@@ -2,33 +2,29 @@
 
 import numpy as np
 import pytest
-from transformers import PreTrainedTokenizerFast
+from tokenizers import Tokenizer
 
 from model2vec import StaticModel
 
 
-def test_initialization(
-    mock_vectors: np.ndarray, mock_tokenizer: PreTrainedTokenizerFast, mock_config: dict[str, str]
-) -> None:
+def test_initialization(mock_vectors: np.ndarray, mock_tokenizer: Tokenizer, mock_config: dict[str, str]) -> None:
     """Test successful initialization of StaticModel."""
     model = StaticModel(vectors=mock_vectors, tokenizer=mock_tokenizer, config=mock_config)
-    assert model.vectors.shape == (5, 2)
+    assert model.embedding.weight.shape == (5, 2)
     assert len(model.tokens) == 5
     assert model.tokenizer == mock_tokenizer
     assert model.config == mock_config
 
 
-def test_initialization_token_vector_mismatch(
-    mock_tokenizer: PreTrainedTokenizerFast, mock_config: dict[str, str]
-) -> None:
+def test_initialization_token_vector_mismatch(mock_tokenizer: Tokenizer, mock_config: dict[str, str]) -> None:
     """Test if error is raised when number of tokens and vectors don't match."""
     mock_vectors = np.array([[0.1, 0.2], [0.2, 0.3]])
     with pytest.raises(ValueError):
         StaticModel(vectors=mock_vectors, tokenizer=mock_tokenizer, config=mock_config)
 
 
 def test_encode_single_sentence(
-    mock_vectors: np.ndarray, mock_tokenizer: PreTrainedTokenizerFast, mock_config: dict[str, str]
+    mock_vectors: np.ndarray, mock_tokenizer: Tokenizer, mock_config: dict[str, str]
 ) -> None:
     """Test encoding of a single sentence."""
     model = StaticModel(vectors=mock_vectors, tokenizer=mock_tokenizer, config=mock_config)
@@ -37,7 +33,7 @@ def test_encode_single_sentence(
 
 
 def test_encode_multiple_sentences(
-    mock_vectors: np.ndarray, mock_tokenizer: PreTrainedTokenizerFast, mock_config: dict[str, str]
+    mock_vectors: np.ndarray, mock_tokenizer: Tokenizer, mock_config: dict[str, str]
 ) -> None:
     """Test encoding of multiple sentences."""
     model = StaticModel(vectors=mock_vectors, tokenizer=mock_tokenizer, config=mock_config)
@@ -46,24 +42,29 @@ def test_encode_multiple_sentences(
 
 
 def test_encode_empty_sentence(
-    mock_vectors: np.ndarray, mock_tokenizer: PreTrainedTokenizerFast, mock_config: dict[str, str]
+    mock_vectors: np.ndarray, mock_tokenizer: Tokenizer, mock_config: dict[str, str]
 ) -> None:
     """Test encoding with an empty sentence."""
     model = StaticModel(vectors=mock_vectors, tokenizer=mock_tokenizer, config=mock_config)
     encoded = model.encode("")
     assert np.array_equal(encoded, np.zeros((2,)))
 
 
-def test_normalize() -> None:
+def test_normalize(mock_vectors: np.ndarray, mock_tokenizer: Tokenizer, mock_config: dict[str, str]) -> None:
     """Test normalization of vectors."""
-    X = np.array([[3, 4], [1, 2], [0, 0]])
-    normalized = StaticModel.normalize(X)
-    expected = np.array([[0.6, 0.8], [0.4472136, 0.89442719], [0, 0]])
+    s = "word1 word2 word3"
+    model = StaticModel(vectors=mock_vectors, tokenizer=mock_tokenizer, config=mock_config, normalize=False)
+    X = model.encode(s)
+    model = StaticModel(vectors=mock_vectors, tokenizer=mock_tokenizer, config=mock_config, normalize=True)
+    normalized = model.encode(s)
+
+    expected = X / np.linalg.norm(X)
+
     np.testing.assert_almost_equal(normalized, expected)
 
 
 def test_save_pretrained(
-    tmp_path: Path, mock_vectors: np.ndarray, mock_tokenizer: PreTrainedTokenizerFast, mock_config: dict[str, str]
+    tmp_path: Path, mock_vectors: np.ndarray, mock_tokenizer: Tokenizer, mock_config: dict[str, str]
 ) -> None:
     """Test saving a pretrained model."""
     model = StaticModel(vectors=mock_vectors, tokenizer=mock_tokenizer, config=mock_config)
@@ -80,7 +81,7 @@ def test_save_pretrained(
 
 
 def test_load_pretrained(
-    tmp_path: Path, mock_vectors: np.ndarray, mock_tokenizer: PreTrainedTokenizerFast, mock_config: dict[str, str]
+    tmp_path: Path, mock_vectors: np.ndarray, mock_tokenizer: Tokenizer, mock_config: dict[str, str]
 ) -> None:
     """Test loading a pretrained model after saving it."""
     # Save the model to a temporary path
@@ -92,6 +93,6 @@ def test_load_pretrained(
     loaded_model = StaticModel.from_pretrained(save_path)
 
     # Assert that the loaded model has the same properties as the original one
-    np.testing.assert_array_equal(loaded_model.vectors, mock_vectors)
+    np.testing.assert_array_equal(loaded_model.embedding.weight.numpy(), mock_vectors)
     assert loaded_model.tokenizer.get_vocab() == mock_tokenizer.get_vocab()
     assert loaded_model.config == mock_config