Added embedding_dtype and vocabulary_quantization to config

Author: Pringled

model2vec/hf_utils.py

@@ -53,6 +53,9 @@ def save_pretrained(
 
     save_file(model_weights, folder_path / "model.safetensors")
     tokenizer.save(str(folder_path / "tokenizer.json"), pretty=False)
+
+    # Add embedding dtype to config
+    config["embedding_dtype"] = np.dtype(embeddings.dtype).name
     json.dump(config, open(folder_path / "config.json", "w"), indent=4)
 
     # Create modules.json

model2vec/model.py

@@ -506,6 +506,7 @@ def quantize_model(
         embeddings, token_mapping, weights = quantize_vocabulary(
             n_clusters=vocabulary_quantization, weights=model.weights, embeddings=model.embedding
         )
+        model.config["vocabulary_quantization"] = vocabulary_quantization
     else:
         embeddings = model.embedding
         token_mapping = model.token_mapping