diff --git a/model2vec/model.py b/model2vec/model.py index a13ab46..a55dfdb 100644 --- a/model2vec/model.py +++ b/model2vec/model.py @@ -80,6 +80,11 @@ def __init__( self.token_mapping: np.ndarray | None = token_mapping self.tokenizer = copy.deepcopy(tokenizer) + padding = self.tokenizer.padding + if padding is not None: + self.tokenizer.enable_padding( + pad_id=padding["pad_id"], pad_token=padding["pad_token"], pad_type_id=padding["pad_type_id"], length=0 + ) self.unk_token_id = _get_unk_token_id(self.tokenizer) self.median_token_length = int(np.median([len(token) for token in self.tokens])) diff --git a/model2vec/tokenizer/tokenizer.py b/model2vec/tokenizer/tokenizer.py index 62b5623..c1cb194 100644 --- a/model2vec/tokenizer/tokenizer.py +++ b/model2vec/tokenizer/tokenizer.py @@ -5,10 +5,20 @@ from typing import Sequence from skeletoken import TokenizerModel +from skeletoken.padding import Padding logger = logging.getLogger(__name__) +def _replace_padding(model: TokenizerModel) -> TokenizerModel: + """Replaces the model's padding module with a benign one.""" + padding = model.padding + if padding is not None: + model.padding = Padding(pad_id=padding.pad_id, pad_token=padding.pad_token, pad_type_id=0) + + return model + + def clean_and_create_vocabulary( model: TokenizerModel, vocabulary_to_add: Sequence[str], @@ -33,6 +43,7 @@ def clean_and_create_vocabulary( # Remove the post processor. model.post_processor = None + model = _replace_padding(model) internal_tokens: list[str] = model.sorted_vocabulary if token_remove_regex: diff --git a/tests/integration/data/sentence-transformers___all-MiniLM-L6-v2_baseline.json b/tests/integration/data/sentence-transformers___all-MiniLM-L6-v2_baseline.json index 9a7be03..7db8c18 100644 --- a/tests/integration/data/sentence-transformers___all-MiniLM-L6-v2_baseline.json +++ b/tests/integration/data/sentence-transformers___all-MiniLM-L6-v2_baseline.json @@ -33,14 +33,14 @@ "flimzycrag" ], "mteb_sts_scores": { - "BIOSSES": 0.500132, - "SICK-R": 0.565144, - "STS12": 0.594863, - "STS13": 0.653794, - "STS14": 0.596517, - "STS15": 0.611504, - "STS16": 0.569703, - "STSBenchmark": 0.559642 + "BIOSSES": 0.497391, + "SICK-R": 0.567003, + "STS12": 0.590715, + "STS13": 0.663817, + "STS14": 0.600884, + "STS15": 0.614552, + "STS16": 0.574441, + "STSBenchmark": 0.565802 }, "token_order_hash": "09d9aaaedc7d41b2fd1e3ffa64f16794a2cefe9ffc7faf1055ae5d49350c3060" }, @@ -76,14 +76,14 @@ "##\uff5e" ], "mteb_sts_scores": { - "BIOSSES": 0.656286, - "SICK-R": 0.611602, - "STS12": 0.620366, - "STS13": 0.735474, - "STS14": 0.670425, - "STS15": 0.745883, - "STS16": 0.67784, - "STSBenchmark": 0.655636 + "BIOSSES": 0.6499, + "SICK-R": 0.612639, + "STS12": 0.617488, + "STS13": 0.739954, + "STS14": 0.672889, + "STS15": 0.747828, + "STS16": 0.68201, + "STSBenchmark": 0.660814 }, "token_order_hash": "d881fab650dd73240615f87f531ff31a6135f961933ebc32c0dfdd130763e7a8" } diff --git a/tests/test_trainable.py b/tests/test_trainable.py index 2bc744a..2929deb 100644 --- a/tests/test_trainable.py +++ b/tests/test_trainable.py @@ -642,6 +642,20 @@ def test_get_probable_pad_token_id(mock_tokenizer: Tokenizer, caplog: pytest.Log assert "No known pad token found, using 0 as default" in caplog.text +def test_get_probable_pad_token_id_through_static_model(mock_vectors: np.ndarray, mock_tokenizer: Tokenizer) -> None: + """Test that a non-standard pad token survives StaticModel construction.""" + tokenizer_model = TokenizerModel.from_tokenizer(mock_tokenizer) + tokenizer_model.pad_token = "word1" + pad_id = tokenizer_model.pad_token_id + assert pad_id != 0 + + t = tokenizer_model.to_tokenizer() + model = StaticModel(vectors=mock_vectors, tokenizer=t) + + assert model.tokenizer.padding is not None + assert get_probable_pad_token_id(model.tokenizer) == pad_id + + def test_determine_class_weight(mock_trained_pipeline: StaticModelForClassification) -> None: """Test what the class weights are.""" w_dict = dict(zip(mock_trained_pipeline.classes, [0.5, 3]))