Skip to content
Merged
Show file tree
Hide file tree
Changes from all commits
Commits
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
2 changes: 1 addition & 1 deletion model2vec/train/README.md
Original file line number Diff line number Diff line change
Expand Up @@ -22,7 +22,7 @@ from model2vec.train import StaticModelForClassification

# From a distilled model
distilled_model = distill("baai/bge-base-en-v1.5")
classifier = StaticModelForClassification.from_static_model(distilled_model)
classifier = StaticModelForClassification.from_static_model(model=distilled_model)

# From a pre-trained model: potion is the default
classifier = StaticModelForClassification.from_pretrained(model_name="minishlab/potion-base-32m")
Expand Down
6 changes: 3 additions & 3 deletions model2vec/train/base.py
Original file line number Diff line number Diff line change
Expand Up @@ -45,14 +45,14 @@ def construct_head(self) -> nn.Sequential:

@classmethod
def from_pretrained(
cls: type[ModelType], out_dim: int = 2, model_name: str = "minishlab/potion-base-32m", **kwargs: Any
cls: type[ModelType], *, out_dim: int = 2, model_name: str = "minishlab/potion-base-32m", **kwargs: Any
) -> ModelType:
"""Load the model from a pretrained model2vec model."""
model = StaticModel.from_pretrained(model_name)
return cls.from_static_model(model, out_dim, **kwargs)
return cls.from_static_model(model=model, out_dim=out_dim, **kwargs)

@classmethod
def from_static_model(cls: type[ModelType], model: StaticModel, out_dim: int = 2, **kwargs: Any) -> ModelType:
def from_static_model(cls: type[ModelType], *, model: StaticModel, out_dim: int = 2, **kwargs: Any) -> ModelType:
"""Load the model from a static model."""
model.embedding = np.nan_to_num(model.embedding)
embeddings_converted = torch.from_numpy(model.embedding)
Expand Down
4 changes: 2 additions & 2 deletions tests/test_trainable.py
Original file line number Diff line number Diff line change
Expand Up @@ -41,7 +41,7 @@ def test_init_base_class(mock_vectors: np.ndarray, mock_tokenizer: Tokenizer) ->
def test_init_base_from_model(mock_vectors: np.ndarray, mock_tokenizer: Tokenizer) -> None:
"""Test initializion from a static model."""
model = StaticModel(vectors=mock_vectors, tokenizer=mock_tokenizer)
s = FinetunableStaticModel.from_static_model(model)
s = FinetunableStaticModel.from_static_model(model=model)
assert s.vectors.shape == mock_vectors.shape
assert s.w.shape[0] == mock_vectors.shape[0]

Expand All @@ -55,7 +55,7 @@ def test_init_base_from_model(mock_vectors: np.ndarray, mock_tokenizer: Tokenize
def test_init_classifier_from_model(mock_vectors: np.ndarray, mock_tokenizer: Tokenizer) -> None:
"""Test initializion from a static model."""
model = StaticModel(vectors=mock_vectors, tokenizer=mock_tokenizer)
s = StaticModelForClassification.from_static_model(model)
s = StaticModelForClassification.from_static_model(model=model)
assert s.vectors.shape == mock_vectors.shape
assert s.w.shape[0] == mock_vectors.shape[0]

Expand Down