wip

Author: stephantul

model2vec/distill/distillation.py

@@ -30,6 +30,7 @@ def distill_from_model(
     token_remove_pattern: str | None = r"\[unused\d+\]",
     quantize_to: DType | str = DType.Float16,
     use_subword: bool | None = None,
+    vocabulary_quantization: int | None = None,
 ) -> StaticModel:
     """
     Distill a staticmodel from a sentence transformer.
@@ -113,14 +114,21 @@ def distill_from_model(
         tokenized=token_ids, model=model, device=device, pad_token_id=tokenizer.get_vocab()[pad_token]
     )
 
-    _, weights = post_process_embeddings(np.asarray(embeddings), None, sif_coefficient=sif_coefficient)
-    km = KMeans(4096, random_state=42)
-    km.fit(embeddings)
-    clustered_embeddings = km.predict(embeddings)
-    mapping = {idx: x for idx, x in enumerate(clustered_embeddings)}
+    if vocabulary_quantization is not None:
+        _, weights = post_process_embeddings(np.asarray(embeddings), None, sif_coefficient=sif_coefficient)
+        km = KMeans(vocabulary_quantization, random_state=42)
+        km.fit(embeddings)
+        clustered_embeddings = km.predict(embeddings)
+        mapping = {idx: x for idx, x in enumerate(clustered_embeddings)}
 
-    embeddings = km.cluster_centers_
-    embeddings, _ = post_process_embeddings(embeddings, pca_dims, sif_coefficient=sif_coefficient)
+        embeddings = km.cluster_centers_
+        embeddings, _ = post_process_embeddings(embeddings, pca_dims, sif_coefficient=sif_coefficient)
+    else:
+        # Post-process the embeddings.
+        embeddings, weights = post_process_embeddings(
+            np.asarray(embeddings), pca_dims, sif_coefficient=sif_coefficient
+        )
+        mapping = {idx: token.form for idx, token in enumerate(all_tokens)}
     # Quantize the embeddings.
     embeddings = quantize_embeddings(embeddings, quantize_to)
 
@@ -219,6 +227,7 @@ def distill(
     trust_remote_code: bool = False,
     quantize_to: DType | str = DType.Float16,
     use_subword: bool | None = None,
+    vocabulary_quantization: int | None = None,
 ) -> StaticModel:
     """
     Distill a staticmodel from a sentence transformer.
@@ -263,4 +272,5 @@ def distill(
         sif_coefficient=sif_coefficient,
         quantize_to=quantize_to,
         use_subword=use_subword,
+        vocabulary_quantization=vocabulary_quantization,
     )

model2vec/hf_utils.py

@@ -24,6 +24,7 @@ def save_pretrained(
     config: dict[str, Any],
     create_model_card: bool = True,
     subfolder: str | None = None,
+    weights: np.ndarray | None = None,
     **kwargs: Any,
 ) -> None:
     """
@@ -39,7 +40,12 @@ def save_pretrained(
     """
     folder_path = folder_path / subfolder if subfolder else folder_path
     folder_path.mkdir(exist_ok=True, parents=True)
-    save_file({"embeddings": embeddings}, folder_path / "model.safetensors")
+
+    model_weights = {"embeddings": embeddings}
+    if weights is not None:
+        model_weights["weights"] = weights
+
+    save_file(model_weights, folder_path / "model.safetensors")
     tokenizer.save(str(folder_path / "tokenizer.json"), pretty=False)
     json.dump(config, open(folder_path / "config.json", "w"), indent=4)
 
@@ -99,7 +105,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]]:
+) -> tuple[np.ndarray, Tokenizer, dict[str, Any], dict[str, Any], np.ndarray | None]:
     """
     Loads a pretrained model from a folder.
 
@@ -177,8 +183,14 @@ def load_pretrained(
     opened_tensor_file = cast(SafeOpenProtocol, safetensors.safe_open(embeddings_path, framework="numpy"))
     if from_sentence_transformers:
         embeddings = opened_tensor_file.get_tensor("embedding.weight")
+        weights = 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
 
     tokenizer: Tokenizer = Tokenizer.from_file(str(tokenizer_path))
     config = json.load(open(config_path))
@@ -188,7 +200,7 @@ def load_pretrained(
             f"Number of tokens does not match number of embeddings: `{len(tokenizer.get_vocab())}` vs `{len(embeddings)}`"
         )
 
-    return embeddings, tokenizer, config, metadata
+    return embeddings, tokenizer, config, metadata, weights
 
 
 def _get_metadata_from_readme(readme_path: Path) -> dict[str, Any]:

model2vec/model.py

@@ -24,7 +24,7 @@ class StaticModel:
     def __init__(
         self,
         vectors: np.ndarray,
-        weights: np.ndarray,
+        weights: np.ndarray | None,
         token_mapping: dict[int, int],
         tokenizer: Tokenizer,
         config: dict[str, Any] | None = None,
@@ -107,6 +107,8 @@ def save_pretrained(self, path: PathLike, model_name: str | None = None, subfold
         """
         from model2vec.hf_utils import save_pretrained
 
