docs: update docs for training

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 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 `train` 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).
 

model2vec/inference/README.md

@@ -1,31 +1,104 @@
 # 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"])
+
+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.
 
-Let's assume you're using our [potion-edu classifier](https://huggingface.co/minishlab/potion-8m-edu-classifier).
+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': {...}}
+```
+
+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. This requires the `onnx` extra:
+
+```bash
+pip install model2vec[onnx]
 ```
 
-This should just work.
+```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.

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,

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