store all relevant info in safetensors

Author: stephantul

model2vec/hf_utils.py

@@ -25,6 +25,7 @@ def save_pretrained(
     create_model_card: bool = True,
     subfolder: str | None = None,
     weights: np.ndarray | None = None,
+    mapping: np.ndarray | None = None,
     **kwargs: Any,
 ) -> None:
     """
@@ -37,6 +38,7 @@ def save_pretrained(
     :param create_model_card: Whether to create a model card.
     :param subfolder: The subfolder to save the model in.
     :param weights: The weights of the model. If None, no weights are saved.
+    :param mapping: The token mapping of the model. If None, there is no token mapping.
     :param **kwargs: Any additional arguments.
     """
     folder_path = folder_path / subfolder if subfolder else folder_path
@@ -45,6 +47,8 @@ def save_pretrained(
     model_weights = {"embeddings": embeddings}
     if weights is not None:
         model_weights["weights"] = weights
+    if mapping is not None:
+        model_weights["mapping"] = mapping
 
     save_file(model_weights, folder_path / "model.safetensors")
     tokenizer.save(str(folder_path / "tokenizer.json"), pretty=False)
@@ -106,7 +110,7 @@ def load_pretrained(
     subfolder: str | None = None,
     token: str | None = None,
     from_sentence_transformers: bool = False,
-) -> tuple[np.ndarray, Tokenizer, dict[str, Any], dict[str, Any], np.ndarray | None]:
+) -> tuple[np.ndarray, Tokenizer, dict[str, Any], dict[str, Any], np.ndarray | None, np.ndarray | None]:
     """
     Loads a pretrained model from a folder.
 
@@ -185,18 +189,23 @@ def load_pretrained(
     if from_sentence_transformers:
         embeddings = opened_tensor_file.get_tensor("embedding.weight")
         weights = None
+        mapping = None
     else:
         embeddings = opened_tensor_file.get_tensor("embeddings")
         try:
             weights = opened_tensor_file.get_tensor("weights")
         except Exception:
             # Bare except because safetensors does not export its own errors.
             weights = None
+        try:
+            mapping = opened_tensor_file.get_tensor("mapping")
+        except Exception:
+            mapping = None
 
     tokenizer: Tokenizer = Tokenizer.from_file(str(tokenizer_path))
     config = json.load(open(config_path))
 
-    return embeddings, tokenizer, config, metadata, weights
+    return embeddings, tokenizer, config, metadata, weights, mapping
 
 
 def _get_metadata_from_readme(readme_path: Path) -> dict[str, Any]:

model2vec/model.py

@@ -30,7 +30,7 @@ def __init__(
         base_model_name: str | None = None,
         language: list[str] | None = None,
         weights: np.ndarray | None = None,
-        token_mapping: list[int] | None = None,
+        token_mapping: np.ndarray | None = None,
     ) -> None:
         """
         Initialize the StaticModel.