+        self.config["token_mapping"] = list(self.token_mapping.items())
+
         save_pretrained(
             folder_path=Path(path),
             embeddings=self.embedding,
@@ -116,6 +118,7 @@ def save_pretrained(self, path: PathLike, model_name: str | None = None, subfold
             language=self.language,
             model_name=model_name,
             subfolder=subfolder,
+            weights=self.weights,
         )
 
     def tokenize(self, sentences: Sequence[str], max_length: int | None = None) -> list[list[int]]:
@@ -131,8 +134,6 @@ def tokenize(self, sentences: Sequence[str], max_length: int | None = None) -> l
             m = max_length * self.median_token_length
             sentences = [sentence[:m] for sentence in sentences]
 
-        max_len = max([len(sentence) for sentence in sentences])
-        # self.tokenizer.model.max_input_chars_per_word = max_len + 1
         if self._can_encode_fast:
             encodings: list[Encoding] = self.tokenizer.encode_batch_fast(sentences, add_special_tokens=False)
         else:
@@ -159,6 +160,7 @@ def from_pretrained(
         subfolder: str | None = None,
         quantize_to: str | DType | None = None,
         dimensionality: int | None = None,
+        vocabulary_quantization: int | None = None,
     ) -> StaticModel:
         """
         Load a StaticModel from a local path or huggingface hub path.
@@ -178,36 +180,45 @@ def from_pretrained(
         """
         from model2vec.hf_utils import load_pretrained
 
-        embeddings, tokenizer, config, metadata = load_pretrained(
+        embeddings, tokenizer, config, metadata, weights = load_pretrained(
             folder_or_repo_path=path,
             token=token,
             from_sentence_transformers=False,
             subfolder=subfolder,
         )
 
-        weights = np.linalg.norm(embeddings, axis=1, keepdims=True) + 1e-32
-        embeddings = embeddings / weights
-
-        """from sklearn.cluster import KMeans
-        from sklearn.decomposition import PCA
-        km = KMeans(n_clusters=4096, random_state=0)
-        km.fit(embeddings)
-        # Do PCA again?
-        assignments = km.predict(embeddings)
-        embeddings = km.cluster_centers_
-
-        p = PCA(n_components=dimensionality)
-        embeddings = p.fit_transform(embeddings)
-
-        token_mapping = {i: x for i, x in enumerate(assignments)}"""
-        token_mapping = {i: i for i in range(len(embeddings))}
-
         embeddings = quantize_and_reduce_dim(
             embeddings=embeddings,
             quantize_to=quantize_to,
             dimensionality=dimensionality,
         )
 
+        if vocabulary_quantization is not None:
+            if len(embeddings) != len(tokenizer.get_vocab()):
+                raise ValueError(
+                    "Already quantized. "
+                )
+            
+            if weights is None:
+                weights = np.linalg.norm(embeddings, axis=1, keepdims=True) + 1e-32
+                embeddings = embeddings / weights
+
+            # Quantize the vocabulary
+            from sklearn.cluster import KMeans
+            kmeans = KMeans(n_clusters=vocabulary_quantization, random_state=42)
+            kmeans.fit(embeddings)
+            token_mapping = {idx: x for idx, x in enumerate(kmeans.predict(embeddings))}
+            embeddings = kmeans.cluster_centers_
+
+        else:
+            token_mapping = config.pop("token_mapping", None)
+            if isinstance(token_mapping, list):
+                # If the token mapping is a list, convert it to a dict
+                token_mapping = {int(k): int(v) for k, v in token_mapping}
+            elif token_mapping is None:
+                # If no token mapping is provided, use the default mapping
+                token_mapping = {i: i for i in range(len(embeddings))}
+
         return cls(
             embeddings,
             weights,
@@ -245,7 +256,7 @@ def from_sentence_transformers(
         """
         from model2vec.hf_utils import load_pretrained
 
-        embeddings, tokenizer, config, metadata = load_pretrained(
+        embeddings, tokenizer, config, metadata, weights = load_pretrained(
             folder_or_repo_path=path,
             token=token,
             from_sentence_transformers=True,
@@ -258,9 +269,10 @@ def from_sentence_transformers(
             dimensionality=dimensionality,
         )
 
-        weights = np.linalg.norm(embeddings, axis=1, keepdims=True) + 1e-32
-        embeddings = embeddings / weights
-        token_mapping = {i: i for i in range(len(embeddings))}
+        token_mapping = config.pop("token_mapping", None)
+        if token_mapping is None:
+            # If no token mapping is provided, use the default mapping
+            token_mapping = {i: i for i in range(len(embeddings))}
 
         return cls(
             embeddings,