diff --git a/model2vec/train/README.md b/model2vec/train/README.md index 3191e90..3596cb2 100644 --- a/model2vec/train/README.md +++ b/model2vec/train/README.md @@ -174,19 +174,19 @@ The validation set is passed as `text_a_val` and `text_b_val`, instead of `X_val # Model options -All four models are created with `from_pretrained` or `from_static_model`, which take the same keyword arguments: +All four models are created with `from_pretrained` or `from_static_model`, which take these keyword arguments: | Argument | Default | Description | | --- | --- | --- | | `n_layers` | `1` | The number of hidden layers in the head. | | `hidden_dim` | `512` | The size of the hidden layers. | -| `out_dim` | | The output dimension. For classification, similarity and regression, this is set by `fit`. For pair similarity, it defaults to the embedding dimension. | +| `out_dim` | The embedding dimension | Pair similarity only. The output dimension. For classification, similarity and regression, the output dimension is set by `fit`. | | `max_length` | The static model's | The maximum number of tokens per text, for both training and inference. Pass `None` to disable truncation. | | `freeze` | `False` | Freeze the token embeddings, so that only the head and token weights are trained. | | `freeze_weights` | `None` | Controls the token weights. `None` trains the model's own token weights, and leaves a model without token weights without them. `False` also gives a model without token weights learnable weights, starting at 1. `True` freezes them. | | `normalize` | `True` | Normalize the embeddings before the head. | -With `n_layers=0`, a classifier gets a single linear layer. A similarity or pair model gets no head at all if `out_dim` equals the embedding dimension, and a single linear layer otherwise. +With `n_layers=0`, a classifier gets a single linear layer. A similarity or pair model gets no head at all if its output dimension equals the embedding dimension, and a single linear layer otherwise. # Training options diff --git a/model2vec/train/base.py b/model2vec/train/base.py index 4615545..386b850 100644 --- a/model2vec/train/base.py +++ b/model2vec/train/base.py @@ -46,7 +46,7 @@ def __init__( tokenizer: Tokenizer, hidden_dim: int = 256, n_layers: int = 0, - out_dim: int = 2, + out_dim: int | None = None, token_mapping: list[int] | None = None, weights: torch.Tensor | None = None, freeze: bool = False, @@ -99,7 +99,7 @@ def __init__( self.freeze = freeze self._weights = weights self.embeddings = self.construct_embeddings() - self.head = self.construct_head() + self._head: nn.Sequential | None = None if self.out_dim is None else self.construct_head() self.w = self.construct_weights() # Truncation happens here through `max_length`; a StaticModel's tokenizer carries its own setting. self.tokenizer = copy.deepcopy(tokenizer) @@ -146,8 +146,25 @@ def construct_weights(self) -> nn.Parameter | None: w = self._weights.clone().float() return nn.Parameter(w, requires_grad=not self.freeze_weights) + @property + def head(self) -> nn.Sequential: + """The head of the model. + + :return: The head. + :raises ValueError: If the head has not been constructed yet. + """ + if self._head is None: + raise ValueError("The model has no head yet. Call `fit`.") + return self._head + def construct_head(self) -> nn.Sequential: - """Constructs a simple head, which is empty if it has no layers and doesn't change the dimension.""" + """Constructs a simple head, which is empty if it has no layers and doesn't change the dimension. + + :return: The head. + :raises ValueError: If `out_dim` is not set. + """ + if self.out_dim is None: + raise ValueError("Cannot construct a head without `out_dim`.") if self.n_layers == 0 and self.embed_dim == self.out_dim: return nn.Sequential() modules: list[nn.Module] = [] @@ -178,7 +195,7 @@ def construct_head(self) -> nn.Sequential: def _initialize(self) -> None: """Initialize the classifier for training.""" - self.head = self.construct_head() + self._head = self.construct_head() self.embeddings = self.construct_embeddings() self.w = self.construct_weights() self.train() @@ -193,7 +210,7 @@ def from_pretrained( max_length: int | None | _UnsetType = _UNSET, n_layers: int = 0, hidden_dim: int = 256, - out_dim: int = 2, + out_dim: int | None = None, freeze: bool = False, normalize: bool = True, freeze_weights: bool | None = None, @@ -237,7 +254,7 @@ def from_static_model( max_length: int | None | _UnsetType = _UNSET, n_layers: int = 0, hidden_dim: int = 256, - out_dim: int = 2, + out_dim: int | None = None, freeze: bool = False, normalize: bool = True, freeze_weights: bool | None = None, @@ -250,7 +267,8 @@ def from_static_model( static model's `max_length` is used. Pass None to disable truncation. :param n_layers: The number of layers in the head. :param hidden_dim: The hidden dimension of the head. - :param out_dim: The output dimension of the head. + :param out_dim: The output dimension of the head. If None, the head is constructed when `out_dim` is set, + for example by `fit`. :param freeze: Whether to freeze the embeddings. :param normalize: Whether to normalize the embeddings. :param freeze_weights: Whether to freeze the token weights. If None, the model's own weights are trained, diff --git a/model2vec/train/classifier.py b/model2vec/train/classifier.py index ccc33a2..2df5e29 100644 --- a/model2vec/train/classifier.py +++ b/model2vec/train/classifier.py @@ -148,7 +148,6 @@ def __init__( tokenizer: Tokenizer, n_layers: int = 1, hidden_dim: int = 512, - out_dim: int = 2, token_mapping: list[int] | None = None, weights: torch.Tensor | None = None, freeze: bool = False, @@ -157,13 +156,12 @@ def __init__( max_length: int | None = DEFAULT_MAX_LENGTH, ) -> None: """Initialize a standard classifier model.""" - # Alias: Follows scikit-learn. Set to dummy classes - self.classes_: list[str] = [str(x) for x in range(out_dim)] + # Alias: Follows scikit-learn. Set by fit. + self.classes_: list[str] = [] # multilabel flag will be set based on the type of `y` passed to fit. self.multilabel: bool = False super().__init__( vectors=vectors, - out_dim=out_dim, tokenizer=tokenizer, token_mapping=token_mapping, weights=weights, @@ -185,7 +183,6 @@ def from_pretrained( max_length: int | None | _UnsetType = _UNSET, n_layers: int = 1, hidden_dim: int = 512, - out_dim: int = 2, freeze: bool = False, normalize: bool = True, freeze_weights: bool | None = None, @@ -200,7 +197,6 @@ def from_pretrained( static model's `max_length` is used. Pass None to disable truncation. :param n_layers: The number of hidden layers in the head. :param hidden_dim: The hidden dimension of the head. - :param out_dim: The number of classes. This is reset when calling `fit`. :param freeze: Whether to freeze the embeddings. :param normalize: Whether to normalize the embeddings. :param freeze_weights: Whether to freeze the token weights. If None, the model's own weights are trained, @@ -214,7 +210,6 @@ def from_pretrained( max_length=max_length, n_layers=n_layers, hidden_dim=hidden_dim, - out_dim=out_dim, freeze=freeze, normalize=normalize, freeze_weights=freeze_weights, @@ -229,7 +224,6 @@ def from_static_model( max_length: int | None | _UnsetType = _UNSET, n_layers: int = 1, hidden_dim: int = 512, - out_dim: int = 2, freeze: bool = False, normalize: bool = True, freeze_weights: bool | None = None, @@ -242,7 +236,6 @@ def from_static_model( static model's `max_length` is used. Pass None to disable truncation. :param n_layers: The number of hidden layers in the head. :param hidden_dim: The hidden dimension of the head. - :param out_dim: The number of classes. This is reset when calling `fit`. :param freeze: Whether to freeze the embeddings. :param normalize: Whether to normalize the embeddings. :param freeze_weights: Whether to freeze the token weights. If None, the model's own weights are trained, @@ -254,7 +247,6 @@ def from_static_model( **_static_model_arguments(model, max_length=max_length), n_layers=n_layers, hidden_dim=hidden_dim, - out_dim=out_dim, freeze=freeze, normalize=normalize, freeze_weights=freeze_weights, @@ -268,7 +260,7 @@ def classes(self) -> np.ndarray: def construct_head(self) -> nn.Sequential: """Constructs a classifier head, which always has at least one linear layer.""" - if self.n_layers == 0: + if self.n_layers == 0 and self.out_dim is not None: linear = nn.Linear(self.embed_dim, self.out_dim) nn.init.xavier_uniform_(linear.weight) nn.init.zeros_(linear.bias) diff --git a/model2vec/train/similarity.py b/model2vec/train/similarity.py index 3921294..fff1721 100644 --- a/model2vec/train/similarity.py +++ b/model2vec/train/similarity.py @@ -73,7 +73,6 @@ def __init__( tokenizer: Tokenizer, n_layers: int = 1, hidden_dim: int = 512, - out_dim: int = 2, token_mapping: list[int] | None = None, weights: torch.Tensor | None = None, freeze: bool = False, @@ -84,7 +83,6 @@ def __init__( """Initialize a standard similarity model.""" super().__init__( vectors=vectors, - out_dim=out_dim, tokenizer=tokenizer, token_mapping=token_mapping, weights=weights, @@ -106,7 +104,6 @@ def from_pretrained( max_length: int | None | _UnsetType = _UNSET, n_layers: int = 1, hidden_dim: int = 512, - out_dim: int = 2, freeze: bool = False, normalize: bool = True, freeze_weights: bool | None = None, @@ -121,7 +118,6 @@ def from_pretrained( static model's `max_length` is used. Pass None to disable truncation. :param n_layers: The number of hidden layers in the head. :param hidden_dim: The hidden dimension of the head. - :param out_dim: The output dimension of the head. This is reset when calling `fit`. :param freeze: Whether to freeze the embeddings. :param normalize: Whether to normalize the embeddings. :param freeze_weights: Whether to freeze the token weights. If None, the model's own weights are trained, @@ -135,7 +131,6 @@ def from_pretrained( max_length=max_length, n_layers=n_layers, hidden_dim=hidden_dim, - out_dim=out_dim, freeze=freeze, normalize=normalize, freeze_weights=freeze_weights, @@ -150,7 +145,6 @@ def from_static_model( max_length: int | None | _UnsetType = _UNSET, n_layers: int = 1, hidden_dim: int = 512, - out_dim: int = 2, freeze: bool = False, normalize: bool = True, freeze_weights: bool | None = None, @@ -163,7 +157,6 @@ def from_static_model( static model's `max_length` is used. Pass None to disable truncation. :param n_layers: The number of hidden layers in the head. :param hidden_dim: The hidden dimension of the head. - :param out_dim: The output dimension of the head. This is reset when calling `fit`. :param freeze: Whether to freeze the embeddings. :param normalize: Whether to normalize the embeddings. :param freeze_weights: Whether to freeze the token weights. If None, the model's own weights are trained, @@ -175,7 +168,6 @@ def from_static_model( **_static_model_arguments(model, max_length=max_length), n_layers=n_layers, hidden_dim=hidden_dim, - out_dim=out_dim, freeze=freeze, normalize=normalize, freeze_weights=freeze_weights, diff --git a/tests/test_trainable.py b/tests/test_trainable.py index 060de7c..a8a4b21 100644 --- a/tests/test_trainable.py +++ b/tests/test_trainable.py @@ -2,7 +2,7 @@ import logging from collections import Counter, UserList from tempfile import TemporaryDirectory -from typing import Any +from typing import Any, TypeVar from unittest.mock import patch import numpy as np @@ -37,6 +37,15 @@ split_indices, ) +M = TypeVar("M", bound=BaseFinetuneable) + + +def _with_out_dim(model: M, out_dim: int) -> M: + """Set the output dimension of a model and construct its head.""" + model.out_dim = out_dim + model._initialize() + return model + @pytest.mark.parametrize("n_layers", [0, 1, 2, 3]) def test_init_predict(n_layers: int, mock_vectors: np.ndarray, mock_tokenizer: Tokenizer) -> None: @@ -46,8 +55,9 @@ def test_init_predict(n_layers: int, mock_vectors: np.ndarray, mock_tokenizer: T assert s.vectors.shape == mock_vectors.shape assert s.w is None assert list(s.classes) == s.classes_ - assert list(s.classes) == ["0", "1"] + assert list(s.classes) == [] + s.out_dim = 2 head = s.construct_head() assert head[0].in_features == mock_vectors.shape[1] head = s.construct_head() @@ -55,6 +65,20 @@ def test_init_predict(n_layers: int, mock_vectors: np.ndarray, mock_tokenizer: T assert head[-1].out_features == 2 +@pytest.mark.parametrize("cls", [StaticModelForClassification, StaticModelForSimilarity, BaseFinetuneable]) +def test_init_without_out_dim(cls: type[BaseFinetuneable], mock_vectors: np.ndarray, mock_tokenizer: Tokenizer) -> None: + """A model without `out_dim` has no head until `out_dim` is set.""" + s = cls(vectors=torch.from_numpy(mock_vectors), tokenizer=mock_tokenizer) + assert s.out_dim is None + assert not any(name.startswith("_head") for name in s.state_dict()) + with pytest.raises(ValueError, match="no head yet"): + s.head + with pytest.raises(ValueError, match="no head yet"): + s.encode(["word1 word2"]) + with pytest.raises(ValueError, match="without `out_dim`"): + s.construct_head() + + def test_init_base_class(mock_vectors: np.ndarray, mock_tokenizer: Tokenizer) -> None: """Test successful initialization of the base class.""" vectors_torched = torch.from_numpy(mock_vectors) @@ -79,6 +103,8 @@ def test_empty_texts_have_finite_gradients(mock_vectors: np.ndarray, mock_tokeni model = StaticModelForClassification( vectors=torch.from_numpy(mock_vectors).float() * 1e20, tokenizer=mock_tokenizer, n_layers=0 ) + model._initialize_on_labels(["0", "1"]) + model._initialize() dataset = model._text_dataset(ColumnRows(text=["word1 word2", ""], label=["0", "1"])) batch, y = next(iter(dataset.to_dataloader(shuffle=False, batch_size=2))) @@ -156,6 +182,7 @@ def test_from_pretrained_explicit_arguments( ) -> None: """Test that from_pretrained passes each argument on to the model.""" model = StaticModel(vectors=mock_vectors, tokenizer=mock_tokenizer, weights=np.ones(len(mock_vectors))) + out_dim_kwargs = {"out_dim": 4} if "out_dim" in inspect.signature(cls.from_pretrained).parameters else {} with TemporaryDirectory() as temp_dir: model.save_pretrained(temp_dir) s = cls.from_pretrained( @@ -163,15 +190,15 @@ def test_from_pretrained_explicit_arguments( max_length=7, n_layers=2, hidden_dim=3, - out_dim=4, freeze=True, normalize=False, freeze_weights=True, + **out_dim_kwargs, ) assert s.max_length == 7 assert s.n_layers == 2 assert s.hidden_dim == 3 - assert s.out_dim == 4 + assert s.out_dim == out_dim_kwargs.get("out_dim") assert s.freeze assert not s.embeddings.weight.requires_grad assert not s.normalize @@ -460,6 +487,7 @@ def test_column_rows() -> None: def test_training_batch_matches_tokenize(mock_vectors: np.ndarray, mock_tokenizer: Tokenizer) -> None: """Training batches hold the same tokens as `tokenize`, and a text encodes the same regardless of its batch.""" s = StaticModelForClassification(vectors=torch.from_numpy(mock_vectors).float(), tokenizer=mock_tokenizer) + s._initialize_on_labels(["0", "1"]) texts = ["word2", "word2 word3"] dataset = s._text_dataset(ColumnRows(text=texts, label=["0", "1"])) @@ -709,11 +737,9 @@ def test_pair_similarity_out_dim_defaults_to_embed_dim(mock_vectors: np.ndarray, ) def test_no_head_without_layers(model_class: Any, mock_vectors: np.ndarray, mock_tokenizer: Tokenizer) -> None: """Without layers and with an unchanged dimension, the model has no head and matches its static model.""" - model = model_class( - vectors=torch.from_numpy(mock_vectors).float(), - tokenizer=mock_tokenizer, - n_layers=0, - out_dim=mock_vectors.shape[1], + model = _with_out_dim( + model_class(vectors=torch.from_numpy(mock_vectors).float(), tokenizer=mock_tokenizer, n_layers=0), + mock_vectors.shape[1], ) assert len(model.head) == 0 @@ -729,11 +755,9 @@ def test_head_without_layers_changes_dimension( model_class: Any, mock_vectors: np.ndarray, mock_tokenizer: Tokenizer ) -> None: """Without layers but with a different output dimension, the head is a single linear layer.""" - model = model_class( - vectors=torch.from_numpy(mock_vectors).float(), - tokenizer=mock_tokenizer, - n_layers=0, - out_dim=mock_vectors.shape[1] + 1, + model = _with_out_dim( + model_class(vectors=torch.from_numpy(mock_vectors).float(), tokenizer=mock_tokenizer, n_layers=0), + mock_vectors.shape[1] + 1, ) assert len(model.head) == 1 assert isinstance(model.head[0], torch.nn.Linear) @@ -754,8 +778,11 @@ def test_similarity_fit_without_layers(mock_vectors: np.ndarray, mock_tokenizer: def test_classifier_keeps_head_when_dimensions_match(mock_vectors: np.ndarray, mock_tokenizer: Tokenizer) -> None: """A classifier without layers keeps its linear layer, even if the number of classes equals the dimension.""" - model = StaticModelForClassification( - vectors=torch.from_numpy(mock_vectors).float(), tokenizer=mock_tokenizer, n_layers=0 + model = _with_out_dim( + StaticModelForClassification( + vectors=torch.from_numpy(mock_vectors).float(), tokenizer=mock_tokenizer, n_layers=0 + ), + mock_vectors.shape[1], ) assert model.out_dim == mock_vectors.shape[1] assert isinstance(model.head[0], torch.nn.Linear)