@@ -63,7 +63,7 @@ def __init__(
         self.weights = weights
         # Convert to an array for fast lookups
         # We can't use or short circuit here because np.ndarray as booleans are ambiguous.
-        self.token_mapping: np.ndarray | None = None if token_mapping is None else np.asarray(token_mapping)
+        self.token_mapping: np.ndarray | None = token_mapping
 
         self.tokenizer = tokenizer
         self.unk_token_id: int | None
@@ -121,9 +121,6 @@ def save_pretrained(self, path: PathLike, model_name: str | None = None, subfold
         """
         from model2vec.hf_utils import save_pretrained
 
-        if self.token_mapping is not None:
-            self.config["token_mapping"] = self.token_mapping.tolist()
-
         save_pretrained(
             folder_path=Path(path),
             embeddings=self.embedding,
@@ -134,6 +131,7 @@ def save_pretrained(self, path: PathLike, model_name: str | None = None, subfold
             model_name=model_name,
             subfolder=subfolder,
             weights=self.weights,
+            mapping=self.token_mapping,
         )
 
     def tokenize(self, sentences: Sequence[str], max_length: int | None = None) -> list[list[int]]:
@@ -490,11 +488,10 @@ def load_local(cls: type[StaticModel], path: PathLike) -> StaticModel:
         if not path.is_dir():
             raise ValueError(f"Path {path} is not a directory.")
 
-        embeddings, tokenizer, config, weights = load_local_model(path)
-        token_mapping = cast(list[int], config.pop("token_mapping", None))
+        embeddings, tokenizer, config, weights, mapping = load_local_model(path)
 
         return StaticModel(
-            vectors=embeddings, tokenizer=tokenizer, config=config, weights=weights, token_mapping=token_mapping
+            vectors=embeddings, tokenizer=tokenizer, config=config, weights=weights, token_mapping=mapping
         )
 
 
@@ -517,7 +514,7 @@ def quantize_model(
     """
     from model2vec.quantization import quantize_and_reduce_dim
 
-    token_mapping: list[int] | None
+    token_mapping: np.ndarray | None
     weights: np.ndarray | None
     if vocabulary_quantization is not None:
         from model2vec.vocabulary_quantization import quantize_vocabulary
@@ -530,7 +527,7 @@ def quantize_model(
         )
     else:
         embeddings = model.embedding
-        token_mapping = cast(list[int], model.token_mapping.tolist()) if model.token_mapping is not None else None
+        token_mapping = model.token_mapping
         weights = model.weights
     if quantize_to is not None or dimensionality is not None:
         embeddings = quantize_and_reduce_dim(
@@ -568,20 +565,18 @@ def _loading_helper(
     if from_sentence_transformers and subfolder is not None:
         raise ValueError("Subfolder is not supported for sentence transformers models.")
 
-    embeddings, tokenizer, config, metadata, weights = load_pretrained(
+    embeddings, tokenizer, config, metadata, weights, mapping = load_pretrained(
         folder_or_repo_path=path,
         token=token,
         from_sentence_transformers=from_sentence_transformers,
         subfolder=subfolder,
     )
 
-    token_mapping = config.pop("token_mapping", None)
-
     model = cls(
         vectors=embeddings,
         tokenizer=tokenizer,
         weights=weights,
-        token_mapping=token_mapping,
+        token_mapping=mapping,
         config=config,
         normalize=normalize,
         base_model_name=metadata.get("base_model"),

model2vec/train/base.py

@@ -150,7 +150,7 @@ def to_static_model(self) -> StaticModel:
         """Convert the model to a static model."""
         emb = self.embeddings.weight.detach().cpu().numpy()
         w = torch.sigmoid(self.w).detach().cpu().numpy()
-        token_mapping = self.token_mapping.tolist()
+        token_mapping = self.token_mapping.numpy()
 
         return StaticModel(
             vectors=emb, weights=w, tokenizer=self.tokenizer, normalize=True, token_mapping=token_mapping

model2vec/utils.py

@@ -104,7 +104,9 @@ def setup_logging() -> None:
     )
 
 
-def load_local_model(folder: Path) -> tuple[np.ndarray, Tokenizer, dict[str, str], np.ndarray | None]:
+def load_local_model(
+    folder: Path,
+) -> tuple[np.ndarray, Tokenizer, dict[str, str], np.ndarray | None, np.ndarray | None]:
     """Load a local model."""
     embeddings_path = folder / "model.safetensors"
     tokenizer_path = folder / "tokenizer.json"
@@ -117,6 +119,10 @@ def load_local_model(folder: Path) -> tuple[np.ndarray, Tokenizer, dict[str, str
     except Exception:
         # Bare except because safetensors does not export its own errors.
         weights = None
+    try:
+        mapping = opened_tensor_file.get_tensor("mapping")
+    except Exception:
+        mapping = None
 
     if config_path.exists():
         config = json.load(open(config_path))
@@ -125,4 +131,4 @@ def load_local_model(folder: Path) -> tuple[np.ndarray, Tokenizer, dict[str, str
 
     tokenizer: Tokenizer = Tokenizer.from_file(str(tokenizer_path))
 
-    return embeddings, tokenizer, config, weights
+    return embeddings, tokenizer, config, weights, mapping

model2vec/vocabulary_quantization.py

@@ -20,7 +20,7 @@
 
 def quantize_vocabulary(
     n_clusters: int, weights: np.ndarray | None, embeddings: np.ndarray
-) -> tuple[np.ndarray, list[int], np.ndarray]:
+) -> tuple[np.ndarray, np.ndarray, np.ndarray]:
     """Quantize the vocabulary of embeddings using KMeans clustering."""
     logger.info(f"Quantizing vocabulary to {n_clusters} clusters.")
     # If the model does not have weights, we assume the norm to be informative.
@@ -38,8 +38,7 @@ def quantize_vocabulary(
     # Fit KMeans to the embeddings
     kmeans.fit(cast_embeddings)
     # Create a mapping from the original token index to the cluster index
-    # Make sure to convert to list, otherwise we get np.int32 which is not jsonable.
-    token_mapping = cast(list[int], kmeans.predict(cast_embeddings).tolist())
+    token_mapping = kmeans.predict(cast_embeddings)
     # The cluster centers are the new embeddings.
     # Convert them back to the original dtype
     embeddings = kmeans.cluster_centers_.astype(orig_dtype)

tests/test_utils.py

@@ -98,7 +98,7 @@ def test_local_load(mock_tokenizer: Tokenizer, config: dict[str, Any], expected:
         mock_tokenizer.save(str(Path(tempdir) / "tokenizer.json"))
         if config is not None:
             json.dump(config, open(tempdir_path / "config.json", "w"))
-        arr, tokenizer, config, weights = load_local_model(tempdir_path)
+        arr, tokenizer, config, weights, _ = load_local_model(tempdir_path)
         assert config == expected
         assert tokenizer.to_str() == mock_tokenizer.to_str()
         assert arr.shape == x.shape