Skip to content

feat: make out dim a lazy assignment #399

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Merged
stephantul merged 2 commits into from
Oct 13, 2026
Merged

feat: make out dim a lazy assignment #399

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
2 commits
Select commit Hold shift + click to select a range
0abcfcd
make out dim a lazy assignment
stephantul Oct 13, 2026
8f523b3
fix comment
stephantul Oct 13, 2026
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
6 changes: 3 additions & 3 deletions model2vec/train/README.md
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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

Expand Down
32 changes: 25 additions & 7 deletions model2vec/train/base.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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,
Expand Down Expand Up @@ -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)
Expand Down Expand Up @@ -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] = []
Expand Down Expand Up @@ -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()
Expand All @@ -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,
Expand Down Expand Up @@ -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,
Expand All @@ -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,
Expand Down
14 changes: 3 additions & 11 deletions model2vec/train/classifier.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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,
Expand All @@ -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,
Expand All @@ -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,
Expand All @@ -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,
Expand All @@ -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,
Expand All @@ -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,
Expand All @@ -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,
Expand All @@ -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,
Expand All @@ -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)
Expand Down
8 changes: 0 additions & 8 deletions model2vec/train/similarity.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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,
Expand All @@ -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,
Expand All @@ -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,
Expand All @@ -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,
Expand All @@ -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,
Expand All @@ -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,
Expand All @@ -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,
Expand All @@ -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,
Expand Down
59 changes: 43 additions & 16 deletions tests/test_trainable.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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
Expand Down Expand Up @@ -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:
Expand All @@ -46,15 +55,30 @@ 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()
assert head[0].in_features == mock_vectors.shape[1]
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)
Expand All @@ -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)))

Expand Down Expand Up @@ -156,22 +182,23 @@ 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(
temp_dir,
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
Expand Down Expand Up @@ -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"]))
Expand Down Expand Up @@ -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

Expand All @@ -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)
Expand All @@ -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)
Expand Down
Loading