From 843f5f0beab1cc2bbb94ca53639206524e788244 Mon Sep 17 00:00:00 2001 From: stephantul Date: Fri, 9 Oct 2026 08:56:16 +0200 Subject: [PATCH 1/7] first pass docs --- README.md | 7 +- model2vec/train/README.md | 174 +++++++++++++++++++++++++++++++------- 2 files changed, 146 insertions(+), 35 deletions(-) diff --git a/README.md b/README.md index 2248be1..1a7b169 100644 --- a/README.md +++ b/README.md @@ -74,7 +74,7 @@ m2v_model = distill(model_name="BAAI/bge-base-en-v1.5") m2v_model.save_pretrained("m2v_model") ``` -For advanced usage, see our [distillation docs](https://minish.ai/packages/model2vec/distillation), which includes some [distillation best practices](https://minish.ai/packages/model2vec/distillation#distillation-best-practices). After distillation, you can also fine-tune your own classification models on top of the distilled model, or on a pre-trained model. First, make sure you install the `training` extras with: +For advanced usage, see our [distillation docs](https://minish.ai/packages/model2vec/distillation), which includes some [distillation best practices](https://minish.ai/packages/model2vec/distillation#distillation-best-practices). After distillation, you can also fine-tune the distilled model, or a pre-trained model, for classification, similarity, regression, or pair similarity (e.g. query-document retrieval). First, make sure you install the `training` extras with: ```bash pip install model2vec[train] @@ -83,12 +83,11 @@ pip install model2vec[train] Then, you can fine-tune a model as follows: ```python -import numpy as np from datasets import load_dataset from model2vec.train import StaticModelForClassification # Initialize a classifier from a pre-trained model -classifier = StaticModelForClassification.from_pretrained(model_name="minishlab/potion-base-32M") +classifier = StaticModelForClassification.from_pretrained(path="minishlab/potion-base-32M") # Load a dataset. Note: both single and multi-label classification datasets are supported ds = load_dataset("setfit/subj") @@ -110,7 +109,7 @@ For advanced usage, see our [training docs](https://minish.ai/packages/model2vec - **Lightweight Dependencies**: the base package's only major dependency is `numpy`. - **Lightning-fast Inference**: up to 500 times faster on CPU than the original model. - **Fast, Dataset-free Distillation**: distill your own model in 30 seconds on a CPU, without a dataset. -- **Fine-tuning**: fine-tune your own classification models on top of Model2Vec models. +- **Fine-tuning**: fine-tune Model2Vec models for classification, similarity, regression, or pair similarity, and run them without `torch`. - **Integrated in many popular libraries**: Model2Vec is integrated direclty into popular libraries such as [Sentence Transformers](https://github.com/UKPLab/sentence-transformers) and [LangChain](https://github.com/langchain-ai/langchain). For more information, see our [integrations documentation](https://minish.ai/packages/model2vec/integrations). - **Tightly integrated with HuggingFace hub**: easily share and load models from the HuggingFace hub, using the familiar `from_pretrained` and `push_to_hub`. Our own models can be found [here](https://huggingface.co/minishlab). diff --git a/model2vec/train/README.md b/model2vec/train/README.md index 0dea77b..736481e 100644 --- a/model2vec/train/README.md +++ b/model2vec/train/README.md @@ -1,8 +1,15 @@ # Training -Aside from [distillation](../../README.md#distillation), `model2vec` also supports training simple classifiers on top of static models, using [pytorch](https://pytorch.org/). +Aside from [distillation](../../README.md#distillation), `model2vec` also supports fine-tuning static models with [pytorch](https://pytorch.org/). Every trainable model consists of the static embeddings of a `StaticModel`, optionally followed by a small MLP head. Both the embeddings and the head are trained. -We support both single and multi-label classification, which work seamlessly based on the labels you provide. +There are four trainable models, one per task: + +| Class | Task | Targets | Loss | +| --- | --- | --- | --- | +| `StaticModelForClassification` | Single- or multi-label classification | A label, or a list of labels, per text | (Binary) cross-entropy, or focal loss | +| `StaticModelForSimilarity` | Mapping texts to target vectors | A vector per text | Cosine distance | +| `StaticModelForRegression` | Regressing vectors of numbers | A vector per text | Mean squared error | +| `StaticModelForPairSimilarity` | Embedding related texts close together | A paired text per text | InfoNCE with in-batch negatives | # Installation @@ -12,28 +19,29 @@ To train, make sure you install the training extra: pip install model2vec[train] ``` +This installs `torch` and `datasets`. + # Quickstart -To train a model, simply initialize it using a `StaticModel`, or from a pre-trained model, as follows: +To train a model, initialize it from a pre-trained model, or from a `StaticModel`, for example one you distilled yourself: ```python from model2vec.distill import distill from model2vec.train import StaticModelForClassification +# From a pre-trained model: potion-base-32m is the default +classifier = StaticModelForClassification.from_pretrained(path="minishlab/potion-base-32M") + # From a distilled model distilled_model = distill("baai/bge-base-en-v1.5") classifier = StaticModelForClassification.from_static_model(model=distilled_model) - -# From a pre-trained model: potion is the default -classifier = StaticModelForClassification.from_pretrained(path="minishlab/potion-base-32m") ``` -This creates a very simple classifier: a StaticModel with a single 512-unit hidden layer on top. You can adjust the number of hidden layers and the number units through some parameters on both functions. Note that the default for `from_pretrained` is [potion-base-32m](https://huggingface.co/minishlab/potion-base-32M), our best model to date. This is our recommended path if you're working with general English data. +This creates a classifier with a single 512-unit hidden layer on top of the static embeddings. See [Model options](#model-options) for how to change this. The default for `from_pretrained` is [potion-base-32m](https://huggingface.co/minishlab/potion-base-32M), our best model to date. This is our recommended path if you're working with general English data. -Now that you have created the classifier, let's just train a model. The example below assumes you have the [`datasets`](https://github.com/huggingface/datasets) library installed. +Now let's train it. The example below uses the [`datasets`](https://github.com/huggingface/datasets) library, which is installed with the training extra. ```python -import numpy as np from datasets import load_dataset from time import perf_counter @@ -47,26 +55,26 @@ classifier = classifier.fit(train["text"], train["label"]) print(f"Training took {int(perf_counter() - s)} seconds.") # Training took 31 seconds -classification_report = classifier.evaluate(ds["test"]["text"], ds["test"]["label"]) +classification_report = classifier.evaluate(test["text"], test["label"]) print(classification_report) # Achieved 92.0 test accuracy ``` As you can see, we got a pretty nice 92% accuracy, with only 31 seconds of training. -The training loop is a plain PyTorch loop (see [`model2vec/train/trainer.py`](trainer.py)). By default the training loop splits the data into a train and validation split, with 90% of the data being used for training and 10% for validation. By default, it runs with early stopping on the validation set accuracy, with a patience of 5. +The classes are taken from the labels you pass to `fit`, so labels can be strings or integers, and `out_dim` does not need to be set beforehand. After training, `classifier.classes` holds the classes in the order of the model's outputs. -Note that this model is as fast as you're used to from us: +Inference is as fast as you're used to from us: ```python -from time import perf_counter - s = perf_counter() classifier.predict(test["text"]) print(f"Took {int((perf_counter() - s) * 1000)} milliseconds for {len(test)} instances on CPU.") # Took 66 milliseconds for 2000 instances on CPU. ``` +`predict_proba` returns the probability of each class instead. + ## Multi-label classification Multi-label classification is supported out of the box. Just pass a list of lists to the `fit` function (e.g. `[[label1, label2], [label1, label3]]`), and a multi-label classifier will be trained. For example, the following code trains a multi-label classifier on the [go_emotions](https://huggingface.co/datasets/google-research-datasets/go_emotions) dataset: @@ -97,8 +105,50 @@ print(classification_report) # {'accuracy': 0.41, 'macro avg': {'precision': 0.527, 'recall': 0.41, 'f1-score': 0.439, ...}, ...} ``` +The `threshold` is the minimum probability for a label to be predicted, and can also be passed to `predict`. During training, the validation accuracy of a multi-label classifier is the Jaccard similarity between the predicted and true labels, at a threshold of 0.5. + The scores are competitive with the popular [roberta-base-go_emotions](https://huggingface.co/SamLowe/roberta-base-go_emotions) model, while our model is orders of magnitude faster. +## Imbalanced data + +If some classes are much rarer than others, there are two options, which can be combined. + +`class_weight` weights the loss of each class. Pass `"balanced"` to weight each class by its inverse frequency, or a dict that maps each class to its weight: + +```python +classifier.fit(X, y, class_weight="balanced") +classifier.fit(X, y, class_weight={"positive": 3.0, "negative": 1.0}) +``` + +For multi-label classification, the weight of a class only applies to the texts that have that class. + +`focal_gamma` switches the loss to a [focal loss](https://arxiv.org/abs/1708.02002), which down-weights the texts the model already classifies correctly with high confidence. The default, `0.0`, is equal to plain cross-entropy. A common value is `2.0`: + +```python +classifier.fit(X, y, focal_gamma=2.0) +``` + +## Similarity and regression + +`StaticModelForSimilarity` and `StaticModelForRegression` learn to map each text to a target vector. Pass the vectors as `y`, as a 2D array or tensor, a list of lists, or a dataset column that holds lists of numbers. The output dimension of the model is set to the dimension of the vectors. + +`StaticModelForSimilarity` is trained with a cosine distance loss, so only the direction of the output matters. A typical use is to train a static model to mimic the embeddings of a larger model: + +```python +from sentence_transformers import SentenceTransformer +from model2vec.train import StaticModelForSimilarity + +teacher = SentenceTransformer("baai/bge-base-en-v1.5") +targets = teacher.encode(texts) + +model = StaticModelForSimilarity.from_pretrained(path="minishlab/potion-base-32M") +model.fit(texts, targets) +``` + +`StaticModelForRegression` is trained with a mean squared error loss, so the output has to match the targets exactly. To regress a single number, pass a vector of length one for each text, e.g. `[[0.3], [1.2], ...]`. + +Both models stop early on the validation loss. + ## Pair similarity `StaticModelForPairSimilarity` trains a model to embed pairs of related texts (e.g. queries and their matching documents) close together, by encoding both sides with the same model. It is trained with an InfoNCE loss with in-batch negatives: each `text_a` is pulled towards its paired `text_b` and pushed away from every other `text_b` in the batch: @@ -106,14 +156,69 @@ The scores are competitive with the popular [roberta-base-go_emotions](https://h ```python from model2vec.train import StaticModelForPairSimilarity -model = StaticModelForPairSimilarity.from_pretrained(path="minishlab/potion-base-32M") +model = StaticModelForPairSimilarity.from_pretrained(path="minishlab/potion-base-32M", n_layers=0) model.fit(text_a=queries, text_b=documents) + +static_model = model.to_static_model() ``` -Because the other pairs in a batch serve as negatives, the training and validation sets each need at least two pairs. Pairs with the same `text_a` are treated as alternative positives for that text, so they don't serve as negatives for each other. +With `n_layers=0`, the model has no head, so the trained embeddings can be turned back into a regular `StaticModel`. With the default `n_layers=1`, the model has an MLP head, and you need [`to_pipeline`](#using-a-trained-model) to keep it. + +Because the other pairs in a batch serve as negatives, the training and validation sets each need at least two pairs, and the batch size must be at least two. Larger batches give more negatives per pair. Pairs with the same `text_a`, or the same `text_b`, are not used as negatives for each other. The InfoNCE temperature can be set with `temperature` (default `0.05`). It must be positive. +The validation set is passed as `text_a_val` and `text_b_val`, instead of `X_val` and `y_val`. + +# Model options + +All four models are created with `from_pretrained` or `from_static_model`, which take the same 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. | +| `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. + +# Training options + +`fit` takes the following options, which are shared by all models, unless noted otherwise: + +| Argument | Default | Description | +| --- | --- | --- | +| `learning_rate` | `1e-3` | The learning rate of the Adam optimizer. | +| `batch_size` | `None` | The batch size. If `None`, a multiple of 32 between 32 and 512 is chosen based on the size of the training set. | +| `min_epochs` | `None` | The minimum number of epochs before early stopping can stop training. | +| `max_epochs` | `-1` | The maximum number of epochs. If `-1`, training continues until early stopping triggers. | +| `early_stopping_patience` | `5` | The number of validation checks without improvement before training stops. If `None`, early stopping is disabled. | +| `test_size` | `0.1` | The size of the validation split, if no validation set is passed. See [Validation](#validation). | +| `X_val`, `y_val` | `None` | An explicit validation set. `text_a_val` and `text_b_val` for pair similarity. | +| `validation_steps` | `None` | Validate every this many training steps. See [Validation](#validation). | +| `device` | `"auto"` | The device to train on. `"auto"` picks CUDA, then MPS, then CPU. | +| `random_seed` | `42` | The seed for initialization and shuffling. The validation split always uses a seed of 42. | +| `token_dropout` | `0.0` | The fraction of tokens to randomly drop from each training text. This is a form of data augmentation, and has no effect during validation. Must be in `[0, 1)`. | +| `class_weight` | `None` | Classification only. See [Imbalanced data](#imbalanced-data). | +| `focal_gamma` | `0.0` | Classification only. See [Imbalanced data](#imbalanced-data). | +| `temperature` | `0.05` | Pair similarity only. The temperature of the InfoNCE loss. | + +Calling `fit` re-initializes the head, the embeddings and the token weights from the static model, so calling it twice does not continue training. + +## Validation + +Without an explicit validation set, `fit` holds out `test_size` of the data for validation, capped at 10,000 rows; pass an int to hold out an exact number of rows. For single-label classification, the split is stratified by class. + +By default, the model is validated once per epoch. If an epoch has more than 250 batches, it is instead validated about four times per epoch, but at most once every 250 steps. `validation_steps` sets a fixed number of steps between validations instead. + +Classifiers stop early on the validation accuracy, the other models on the validation loss. Whenever the validation loss stops improving for a few epochs, the learning rate is halved. After training, the weights from the best validation check are loaded back into the model. + +The training loop shows the loss, and for classifiers the accuracy, of the latest training batches and of the latest validation check in the progress bar. + ## Large datasets Training data is read and tokenized per batch, so `fit` also accepts the columns of a Hugging Face dataset. These are not loaded into memory: @@ -125,25 +230,23 @@ dataset = load_dataset("sentence-transformers/gooaq", split="train") model.fit(text_a=dataset["question"], text_b=dataset["answer"]) ``` -Without an explicit validation set, `fit` holds out `test_size` of the data for validation, capped at 10,000 rows; pass an int to hold out an exact number of rows. For single-label classification, the split is stratified by class. - Because batches are shuffled, training reads the rows of a dataset in random order. If the dataset is stored on disk and does not fit in memory, this can be slow, especially on a network file system. Columns of a dataset with a transform, set with `with_transform`, are not accepted. Apply the transform first with `dataset.map(transform, batched=True)`. For a dataset loaded from disk or the Hub, this writes the result to the cache on disk, so it is still not loaded into memory. -Columns of an iterable dataset, such as one loaded with `streaming=True`, are not accepted. +Columns of an iterable dataset, such as one loaded with `streaming=True`, are not accepted. Neither is a whole `Dataset`: pass its columns, such as `dataset["text"]`. -# Persistence +# Using a trained model -You can turn a classifier into a lightweight inference pipeline, as follows: +You can turn any trained model into a lightweight inference pipeline: ```python pipeline = classifier.to_pipeline() ``` -This strips away `torch`: the head's weights are plain `numpy` arrays, so the resulting `StaticModelPipeline` can be used for inference without installing `torch`. +This strips away `torch`: the head's weights are plain `numpy` arrays, so the resulting `StaticModelPipeline` can be used for inference without installing `torch`. For a classifier, `pipeline.predict` returns labels, and `pipeline.predict_proba` and `pipeline.evaluate` work as on the classifier. For the other models, `pipeline.predict` returns the output vectors. -If you want to persist your pipeline to the Hugging Face hub, you can use our built-in functions: +If you want to persist your pipeline locally, or to the Hugging Face Hub, you can use our built-in functions: ```python pipeline.save_pretrained(path) @@ -160,19 +263,28 @@ pipeline = StaticModelPipeline.from_pretrained("my_cool/project") Loading pipelines in this way is _extremely_ fast. It takes only 30ms to load a pipeline from disk. +A pipeline can also be exported to ONNX, including its head: -# Bring your own architecture +```python +from model2vec.onnx import export_model_to_onnx + +export_model_to_onnx(pipeline, "my_onnx_model") +``` + +`to_static_model` turns a trained model into a regular `StaticModel`, without its head, but with the trained embeddings and token weights. This is mainly useful for models without a head, such as a pair similarity model with `n_layers=0`. -Our training architecture is set up to be extensible, with each task having a specific class. Right now, we only offer `StaticModelForClassification`, but in the future we'll also offer regression, etc. +# Bring your own architecture -The core functionality of the `StaticModelForClassification` is contained in a couple of functions: +Our training architecture is set up to be extensible, with each task having a specific class. All of them subclass `BaseFinetuneable` (in [`base.py`](base.py)), which contains the shared functionality: -* `construct_head`: This function constructs the classifier on top of the staticmodel. For example, if you want to create a model that has LayerNorm, just subclass, and replace this function. This should be the main function to update if you want to change model behavior. -* `_create_datasets`: splits off the validation data, and creates the `torch.Dataset`s that will be used in the `Dataloader` during training. -* `_encode`: The encoding function used in the model. -* `fit`: contains all the fitting logic. +* `construct_head`: constructs the head on top of the static model. For example, if you want to create a model that has LayerNorm, just subclass, and replace this function. This should be the main function to update if you want to change model behavior. +* `construct_embeddings` and `construct_weights`: construct the token embeddings and the token weights. +* `_encode`: the encoding function used in the model: a weighted mean over the token embeddings, followed by normalization. +* `_create_datasets`: splits off the validation data, and creates the datasets that are tokenized per batch during training. +* `_to_targets`: turns a batch of labels into the targets of the loss. +* `fit`: contains all the fitting logic, and is implemented by each task. -The training loop itself is defined in `model2vec.train.trainer.run_training_loop`, a plain torch loop that is fairly basic and easy to modify. Each task passes in its own loss function (and, for classification, a small function that computes extra validation metrics like accuracy). +The training loop itself is defined in `model2vec.train.trainer.run_training_loop`, a plain torch loop that is fairly basic and easy to modify. Each task passes in its own loss function, and, for classification, a small function that computes extra metrics like accuracy. `StaticModelForRegression`, for example, only replaces the loss function of `StaticModelForSimilarity`. # Results From 485f8232913cb92be24562287d1f1db37f781c52 Mon Sep 17 00:00:00 2001 From: stephantul Date: Fri, 9 Oct 2026 10:24:51 +0200 Subject: [PATCH 2/7] updeet --- model2vec/train/README.md | 8 ++++---- 1 file changed, 4 insertions(+), 4 deletions(-) diff --git a/model2vec/train/README.md b/model2vec/train/README.md index 736481e..d7aa835 100644 --- a/model2vec/train/README.md +++ b/model2vec/train/README.md @@ -156,13 +156,13 @@ Both models stop early on the validation loss. ```python from model2vec.train import StaticModelForPairSimilarity -model = StaticModelForPairSimilarity.from_pretrained(path="minishlab/potion-base-32M", n_layers=0) +model = StaticModelForPairSimilarity.from_pretrained(path="minishlab/potion-base-32M") model.fit(text_a=queries, text_b=documents) static_model = model.to_static_model() ``` -With `n_layers=0`, the model has no head, so the trained embeddings can be turned back into a regular `StaticModel`. With the default `n_layers=1`, the model has an MLP head, and you need [`to_pipeline`](#using-a-trained-model) to keep it. +`to_static_model` turns the trained model into a regular `StaticModel`, with the trained embeddings and token weights. The head is only used during training, and is discarded. If you want to keep the head, use [`to_pipeline`](#using-a-trained-model) instead. Because the other pairs in a batch serve as negatives, the training and validation sets each need at least two pairs, and the batch size must be at least two. Larger batches give more negatives per pair. Pairs with the same `text_a`, or the same `text_b`, are not used as negatives for each other. @@ -201,7 +201,7 @@ With `n_layers=0`, a classifier gets a single linear layer. A similarity or pair | `X_val`, `y_val` | `None` | An explicit validation set. `text_a_val` and `text_b_val` for pair similarity. | | `validation_steps` | `None` | Validate every this many training steps. See [Validation](#validation). | | `device` | `"auto"` | The device to train on. `"auto"` picks CUDA, then MPS, then CPU. | -| `random_seed` | `42` | The seed for initialization and shuffling. The validation split always uses a seed of 42. | +| `random_seed` | `42` | The seed for the validation split, initialization and shuffling. | | `token_dropout` | `0.0` | The fraction of tokens to randomly drop from each training text. This is a form of data augmentation, and has no effect during validation. Must be in `[0, 1)`. | | `class_weight` | `None` | Classification only. See [Imbalanced data](#imbalanced-data). | | `focal_gamma` | `0.0` | Classification only. See [Imbalanced data](#imbalanced-data). | @@ -271,7 +271,7 @@ from model2vec.onnx import export_model_to_onnx export_model_to_onnx(pipeline, "my_onnx_model") ``` -`to_static_model` turns a trained model into a regular `StaticModel`, without its head, but with the trained embeddings and token weights. This is mainly useful for models without a head, such as a pair similarity model with `n_layers=0`. +`to_static_model` turns a trained model into a regular `StaticModel`, without its head, but with the trained embeddings and token weights. This is how you get a regular embedding model out of a pair similarity model. # Bring your own architecture From 638e3bfb4991e6415c2018cffe018eb898e1e8d4 Mon Sep 17 00:00:00 2001 From: stephantul Date: Fri, 9 Oct 2026 10:50:29 +0200 Subject: [PATCH 3/7] more docs, small fix --- model2vec/inference/README.md | 89 ++++++++++++++++++++++++---- model2vec/inference/model.py | 2 +- model2vec/persistence/persistence.py | 9 ++- model2vec/train/README.md | 2 +- tests/test_inference.py | 4 +- 5 files changed, 89 insertions(+), 17 deletions(-) diff --git a/model2vec/inference/README.md b/model2vec/inference/README.md index 74f3c5b..600cf7b 100644 --- a/model2vec/inference/README.md +++ b/model2vec/inference/README.md @@ -1,31 +1,100 @@ # Inference -This subpackage mainly contains helper functions for inference with trained classifier/projector heads, persisted as a `safetensors` file and `config.json` metadata. +This subpackage contains `StaticModelPipeline`, which runs a trained model: a `StaticModel` followed by a small MLP head. The head is implemented in `numpy`, so a pipeline only needs the base `model2vec` package: no `torch` or `scikit-learn`. -If you're looking for information on how to train a model, see [here](../train/README.md). +A pipeline is created from any of the trainable models in [`model2vec.train`](../train/README.md). For a classifier, the pipeline predicts labels. For the other models, it predicts vectors. -# Usage +# Creating a pipeline + +To create a pipeline, train a model and call `to_pipeline`: + +```python +from datasets import load_dataset +from model2vec.train import StaticModelForClassification + +ds = load_dataset("setfit/subj") +classifier = StaticModelForClassification.from_pretrained(path="minishlab/potion-base-32M") +classifier.fit(ds["train"]["text"], ds["train"]["label"]) -Let's assume you're using our [potion-edu classifier](https://huggingface.co/minishlab/potion-8m-edu-classifier). +pipeline = classifier.to_pipeline() +``` + +# Saving and loading + +Pipelines can be saved to a folder, or pushed to the Hugging Face Hub: + +```python +pipeline.save_pretrained("my_pipeline") +pipeline.push_to_hub("my_cool/project") +``` + +`push_to_hub` also takes a `token`, a `subfolder` to push to, and `private=True` to create a private repository. + +Later, you can load these as follows: ```python from model2vec.inference import StaticModelPipeline -classifier = StaticModelPipeline.from_pretrained("minishlab/potion-8m-edu-classifier") -label = classifier.predict("Attitudes towards cattle in the Alps: a study in letting go.") +pipeline = StaticModelPipeline.from_pretrained("my_cool/project") +``` + +To load a private pipeline, pass a `token`. Loading a pipeline is _extremely_ fast: it takes only 30ms to load one from disk. + +A saved pipeline is a regular model2vec model, with two additions: the weights of the head are stored in `head.safetensors`, and the head's activation and classes are stored under `head_config` in `config.json`. This means the folder can also be loaded as a plain `StaticModel`, which gives you the embeddings without the head. + +# Usage + +`predict` takes a list of texts, or a single text: + +```python +pipeline.predict(["This movie was great.", "The plot follows a detective in Paris."]) +# array([0, 1]) +``` + +What it returns depends on the head: + +* **Single-label classification**: an array with one label per text. The labels have the same type as the labels the model was trained on. +* **Multi-label classification**: an object array with an array of labels per text. A label is predicted if its probability is higher than `threshold`, which defaults to `0.5`. +* **Similarity, regression and pair similarity**: an array with one output vector per text. + +For classifiers, `predict_proba` returns the probability of each class, in the order of `pipeline.classes_`. These are softmax probabilities for single-label classification, and sigmoid probabilities for multi-label classification. + +`evaluate` predicts labels for a set of texts, and compares them with the true labels: + +```python +report = pipeline.evaluate(ds["test"]["text"], ds["test"]["label"]) +# {'0': {'precision': ..., 'recall': ..., 'f1-score': ..., 'support': ...}, ..., 'accuracy': ..., 'macro avg': {...}, 'weighted avg': {...}} ``` -This should just work. +The report has the same format as the dictionary returned by scikit-learn's `classification_report`, but does not need scikit-learn. For multi-label classification, accuracy is the fraction of texts for which all labels are predicted correctly, and `threshold` can be passed as with `predict`. The same function is available as `evaluate_single_or_multi_label`, which takes predictions and true labels. + +`predict_proba` and `evaluate` raise an error for pipelines that predict vectors. + +`predict` and `predict_proba` take the same encoding options as `StaticModel.encode`: `batch_size`, `max_length`, `show_progress_bar`, `use_multiprocessing` and `multiprocessing_threshold`. If `max_length` is not passed, the model's own `max_length` is used. Pass `max_length=None` to disable truncation. + +The underlying `StaticModel` is available as `pipeline.model`, and the head as `pipeline.head`. + +# ONNX export + +A pipeline can be exported to ONNX, including its head: + +```python +from model2vec.onnx import export_model_to_onnx + +export_model_to_onnx(pipeline, "my_onnx_pipeline") +``` # Migrating a legacy pipeline -Pipelines saved by older versions of model2vec store the head as a `scikit-learn`/`skops` `pipeline.skops` file instead of `head.safetensors`. `from_pretrained` still loads these automatically, falling back to the legacy format and emitting a warning. This requires `scikit-learn` and `skops` to be installed. +Pipelines saved by older versions of model2vec store the head as a `scikit-learn`/`skops` `pipeline.skops` file instead of `head.safetensors`. An example is our [potion-edu classifier](https://huggingface.co/minishlab/potion-8m-edu-classifier). `from_pretrained` still loads these automatically, falling back to the legacy format and emitting a warning. This requires `scikit-learn` and `skops` to be installed. To upgrade a pipeline to the current format (and silence the warning), convert it with `convert_legacy_pipeline` and save the result: ```python from model2vec.inference import convert_legacy_pipeline -pipeline = convert_legacy_pipeline("path/or/repo-id/of/legacy/pipeline") -pipeline.save_pretrained("path/to/save") +pipeline = convert_legacy_pipeline("minishlab/potion-8m-edu-classifier") +pipeline.save_pretrained("potion-8m-edu-classifier") ``` + +Only heads that are a scikit-learn `MLPClassifier` or `MLPRegressor` with ReLU activations can be converted. By default, the `skops` file may only contain `scikit-learn` types. Pass `trust_remote_code=True` to load other types, but only for files you trust. diff --git a/model2vec/inference/model.py b/model2vec/inference/model.py index 3e3f17b..eed28f7 100644 --- a/model2vec/inference/model.py +++ b/model2vec/inference/model.py @@ -159,7 +159,7 @@ def predict_proba( :raises ValueError: If the classifier type is projector. """ if self.head.activation == Activation.IDENTITY: - raise ValueError("You are using evaluate on a projector model. This is not supported.") + raise ValueError("You are using predict_proba on a projector model. This is not supported.") encoded = self._encode_and_coerce_to_2d( X, show_progress_bar=show_progress_bar, diff --git a/model2vec/persistence/persistence.py b/model2vec/persistence/persistence.py index 20248be..1a8e4de 100644 --- a/model2vec/persistence/persistence.py +++ b/model2vec/persistence/persistence.py @@ -64,14 +64,16 @@ def save_pretrained( cfg["vocabulary_quantization"] = int(embeddings.shape[0]) else: cfg.pop("vocabulary_quantization", None) - json.dump(cfg, open(folder_path / "config.json", "w"), indent=4) + with open(folder_path / "config.json", "w") as config_file: + json.dump(cfg, config_file, indent=4) # Create modules.json modules = [{"idx": 0, "name": "0", "path": ".", "type": "sentence_transformers.models.StaticEmbedding"}] if cfg.get("normalize"): # If normalize=True, add sentence_transformers.models.Normalize modules.append({"idx": 1, "name": "1", "path": "1_Normalize", "type": "sentence_transformers.models.Normalize"}) - json.dump(modules, open(folder_path / "modules.json", "w"), indent=4) + with open(folder_path / "modules.json", "w") as modules_file: + json.dump(modules, modules_file, indent=4) logger.info(f"Saved model to {folder_path}") @@ -131,7 +133,8 @@ def load_pretrained( metadata = {} tokenizer: Tokenizer = Tokenizer.from_file(str(selected_layout.tokenizer)) - config = json.load(open(selected_layout.config)) + with open(selected_layout.config) as config_file: + config = json.load(config_file) return embeddings, tokenizer, config, metadata, weights, mapping diff --git a/model2vec/train/README.md b/model2vec/train/README.md index d7aa835..c3806ce 100644 --- a/model2vec/train/README.md +++ b/model2vec/train/README.md @@ -244,7 +244,7 @@ You can turn any trained model into a lightweight inference pipeline: pipeline = classifier.to_pipeline() ``` -This strips away `torch`: the head's weights are plain `numpy` arrays, so the resulting `StaticModelPipeline` can be used for inference without installing `torch`. For a classifier, `pipeline.predict` returns labels, and `pipeline.predict_proba` and `pipeline.evaluate` work as on the classifier. For the other models, `pipeline.predict` returns the output vectors. +This strips away `torch`: the head's weights are plain `numpy` arrays, so the resulting `StaticModelPipeline` can be used for inference without installing `torch`. For a classifier, `pipeline.predict` returns labels, and `pipeline.predict_proba` and `pipeline.evaluate` work as on the classifier. For the other models, `pipeline.predict` returns the output vectors. See the [inference docs](../inference/README.md) for more. If you want to persist your pipeline locally, or to the Hugging Face Hub, you can use our built-in functions: diff --git a/tests/test_inference.py b/tests/test_inference.py index 48c4d07..3fc8fb5 100644 --- a/tests/test_inference.py +++ b/tests/test_inference.py @@ -60,9 +60,9 @@ def test_init_predict_projector(mock_inference_pipeline_projector: StaticModelPi """Test successful init and predict with StaticModelPipeline.""" assert mock_inference_pipeline_projector.head.activation == Activation.IDENTITY assert mock_inference_pipeline_projector.classes_ is None - with pytest.raises(ValueError): + with pytest.raises(ValueError, match="predict_proba"): mock_inference_pipeline_projector.predict_proba(["dog"]) - with pytest.raises(ValueError): + with pytest.raises(ValueError, match="evaluate"): mock_inference_pipeline_projector.evaluate(["dog"], ["a"]) prediction = mock_inference_pipeline_projector.predict(["dog"]) From 41a275292b93e21971022c14619626dd1822aa82 Mon Sep 17 00:00:00 2001 From: stephantul Date: Fri, 9 Oct 2026 14:36:57 +0200 Subject: [PATCH 4/7] update results --- results/README.md | 100 ++++++++++++++++++++++++---------------------- 1 file changed, 52 insertions(+), 48 deletions(-) diff --git a/results/README.md b/results/README.md index 98f9845..bb253ab 100644 --- a/results/README.md +++ b/results/README.md @@ -96,56 +96,60 @@ As can be seen, [potion-retrieval-32M](https://huggingface.co/minishlab/potion-r The main results for Model2Vec training are outlined in this section. -We compare five different architectures for our main results: -- `model2vec + logreg`: A model2vec model with a scikit-learn `LogisticRegressionCV` on top. -- `model2vec full finetune`: A model2vec classifier with the full model finetuned. This uses our `StaticModelForClassification`. +We compare six different architectures for our main results: - `tfidf`: A TF-IDF model with a scikit-learn `LogisticRegressionCV` on top. -- `setfit`: A [SetFit](https://github.com/huggingface/setfit/tree/main) model trained using [all-minilm-l6-v2](sentence-transformers/all-MiniLM-L6-v2) as a base model. -- `bge-base + logreg`: A [BGE-base](https://huggingface.co/BAAI/bge-base-en-v1.5) encoder model with a scikit-learn `LogisticRegressionCV` on top. +- `fasttext`: A supervised [fastText](https://fasttext.cc/) classifier. +- `model2vec + logreg`: A model2vec model ([potion-base-32M](https://huggingface.co/minishlab/potion-base-32M)) with a scikit-learn `LogisticRegressionCV` on top. +- `model2vec full finetune`: A model2vec classifier with the full model finetuned, starting from [potion-base-32M](https://huggingface.co/minishlab/potion-base-32M). This uses our `StaticModelForClassification`. +- `setfit`: A [SetFit](https://github.com/huggingface/setfit/tree/main) model trained using [all-MiniLM-L6-v2](https://huggingface.co/sentence-transformers/all-MiniLM-L6-v2) as a base model. +- `minilm full finetune`: A [all-MiniLM-L6-v2](https://huggingface.co/sentence-transformers/all-MiniLM-L6-v2) model, fully finetuned as a sequence classifier using the Hugging Face `Trainer`. + +We use 26 classification datasets, using 1000 examples from the train set, and the full test set. The scores are weighted F1 scores on the test set. No parameters were tuned on any validation set. All datasets except `banking77` ([mteb/banking77](https://huggingface.co/datasets/mteb/banking77)) and `clinc_oos` ([clinc/clinc_oos](https://huggingface.co/datasets/clinc/clinc_oos)) were taken from the [Setfit organization on Hugging Face](https://huggingface.co/datasets/SetFit). + +| dataset | tfidf | fasttext | model2vec + logreg | model2vec full finetune | setfit | minilm full finetune | +|:---------------------------|------:|---------:|-------------------:|------------------------:|-------:|---------------------:| +| 20_newgroups | 52.13 | 28.27 | 56.24 | 58.14 | 61.28 | 58.49 | +| ade | 79.76 | 76.25 | 79.20 | 75.65 | 80.23 | 82.81 | +| ag_news | 82.91 | 77.09 | 86.70 | 87.00 | 87.69 | 87.95 | +| amazon_counterfactual | 92.77 | 93.00 | 90.96 | 91.31 | 90.81 | 94.85 | +| amazon_polarity | 78.36 | 74.30 | 82.04 | 83.25 | 78.87 | 87.77 | +| banking77 | 69.60 | 53.23 | 79.95 | 75.61 | 56.42 | 48.72 | +| bbc | 96.60 | 95.30 | 95.80 | 96.41 | 94.59 | 96.70 | +| clinc_oos | 56.71 | 39.67 | 68.11 | 68.89 | 41.82 | 38.53 | +| emotion | 61.88 | 56.48 | 65.57 | 66.46 | 60.73 | 81.22 | +| enron_spam | 96.30 | 95.50 | 96.40 | 96.95 | 96.05 | 96.85 | +| ethos | 61.82 | 58.31 | 69.38 | 63.11 | 67.66 | 74.70 | +| hatespeech_offensive | 80.59 | 77.78 | 83.54 | 84.90 | 85.50 | 86.20 | +| imdb | 81.78 | 78.03 | 85.34 | 86.08 | 82.92 | 81.30 | +| insincere_questions | 90.79 | 91.89 | 93.50 | 93.24 | 94.43 | 94.09 | +| massive_intent | 67.54 | 63.99 | 73.83 | 72.00 | 62.32 | 67.00 | +| massive_scenario | 79.60 | 77.37 | 82.86 | 83.91 | 81.24 | 86.53 | +| senteval_cr | 74.89 | 73.78 | 77.03 | 79.11 | 82.52 | 86.00 | +| sst2 | 72.16 | 70.70 | 79.76 | 81.33 | 82.81 | 82.26 | +| sst5 | 31.12 | 33.02 | 32.34 | 37.71 | 40.96 | 35.65 | +| student | 79.74 | 75.39 | 83.20 | 85.60 | 87.27 | 89.22 | +| subj | 86.95 | 87.39 | 89.20 | 89.79 | 89.90 | 92.65 | +| toxic_conversations | 88.30 | 88.48 | 90.44 | 88.30 | 91.45 | 89.93 | +| trec | 67.68 | 66.50 | 56.34 | 55.97 | 57.79 | 73.58 | +| tweet_sentiment_extraction | 57.88 | 52.53 | 64.96 | 64.45 | 72.07 | 72.79 | +| tweet_stance_abortion | 63.58 | 64.11 | 71.99 | 68.33 | 66.86 | 69.72 | +| yelp_review_full | 45.21 | 41.77 | 48.42 | 50.96 | 48.38 | 48.54 | + + +| | tfidf | fasttext | model2vec + logreg | model2vec full finetune | setfit | minilm full finetune | +|:--------|------:|---------:|-------------------:|------------------------:|-------:|---------------------:| +| average | 72.9 | 68.9 | 76.3 | 76.3 | 74.7 | 77.1 | + + +The fully finetuned MiniLM model has the highest average score, followed by the two model2vec variants, which both outperform `setfit`. Full fine-tuning of model2vec scores higher than logistic regression on 17 of the 26 datasets, but loses heavily on a few, such as `ethos` and `banking77`, which leads to the same average score. Our advice is to test both if you can use `potion-base-32m`, and to use full fine-tuning if you are starting from another base model. + +The table below shows inference speed. The model2vec full finetune is about 80x faster than `setfit` and about 295x faster than the `minilm full finetune`, and comes close to the speed of `tfidf` and `fasttext`. + + +| | tfidf | fasttext | model2vec + logreg | model2vec full finetune | setfit | minilm full finetune | +|:-----------------|------:|---------:|-------------------:|------------------------:|-------:|---------------------:| +| samples / second | 93718 | 96288 | 32869 | 81018 | 1010 | 274 | -We use 14 classification datasets, using 1000 examples from the train set, and the full test set. No parameters were tuned on any validation set. All datasets were taken from the [Setfit organization on Hugging Face](https://huggingface.co/datasets/SetFit). - -| dataset | tfidf | model2vec + logreg | model2vec full finetune | setfit | bge-base + logreg | -|:---------------------------|--------:|---------------------:|--------------------------:|---------:|--------------------:| -| 20_newgroups | 50.71 | 56.24 | 57.94 | 61.29 | 67.39 | -| ade | 71.46 | 79.20 | 79.68 | 83.05 | 86.12 | -| ag_news | 81.68 | 86.70 | 87.20 | 88.01 | 88.95 | -| amazon_counterfactual | 85.18 | 90.96 | 91.93 | 95.51 | 92.74 | -| bbc | 95.09 | 95.80 | 97.21 | 96.60 | 97.50 | -| emotion | 59.28 | 65.57 | 67.11 | 72.86 | 65.63 | -| enron_spam | 96.00 | 96.40 | 96.85 | 97.45 | 97.30 | -| hatespeech_offensive | 66.45 | 83.54 | 85.61 | 87.69 | 84.92 | -| imdb | 80.44 | 85.34 | 85.59 | 86.00 | 92.25 | -| massive_scenario | 77.26 | 82.86 | 84.42 | 83.54 | 87.07 | -| senteval_cr | 65.61 | 77.03 | 79.47 | 86.15 | 90.53 | -| sst5 | 18.52 | 32.34 | 37.95 | 42.31 | 38.49 | -| student | 74.16 | 83.20 | 85.02 | 89.62 | 89.71 | -| subj | 86.39 | 89.20 | 89.85 | 93.80 | 94.55 | -| tweet_sentiment_extraction | 53.20 | 64.96 | 62.65 | 75.15 | 69.48 | - - -| | tfidf | model2vec + logreg | model2vec full finetune | setfit | bge-base + logreg | -|:--------|--------:|---------------------:|--------------------------:|---------:|--------------------:| -| average | 70.8 | 78.0 | 79.2 | 82.6 | 82.8 | - - - - -As can be seen see, full fine-tuning brings modest performance improvements in some cases, but very large ones in other cases, leading to a pretty large increase in average score. Our advice is to test both if you can use `potion-base-32m`, and to use full fine-tuning if you are starting from another base model. - -The speed difference between model2vec and the other models is immeense, with the full finetune being 35x faster than a setfit based on `all-minilm-l6-v2` on CPU and 200x faster than the`bge-base` transformer model. - - -| | tfidf | model2vec + logreg | model2vec full finetune | setfit | bge-base + logreg | -|:-----------------|--------:|---------------------:|--------------------------:|---------:|--------------------:| -| samples / second | 108434 | 17925 | 24744 | 716 | 118 | - - -The figure below shows the relationship between the number of sentences per second and the average training score, where we've included more transformer-based models for comparison. - -| ![Description](../assets/images/training_speed_vs_score.png) | -|:--:| -|*Figure: The average training score plotted against sentences per second (log scale).*| ## Ablations From 641d64f9d446b8961c41f401019c3796e345e7cf Mon Sep 17 00:00:00 2001 From: stephantul Date: Sun, 11 Oct 2026 08:30:58 +0200 Subject: [PATCH 5/7] update wording --- README.md | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/README.md b/README.md index 1a7b169..ff4b9d8 100644 --- a/README.md +++ b/README.md @@ -55,7 +55,7 @@ embeddings = model.encode(["It's dangerous to go alone!", "It's a secret to ever # Make sequences of token embeddings token_embeddings = model.encode_as_sequence(["It's dangerous to go alone!", "It's a secret to everybody."]) ``` -For advanced usage, see our [inference docs](https://minish.ai/packages/model2vec/inference). Instead of using one of our models, you can also distill your own Model2Vec model from a Sentence Transformer model. First, install the `distillation` extras with: +For advanced usage, see our [inference docs](https://minish.ai/packages/model2vec/inference). Instead of using one of our models, you can also distill your own Model2Vec model from a Sentence Transformer model. First, install the `distill` extra with: ```bash pip install model2vec[distill] @@ -74,7 +74,7 @@ m2v_model = distill(model_name="BAAI/bge-base-en-v1.5") m2v_model.save_pretrained("m2v_model") ``` -For advanced usage, see our [distillation docs](https://minish.ai/packages/model2vec/distillation), which includes some [distillation best practices](https://minish.ai/packages/model2vec/distillation#distillation-best-practices). After distillation, you can also fine-tune the distilled model, or a pre-trained model, for classification, similarity, regression, or pair similarity (e.g. query-document retrieval). First, make sure you install the `training` extras with: +For advanced usage, see our [distillation docs](https://minish.ai/packages/model2vec/distillation), which includes some [distillation best practices](https://minish.ai/packages/model2vec/distillation#distillation-best-practices). After distillation, you can also fine-tune the distilled model, or a pre-trained model, for classification, similarity, regression, or pair similarity (e.g. query-document retrieval). First, make sure you install the `train` extras with: ```bash pip install model2vec[train] From 07cd94fc92e05f59e27a0c7b09d9556def34ca52 Mon Sep 17 00:00:00 2001 From: stephantul Date: Sun, 11 Oct 2026 08:36:20 +0200 Subject: [PATCH 6/7] onnx extra --- model2vec/inference/README.md | 6 +++++- model2vec/train/README.md | 6 +++++- 2 files changed, 10 insertions(+), 2 deletions(-) diff --git a/model2vec/inference/README.md b/model2vec/inference/README.md index 600cf7b..1a81099 100644 --- a/model2vec/inference/README.md +++ b/model2vec/inference/README.md @@ -76,7 +76,11 @@ The underlying `StaticModel` is available as `pipeline.model`, and the head as ` # ONNX export -A pipeline can be exported to ONNX, including its head: +A pipeline can be exported to ONNX, including its head. This requires the `onnx` extra: + +```bash +pip install model2vec[onnx] +``` ```python from model2vec.onnx import export_model_to_onnx diff --git a/model2vec/train/README.md b/model2vec/train/README.md index c3806ce..48a1500 100644 --- a/model2vec/train/README.md +++ b/model2vec/train/README.md @@ -263,7 +263,11 @@ pipeline = StaticModelPipeline.from_pretrained("my_cool/project") Loading pipelines in this way is _extremely_ fast. It takes only 30ms to load a pipeline from disk. -A pipeline can also be exported to ONNX, including its head: +A pipeline can also be exported to ONNX, including its head. This requires the `onnx` extra: + +```bash +pip install model2vec[onnx] +``` ```python from model2vec.onnx import export_model_to_onnx From feb021f461c0ff25db01c7ac88f89657b596f7c9 Mon Sep 17 00:00:00 2001 From: stephantul Date: Sun, 11 Oct 2026 09:34:24 +0200 Subject: [PATCH 7/7] update wording --- model2vec/train/README.md | 14 +++++++++++++- 1 file changed, 13 insertions(+), 1 deletion(-) diff --git a/model2vec/train/README.md b/model2vec/train/README.md index 48a1500..3191e90 100644 --- a/model2vec/train/README.md +++ b/model2vec/train/README.md @@ -128,6 +128,8 @@ For multi-label classification, the weight of a class only applies to the texts classifier.fit(X, y, focal_gamma=2.0) ``` +The loss of each text is multiplied by `(1 - p) ** focal_gamma`, where `p` is the predicted probability of the true label. For multi-label classification, this is done for each class separately, based on its binary cross-entropy. `focal_gamma` can be combined with `class_weight`, in which case the class weights are applied on top of the focal loss. + ## Similarity and regression `StaticModelForSimilarity` and `StaticModelForRegression` learn to map each text to a target vector. Pass the vectors as `y`, as a 2D array or tensor, a list of lists, or a dataset column that holds lists of numbers. The output dimension of the model is set to the dimension of the vectors. @@ -202,13 +204,23 @@ With `n_layers=0`, a classifier gets a single linear layer. A similarity or pair | `validation_steps` | `None` | Validate every this many training steps. See [Validation](#validation). | | `device` | `"auto"` | The device to train on. `"auto"` picks CUDA, then MPS, then CPU. | | `random_seed` | `42` | The seed for the validation split, initialization and shuffling. | -| `token_dropout` | `0.0` | The fraction of tokens to randomly drop from each training text. This is a form of data augmentation, and has no effect during validation. Must be in `[0, 1)`. | +| `token_dropout` | `0.0` | See [Token dropout](#token-dropout). | | `class_weight` | `None` | Classification only. See [Imbalanced data](#imbalanced-data). | | `focal_gamma` | `0.0` | Classification only. See [Imbalanced data](#imbalanced-data). | | `temperature` | `0.05` | Pair similarity only. The temperature of the InfoNCE loss. | Calling `fit` re-initializes the head, the embeddings and the token weights from the static model, so calling it twice does not continue training. +## Token dropout + +`token_dropout` is a form of data augmentation. During training, every token of every text is dropped with this probability, and the text is embedded using the remaining tokens. A non-empty text always keeps at least one token. Token dropout has no effect during validation and inference. It must be in `[0, 1)`, and the default, `0.0`, disables it: + +```python +classifier.fit(X, y, token_dropout=0.1) +``` + +Token dropout is available for all models. For pair similarity, it is applied to both texts of a pair. + ## Validation Without an explicit validation set, `fit` holds out `test_size` of the data for validation, capped at 10,000 rows; pass an int to hold out an exact number of rows. For single-label classification, the split is stratified by class.