diff --git a/.pre-commit-config.yaml b/.pre-commit-config.yaml index b79d668..e5bdd79 100644 --- a/.pre-commit-config.yaml +++ b/.pre-commit-config.yaml @@ -38,4 +38,4 @@ repos: language: system types: [python] pass_filenames: false - args: [src/] + args: [model2vec/] diff --git a/model2vec/inference/evaluation.py b/model2vec/inference/evaluation.py index df1dcc6..63ee128 100644 --- a/model2vec/inference/evaluation.py +++ b/model2vec/inference/evaluation.py @@ -56,7 +56,7 @@ def _precision_recall_f1_support( def evaluate_single_or_multi_label( predictions: np.ndarray, - y: list[int] | list[str] | list[list[int]] | list[list[str]], + y: Sequence[Any], ) -> dict[str, dict[str, float]]: """Evaluate the classifier on a given dataset using a classification report. @@ -64,9 +64,10 @@ def evaluate_single_or_multi_label( overall accuracy, macro average, and weighted average. :param predictions: The predictions. - :param y: The ground truth labels. + :param y: The ground truth labels as a sequence. :return: A classification report, as a dictionary. """ + y = [label.tolist() if hasattr(label, "tolist") else label for label in y] if _is_multi_label_shaped(y): y = cast(list[list[str]] | list[list[int]], y) predictions = cast(np.ndarray, predictions) diff --git a/model2vec/model.py b/model2vec/model.py index 4e56589..2a3f6c5 100644 --- a/model2vec/model.py +++ b/model2vec/model.py @@ -459,7 +459,7 @@ def _encode_batch(self, sentences: Sequence[str], normalize: bool) -> np.ndarray ids = self.tokenize(sentences=sentences) dtype = self.embedding.dtype if dtype == np.int8: - dtype = np.float32 + dtype = np.dtype(np.float32) out = np.zeros((len(ids), self.dim), dtype=dtype) if self.token_mapping is None and self.weights is None: diff --git a/model2vec/train/README.md b/model2vec/train/README.md index 63ef8b2..b45839a 100644 --- a/model2vec/train/README.md +++ b/model2vec/train/README.md @@ -25,7 +25,7 @@ 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(model_name="minishlab/potion-base-32m") +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. @@ -35,6 +35,7 @@ Now that you have created the classifier, let's just train a model. The example ```python import numpy as np from datasets import load_dataset +from time import perf_counter # Load the subj dataset ds = load_dataset("setfit/subj") @@ -45,13 +46,13 @@ s = perf_counter() classifier = classifier.fit(train["text"], train["label"]) print(f"Training took {int(perf_counter() - s)} seconds.") -# Training took 81 seconds +# Training took 31 seconds classification_report = classifier.evaluate(ds["test"]["text"], ds["test"]["label"]) print(classification_report) -# Achieved 91.0 test accuracy +# Achieved 92.0 test accuracy ``` -As you can see, we got a pretty nice 91% accuracy, with only 81 seconds of training. +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. @@ -63,7 +64,7 @@ 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 67 milliseconds for 2000 instances on CPU. +# Took 66 milliseconds for 2000 instances on CPU. ``` ## Multi-label classification @@ -75,7 +76,7 @@ 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 multi-label dataset ds = load_dataset("google-research-datasets/go_emotions") @@ -113,6 +114,25 @@ Because the other pairs in a batch serve as negatives, the training and validati The InfoNCE temperature can be set with `temperature` (default `0.05`). It must be positive. +## 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: + +```python +from datasets import load_dataset + +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. + # Persistence You can turn a classifier into a lightweight inference pipeline, as follows: @@ -148,12 +168,11 @@ Our training architecture is set up to be extensible, with each task having a sp The core functionality of the `StaticModelForClassification` is contained in a couple of functions: * `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. -* `train_test_split`: governs the train test split before classification. -* `prepare_dataset`: Selects the `torch.Dataset` that will be used in the `Dataloader` during training. +* `_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. -The training loop itself lives 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 validation metrics like accuracy). # Results diff --git a/model2vec/train/base.py b/model2vec/train/base.py index ebd3d36..4b84317 100644 --- a/model2vec/train/base.py +++ b/model2vec/train/base.py @@ -2,11 +2,13 @@ import copy import logging -from collections.abc import Sequence +from collections.abc import Sequence, Sized from typing import Any, TypeVar import numpy as np import torch +from datasets import Column, DatasetDict, IterableColumn, IterableDataset, IterableDatasetDict +from datasets import Dataset as HFDataset from tokenizers import Encoding, Tokenizer from torch import nn from torch.nn.utils.rnn import pad_sequence @@ -14,12 +16,13 @@ from model2vec.inference import StaticModelPipeline from model2vec.model import DEFAULT_MAX_LENGTH, PathLike, StaticModel, _disable_padding, _get_unk_token_id -from model2vec.train.dataset import PairDataset, TextDataset +from model2vec.train.dataset import ColumnRows, PairDataset, TextDataset, has_only_strings, has_transform from model2vec.train.trainer import MetricsFn, default_metrics, resolve_device, run_training_loop from model2vec.train.utils import ( + MAX_VALIDATION_SIZE, get_probable_pad_token_id, + split_indices, to_pipeline, - train_test_split, ) logger = logging.getLogger(__name__) @@ -97,11 +100,24 @@ def __init__( _disable_padding(self.tokenizer) self.unk_token_id = _get_unk_token_id(self.tokenizer) - def _remove_unk(self, token_ids: list[int]) -> list[int]: - """Drop unknown tokens, mirroring `StaticModel.tokenize`.""" - if self.unk_token_id is None: - return token_ids - return [token_id for token_id in token_ids if token_id != self.unk_token_id] + def _tokenize_ids(self, texts: Sequence[str]) -> list[list[int]]: + """Tokenize texts into lists of token ids, dropping unknown tokens and truncating to `max_length` tokens. + + :param texts: The texts to tokenize. + :return: The token ids of each text. + """ + if self.max_length is not None: + truncate_length = self.max_length * 10 + texts = [text[:truncate_length] for text in texts] + encoded: list[Encoding] = self.tokenizer.encode_batch_fast(texts, add_special_tokens=False) + ids = [encoding.ids for encoding in encoded] + if self.unk_token_id is not None: + ids = [[token_id for token_id in token_ids if token_id != self.unk_token_id] for token_ids in ids] + return [token_ids[: self.max_length] for token_ids in ids] + + def _to_targets(self, labels: Any) -> torch.Tensor: + """Turn a batch of labels, such as vectors, into a float tensor of targets.""" + return torch.as_tensor(labels, dtype=torch.float32) def construct_weights(self) -> nn.Parameter: """Construct the weights for the model.""" @@ -275,11 +291,7 @@ def tokenize(self, texts: list[str]) -> torch.Tensor: :param texts: The texts to tokenize. :return: A 2D padded tensor """ - max_length = self.max_length - encoded: list[Encoding] = self.tokenizer.encode_batch_fast(texts, add_special_tokens=False) - encoded_ids: list[torch.Tensor] = [ - torch.Tensor(self._remove_unk(encoding.ids)[:max_length]).long() for encoding in encoded - ] + encoded_ids: list[torch.Tensor] = [torch.LongTensor(token_ids) for token_ids in self._tokenize_ids(texts)] return pad_sequence(encoded_ids, batch_first=True, padding_value=self.pad_id) @property @@ -327,31 +339,6 @@ def _determine_batch_size(self, batch_size: int | None, train_length: int) -> in return batch_size - def _check_val_split( - self, - X: list[str], - y: list, - X_val: list[str] | None, - y_val: list | None, - test_size: float, - ) -> tuple[list[str], list[str], Sequence, Sequence]: - if (X_val is not None) != (y_val is not None): - raise ValueError("Both X_val and y_val must be provided together, or neither.") - - if X_val is not None and y_val is not None: - # Additional check to ensure y_val is of the same type as y - if type(y_val[0]) != type(y[0]): - raise ValueError("X_val and y_val must be of the same type as X and y.") - - train_texts = X - train_labels = y - validation_texts = X_val - validation_labels = y_val - else: - train_texts, validation_texts, train_labels, validation_labels = train_test_split(X, y, test_size=test_size) - - return train_texts, validation_texts, train_labels, validation_labels - def _train( self, loss_function: nn.Module, @@ -420,59 +407,112 @@ def _determine_val_check_interval( return val_check_interval, check_val_every_epoch - def _tokenize_texts(self, X: list[str], max_length: int | None) -> list[list[int]]: - """Tokenize a list of texts into lists of token ids. + @staticmethod + def _check_inputs(**arguments: object) -> None: + """Check that every argument of `fit` is a sequence, an array, a tensor, or a column of a Hugging Face dataset. + + :param **arguments: The arguments, by name. None is skipped. + :raises ValueError: If an argument is a Hugging Face `Dataset` or `DatasetDict`, an iterable dataset or one of + its columns, a column of a dataset with a transform, a single string, or any other object that is not a + sequence, an array, or a tensor. + """ + for name, value in arguments.items(): + if isinstance(value, (IterableDataset, IterableDatasetDict, IterableColumn)): + raise ValueError( + f"{name} comes from an iterable Hugging Face dataset, which has no length. Pass a column of a " + "regular dataset instead, such as one loaded without `streaming=True`." + ) + if isinstance(value, (HFDataset, DatasetDict)): + raise ValueError( + f"{name} is a Hugging Face dataset. Pass one of its columns instead, such as `dataset['text']`." + ) + if isinstance(value, Column) and has_transform(value): + raise ValueError( + f"{name} is a column of a Hugging Face dataset with a transform. Apply the transform first with " + "`dataset.map(transform, batched=True)`, or pass a list." + ) + if isinstance(value, str) or ( + value is not None and not isinstance(value, (Sequence, np.ndarray, torch.Tensor)) + ): + raise ValueError( + f"{name} must be a list, a tuple, an array, a tensor, or a column of a Hugging Face dataset, got " + f"{type(value).__name__}." + ) + + @staticmethod + def _check_aligned(**columns: Sized) -> None: + """Check that the columns of the training or validation data have the same length. - :param X: The texts to tokenize. - :param max_length: The maximum length of the input in tokens. If this is None, no truncation is done. - :return: The tokenized texts. + :param **columns: The columns, by name. + :raises ValueError: If the columns don't all have the same length. """ - batch_size = 1024 - tokenized: list[list[int]] = [] - for batch_idx in trange(0, len(X), batch_size, desc="Tokenizing data"): - batch = X[batch_idx : batch_idx + batch_size] - if max_length is not None: - truncate_length = max_length * 10 - batch = [x[:truncate_length] for x in batch] - encoded = self.tokenizer.encode_batch_fast(batch, add_special_tokens=False) - tokenized.extend([self._remove_unk(encoding.ids)[:max_length] for encoding in encoded]) - - return tokenized - - def _prepare_dataset(self, X: list[str], y: torch.Tensor, max_length: int | None) -> TextDataset: - """Prepare a dataset. - - :param X: The texts. - :param y: The labels. - :param max_length: The maximum length of the input in tokens. If this is None, no truncation is done. - :return: A TextDataset. + lengths = {name: len(column) for name, column in columns.items()} + if len(set(lengths.values())) > 1: + raise ValueError(f"{' and '.join(lengths)} must have the same length, got {lengths}.") + + @staticmethod + def _check_texts(**texts: Sequence[str] | None) -> None: + """Check that all texts are strings. + + :param **texts: The texts, by name. None is skipped. + :raises ValueError: If a text is missing or is not a string. """ - return TextDataset(self._tokenize_texts(X, max_length), y, pad_id=self.pad_id) + for name, values in texts.items(): + if values is None: + continue + is_valid = ( + has_only_strings(values) + if isinstance(values, Column) + else all(isinstance(text, str) for text in values) + ) + if not is_valid: + raise ValueError(f"All texts in {name} must be strings.") - def _labels_to_tensor(self, labels: Any) -> torch.Tensor: - """Turn the labels into a tensor.""" - return labels + def _text_dataset(self, rows: ColumnRows, indices: np.ndarray | None = None) -> TextDataset: + """Create a dataset of labeled texts that are tokenized per batch. + + :param rows: The labeled texts, in a `text` and a `label` column. + :param indices: The indices of the rows that belong to the dataset. If None, all rows belong to it. + :return: The dataset. + """ + return TextDataset(rows, self._tokenize_ids, self._to_targets, indices, pad_id=self.pad_id) def _create_datasets( self, - X: list[str], + X: Sequence[str], y: Any, - X_val: list[str] | None, + X_val: Sequence[str] | None, y_val: Any | None, - test_size: float, + test_size: float | int, + stratify_by: Sequence[Any] | None = None, ) -> tuple[TextDataset, TextDataset]: - train_texts, validation_texts, train_labels, validation_labels = self._check_val_split( - X, y, X_val, y_val, test_size - ) - y_tensor = self._labels_to_tensor(train_labels) - y_val_tensor = self._labels_to_tensor(validation_labels) - - logger.info("Preparing train dataset.") - train_dataset = self._prepare_dataset(train_texts, y_tensor, self.max_length) - logger.info("Preparing validation dataset.") - val_dataset = self._prepare_dataset(validation_texts, y_val_tensor, self.max_length) + """Create the training and validation datasets. + + :param X: The training texts. + :param y: The training labels. + :param X_val: The validation texts. If None, the validation data is split off from `X` and `y`. + :param y_val: The validation labels. + :param test_size: The size of the validation split if `X_val` is None: a fraction of the data, capped at + `MAX_VALIDATION_SIZE` rows, or a number of rows if it is an int. + :param stratify_by: Validated single labels to stratify the validation split by. If None, the split is not + stratified. + :return: The train and validation datasets. + :raises ValueError: If only one of `X_val` and `y_val` is given, or if the texts and labels have different + lengths. + """ + if (X_val is None) != (y_val is None): + raise ValueError("Both X_val and y_val must be provided together, or neither.") + self._check_aligned(X=X, y=y) + self._check_texts(X=X, X_val=X_val) + rows = ColumnRows(text=X, label=y) + if X_val is not None and y_val is not None: + self._check_aligned(X_val=X_val, y_val=y_val) + return self._text_dataset(rows), self._text_dataset(ColumnRows(text=X_val, label=y_val)) - return train_dataset, val_dataset + train_indices, val_indices = split_indices( + len(rows), test_size, max_test_size=MAX_VALIDATION_SIZE, stratify_by=stratify_by + ) + return self._text_dataset(rows, train_indices), self._text_dataset(rows, val_indices) T = TypeVar("T", bound=BaseFinetuneable) diff --git a/model2vec/train/classifier.py b/model2vec/train/classifier.py index 1f98000..5ca5da7 100644 --- a/model2vec/train/classifier.py +++ b/model2vec/train/classifier.py @@ -2,11 +2,13 @@ import logging from collections import Counter +from collections.abc import Mapping, Sequence from itertools import chain from typing import Any, Literal, cast import numpy as np import torch +from datasets import Column from tokenizers import Tokenizer from torch import nn from tqdm import trange @@ -14,6 +16,7 @@ from model2vec.inference import evaluate_single_or_multi_label from model2vec.model import DEFAULT_MAX_LENGTH from model2vec.train.base import BaseFinetuneable +from model2vec.train.dataset import read_label_column from model2vec.train.utils import DEFAULT_RANDOM_SEED, seed_everything logger = logging.getLogger(__name__) @@ -42,6 +45,32 @@ def _multilabel_classifier_metrics(head_out: torch.Tensor, y: torch.Tensor, loss return {"loss": loss.item(), "accuracy": accuracy} +def _read_labels(y: LabelType, name: str) -> tuple[bool, Counter]: + """Determine whether labels are multi-label, and count the number of times each class occurs. + + :param y: The labels. If the first label is a list, multi-label classification is assumed. A column of a + Hugging Face dataset is read in batches, without converting it to Python objects. + :param name: The name of the labels, used in error messages. + :return: Whether the labels are multi-label, and the number of times each class occurs. + :raises ValueError: If the labels are inconsistent, or are not strings, integers, or lists of those. + """ + if isinstance(y, Column): + return read_label_column(y, name) + if isinstance(y, (np.ndarray, torch.Tensor)): + y = y.tolist() + + if isinstance(y[0], (str, int)): + if not all(isinstance(label, (str, int)) for label in y): + raise ValueError(f"Inconsistent label types in {name}. All labels must be strings or integers.") + return False, Counter(cast(list[str], y)) + if not all(isinstance(label, (list, tuple)) for label in y): + raise ValueError(f"Inconsistent label types in {name}. All labels must be lists or tuples.") + classes = list(chain.from_iterable(cast(list[list[str]], y))) + if not all(isinstance(label, (str, int)) for label in classes): + raise ValueError(f"Inconsistent label types in {name}. All classes must be strings or integers.") + return True, Counter(classes) + + class StaticModelForClassification(BaseFinetuneable): val_metric = "val_accuracy" early_stopping_direction = "max" @@ -142,16 +171,16 @@ def predict_proba(self, X: list[str], show_progress_bar: bool = False, batch_siz def fit( self, - X: list[str], + X: Sequence[str], y: LabelType, learning_rate: float = 1e-3, batch_size: int | None = None, min_epochs: int | None = None, max_epochs: int | None = -1, early_stopping_patience: int | None = 5, - test_size: float = 0.1, + test_size: float | int = 0.1, device: str = "auto", - X_val: list[str] | None = None, + X_val: Sequence[str] | None = None, y_val: LabelType | None = None, class_weight: Literal["balanced"] | dict[str, float] | torch.Tensor | None = None, validation_steps: int | None = None, @@ -170,6 +199,9 @@ def fit( If `X_val` and `y_val` are not provided, the function will automatically split the training data into a train and validation set using `test_size`. + The texts and labels are read and tokenized per batch. They can be lists, or columns of a Hugging Face + dataset, such as `dataset["text"]`, which are not loaded into memory. The dataset must not have a transform. + :param X: The texts to train on. :param y: The labels to train on. If the first element is a list, multi-label classification is assumed. :param learning_rate: The learning rate. @@ -179,7 +211,9 @@ def fit( If this is -1, the model trains until early stopping is triggered. :param early_stopping_patience: The patience for early stopping. If this is None, early stopping is disabled. - :param test_size: The test size for the train-test split. + :param test_size: The size of the validation split if `X_val` is None: a fraction of the data, capped at + 10,000 rows, or a number of rows if it is an int. The split is stratified if `y` holds single + labels. :param device: The device to train on. If this is "auto", the device is chosen automatically. :param X_val: The texts to be used for validation. :param y_val: The labels to be used for validation. @@ -191,27 +225,19 @@ def fit( :param token_dropout: The fraction of tokens to randomly drop from each training sample. Has no effect during validation. Must be in the range [0, 1). :return: The fitted model. - :raises ValueError: If either X_val or y_val are provided, but not both. """ seed_everything(random_seed) logger.info("Re-initializing model.") + self._check_inputs(X=X, y=y, X_val=X_val, y_val=y_val) - # Determine whether the task is multilabel based on the type of y. - self._initialize_on_labels(y) + label_counts = self._initialize_on_labels(y) + if y_val is not None: + self._check_validation_labels(y_val) self._initialize() - - if class_weight is not None: - if isinstance(class_weight, torch.Tensor): - logger.warning("You are passing a tensor as class weight. This will be removed in an upcoming version.") - if len(class_weight) != len(self.classes_): - raise ValueError("class_weight must have the same length as the number of classes.") - class_weight = {self.classes_[idx]: w for idx, w in enumerate(class_weight.tolist())} - resolved_class_weight = self._determine_class_weight(class_weight, y) - else: - resolved_class_weight = None - - train_dataset, val_dataset = self._create_datasets(X, y, X_val, y_val, test_size) - batch_size = self._determine_batch_size(batch_size, len(train_dataset)) + resolved_class_weight = self._resolve_class_weight(class_weight, label_counts) + train_dataset, val_dataset = self._create_datasets( + X, y, X_val, y_val, test_size, stratify_by=None if self.multilabel else y + ) if self.multilabel: loss_function: nn.Module = nn.BCEWithLogitsLoss(pos_weight=resolved_class_weight) @@ -225,7 +251,7 @@ def fit( learning_rate=learning_rate, train_dataset=train_dataset, val_dataset=val_dataset, - batch_size=batch_size, + batch_size=self._determine_batch_size(batch_size, len(train_dataset)), early_stopping_patience=early_stopping_patience, min_epochs=min_epochs, max_epochs=max_epochs, @@ -237,17 +263,32 @@ def fit( return self - def _determine_class_weight( - self, class_weight: dict[str, float] | Literal["balanced"], y: LabelType + def _resolve_class_weight( + self, + class_weight: Literal["balanced"] | dict[str, float] | torch.Tensor | None, + counts: Mapping[Any, int], + ) -> torch.Tensor | None: + """Turn the `class_weight` passed to `fit` into a tensor with one weight per class. + + :param class_weight: The class weight passed to `fit`. + :param counts: The number of times each class occurs. + :return: The weight of each class, or None if `class_weight` is None. + :raises ValueError: If `class_weight` is a tensor with the wrong length. + """ + if class_weight is None: + return None + if isinstance(class_weight, torch.Tensor): + logger.warning("You are passing a tensor as class weight. This will be removed in an upcoming version.") + if len(class_weight) != len(self.classes_): + raise ValueError("class_weight must have the same length as the number of classes.") + class_weight = {self.classes_[idx]: w for idx, w in enumerate(class_weight.tolist())} + return self._class_weight_from_counts(class_weight, counts) + + def _class_weight_from_counts( + self, class_weight: dict[str, float] | Literal["balanced"], counts: Mapping[Any, int] ) -> torch.Tensor: - """Determine the class weight for the classifier.""" + """Determine the class weight for the classifier from the number of times each class occurs.""" if class_weight == "balanced": - if self.multilabel: - y = cast(list[list[str]], y) - counts = Counter(chain.from_iterable(y)) - else: - y = cast(list[str], y) - counts = Counter(y) total = sum(counts.values()) n_classes = len(counts) # Reciprocal weight: upweight rare classes, downweight frequent ones @@ -271,43 +312,50 @@ def evaluate( predictions = self.predict(X, show_progress_bar=True, batch_size=batch_size, threshold=threshold) return evaluate_single_or_multi_label(predictions=predictions, y=y) - def _initialize_on_labels(self, y: LabelType) -> None: - """Sets the output dimensionality, the classes, and initializes the head. + def _initialize_on_labels(self, y: LabelType) -> Mapping[Any, int]: + """Sets the output dimensionality and the classes from the labels. - :param y: The labels. - :raises ValueError: If the labels are inconsistent. + :param y: The labels. A column of a Hugging Face dataset is read in batches, without converting it to + Python objects. + :return: The number of times each class occurs. """ - if isinstance(y[0], (str, int)): - y = cast(list[str], y) - # Check if all labels are strings or integers. - if not all(isinstance(label, (str, int)) for label in y): - raise ValueError("Inconsistent label types in y. All labels must be strings or integers.") - self.multilabel = False - classes = sorted(set(y)) - else: - y = cast(list[list[str]], y) - # Check if all labels are lists or tuples. - if not all(isinstance(label, (list, tuple)) for label in y): - raise ValueError("Inconsistent label types in y. All labels must be lists or tuples.") - self.multilabel = True - classes = sorted(set(chain.from_iterable(y))) - - self.classes_ = classes + self.multilabel, counts = _read_labels(y, "y") + self.classes_ = sorted(counts) self.out_dim = len(self.classes_) + return counts - def _labels_to_tensor(self, labels: Any) -> torch.Tensor: - """Convert a list or list of list of labels to a tensor.""" - if self.multilabel: - # Convert labels to multi-hot vectors - num_classes = len(self.classes_) - labels_tensor = torch.zeros(len(labels), num_classes, dtype=torch.float) - mapping = {label: idx for idx, label in enumerate(self.classes_)} - for i, sample_labels in enumerate(labels): - indices = [mapping[label] for label in sample_labels] - labels_tensor[i, indices] = 1.0 - else: - labels_tensor = torch.tensor( - [self.classes_.index(label) for label in cast(list[str], labels)], dtype=torch.long - ) + def _check_validation_labels(self, y_val: LabelType) -> None: + """Check that the validation labels match the labels the classifier was initialized on. - return labels_tensor + :param y_val: The validation labels. + :raises ValueError: If `y_val` is multi-label and `y` is not, or the other way around, or if `y_val` + contains classes that are not in `y`. + """ + multilabel, counts = _read_labels(y_val, "y_val") + if multilabel != self.multilabel: + raise ValueError("y_val must be multi-label if and only if y is multi-label.") + unknown = set(counts) - set(self.classes_) + if unknown: + raise ValueError(f"y_val contains classes that are not in y: {sorted(unknown, key=str)}.") + + def _to_targets(self, labels: Any) -> torch.Tensor: + """Turn a batch of labels into targets. + + :param labels: The labels. A tensor or an array is converted to a list first. + :return: The class indices, or the multi-hot vectors if the task is multilabel. + :raises ValueError: If a label is not one of the classes. + """ + if isinstance(labels, (torch.Tensor, np.ndarray)): + labels = labels.tolist() + index = {label: i for i, label in enumerate(self.classes_)} + try: + if not self.multilabel: + return torch.tensor([index[label] for label in labels], dtype=torch.long) + targets = torch.zeros(len(labels), len(index), dtype=torch.float) + for row, sample_labels in enumerate(labels): + if isinstance(sample_labels, (torch.Tensor, np.ndarray)): + sample_labels = sample_labels.tolist() + targets[row, [index[label] for label in sample_labels]] = 1.0 + return targets + except KeyError as error: + raise ValueError(f"Label {error.args[0]!r} is not one of the classes {self.classes_}.") from None diff --git a/model2vec/train/dataset.py b/model2vec/train/dataset.py index ec760a8..c5640ec 100644 --- a/model2vec/train/dataset.py +++ b/model2vec/train/dataset.py @@ -1,32 +1,292 @@ +from __future__ import annotations + +import logging +from abc import ABC, abstractmethod +from collections import Counter, defaultdict +from collections.abc import Callable, Iterator, Mapping, Sequence +from typing import Any + +import numpy as np +import pyarrow as pa +import pyarrow.compute as pc import torch +from datasets import Column +from datasets import Dataset as HFDataset from torch.nn.utils.rnn import pad_sequence -from torch.utils.data import DataLoader, Dataset +from torch.utils.data import BatchSampler, DataLoader, Dataset, RandomSampler, SequentialSampler + +logger = logging.getLogger(__name__) + +LABEL_COLUMN = "label" +TEXT_COLUMN = "text" +TEXT_A_COLUMN = "text_a" +TEXT_B_COLUMN = "text_b" + + +def _column_path(column: Column) -> tuple[HFDataset, str, list[str]]: + """Find the dataset a column belongs to, its top-level column, and the struct fields leading to the column.""" + names = [] + source: Any = column + while isinstance(source, Column): + names.append(source.column_name) + source = source.source + *fields, name = names + return source, name, fields[::-1] + + +def _struct_fields(array: pa.ChunkedArray, fields: list[str]) -> pa.ChunkedArray: + """Select nested struct fields from an Arrow array.""" + for field in fields: + array = pc.struct_field(array, field) + return array + + +def has_transform(column: Column) -> bool: + """Check whether a column belongs to a Hugging Face dataset with a transform, set with `with_transform`.""" + dataset, _, _ = _column_path(column) + return dataset.format["type"] == "custom" + + +def iter_column(column: Column, batch_size: int = 10_000) -> Iterator[pa.ChunkedArray]: + """Read a column of a Hugging Face dataset in batches of Arrow arrays, without converting it to Python objects. + + :param column: The column, which can be nested, such as `dataset["metadata"]["label"]`. Its dataset must not + have a transform. + :param batch_size: The number of rows in each batch. + :return: The values of the column, one batch at a time. + """ + dataset, name, fields = _column_path(column) + batches = dataset.select_columns([name]).with_format("arrow").iter(batch_size=batch_size) + return (_struct_fields(batch.column(name), fields) for batch in batches) + + +def column_type(column: Column) -> pa.DataType: + """Get the Arrow type of a column of a Hugging Face dataset, without reading the column. + + :param column: The column, which can be nested, such as `dataset["metadata"]["label"]`. + :return: The type of the values of the column. + """ + dataset, name, fields = _column_path(column) + return _struct_fields(dataset.data.column(name), fields).type + + +def has_only_strings(column: Column) -> bool: + """Check that a column of a Hugging Face dataset only holds strings, reading the strings only if it has nulls. + + :param column: The column, which can be nested, such as `dataset["metadata"]["text"]`. Its dataset must not + have a transform. + :return: Whether every row of the column is a string. + """ + dataset, name, fields = _column_path(column) + array = _struct_fields(dataset.data.column(name), fields) + if not (pa.types.is_string(array.type) or pa.types.is_large_string(array.type)): + return False + return array.null_count == 0 or not any(batch.null_count for batch in iter_column(column)) + + +def read_label_column(labels: Column, name: str) -> tuple[bool, Counter]: + """Determine whether a column of labels is multi-label, and count the number of times each class occurs. + + :param labels: A column of a Hugging Face dataset. If it holds lists, multi-label classification is assumed. + :param name: The name of the labels, used in error messages. + :return: Whether the labels are multi-label, and the number of times each class occurs. + :raises ValueError: If the labels are not strings, integers, or lists of those, or if a label is missing. + """ + label_type = column_type(labels) + multilabel = ( + pa.types.is_list(label_type) or pa.types.is_large_list(label_type) or pa.types.is_fixed_size_list(label_type) + ) + value_type = label_type.value_type if multilabel else label_type + if not (pa.types.is_string(value_type) or pa.types.is_large_string(value_type) or pa.types.is_integer(value_type)): + raise ValueError(f"Labels in {name} must be strings, integers, or lists of those, got {label_type}.") + counts: Counter = Counter() + for array in iter_column(labels): + values = pc.list_flatten(array) if multilabel else array + if array.null_count or values.null_count: + raise ValueError(f"Labels in {name} must not be missing.") + value_counts = pc.value_counts(values) + counts.update(dict(zip(value_counts.field("values").to_pylist(), value_counts.field("counts").to_pylist()))) + return multilabel, counts + + +def get_vector_dims_from_column(vectors: Column, name: str) -> set[int]: + """Get the dimensions of the vectors in a column of a Hugging Face dataset, reading the column in batches. + + :param vectors: A column of a Hugging Face dataset that holds lists of numbers. + :param name: The name of the vectors, used in error messages. + :return: The dimensions found, stopping as soon as more than one is found. + :raises ValueError: If a vector is missing, or if the column doesn't hold lists of numbers. + """ + array_type = column_type(vectors) + is_list = pa.types.is_list(array_type) or pa.types.is_large_list(array_type) + if not (is_list or pa.types.is_fixed_size_list(array_type)): + raise ValueError(f"{name} must hold lists of numbers, got {array_type}.") + value_type = array_type.value_type + if not (pa.types.is_floating(value_type) or pa.types.is_integer(value_type)): + raise ValueError(f"{name} must hold lists of numbers, got {array_type}.") + dims: set[int] = set() + for array in iter_column(vectors): + if array.null_count or pc.list_flatten(array).null_count: + raise ValueError(f"Vectors in {name} must not be missing.") + bounds = pc.min_max(pc.list_value_length(array)) + dims |= {bounds["min"].as_py(), bounds["max"].as_py()} + if len(dims) > 1: + break + return dims + + +def _list_strata(labels: Sequence[Any]) -> list[np.ndarray]: + """Group the indices of a list of labels by label, in order of first occurrence.""" + indices_by_label: dict[Any, list[int]] = defaultdict(list) + for index, label in enumerate(labels): + indices_by_label[label].append(index) + return [np.asarray(indices) for indices in indices_by_label.values()] + + +def _array_strata(labels: np.ndarray) -> list[np.ndarray]: + """Group the indices of an array of labels by label, in order of first occurrence.""" + _, first, codes = np.unique(labels, return_index=True, return_inverse=True) + codes = np.argsort(np.argsort(first))[codes] + return np.split(np.argsort(codes, kind="stable"), np.cumsum(np.bincount(codes))[:-1]) + + +def _column_strata(labels: Column) -> list[np.ndarray]: + """Group the indices of a Hugging Face column of single labels by label, in order of first occurrence. + + :param labels: The labels, which must be strings or integers, none of them missing. + :return: The indices of each label. + """ + classes: dict[Any, int] = {} + batch_codes = [] + for array in iter_column(labels): + encoded = pc.dictionary_encode(array.combine_chunks()) + mapping = np.array([classes.setdefault(label, len(classes)) for label in encoded.dictionary.to_pylist()]) + batch_codes.append(mapping[encoded.indices.to_numpy(zero_copy_only=False)]) + codes = np.concatenate(batch_codes) + return np.split(np.argsort(codes, kind="stable"), np.cumsum(np.bincount(codes))[:-1]) + + +def stratify_indices(labels: Sequence[Any]) -> list[np.ndarray] | None: + """Group the indices of single labels by label, or return None if there are no labels or a label occurs once. + :param labels: Single labels that have been validated: strings or integers, none of them missing. + :return: The indices of each label, in order of first occurrence, or None if the labels can't be stratified. + """ + if not len(labels): + return None + if isinstance(labels, Column): + strata = _column_strata(labels) + elif isinstance(labels, torch.Tensor): + strata = _array_strata(labels.cpu().numpy()) + elif isinstance(labels, np.ndarray): + strata = _array_strata(labels) + else: + strata = _list_strata(labels) + if min(len(indices) for indices in strata) < 2: + logger.info("Some classes have fewer than 2 samples. Stratification is disabled.") + return None + return strata -class TextDataset(Dataset): - def __init__(self, tokenized_texts: list[list[int]], targets: torch.Tensor, pad_id: int = 0) -> None: - """A dataset of texts. - :param tokenized_texts: The tokenized texts. Each text is a list of token ids. - :param targets: The targets. +class ColumnRows: + def __init__(self, **columns: Sequence[Any] | torch.Tensor) -> None: + """Rows made up of aligned columns, which are only read for the rows that are fetched. + + :param **columns: The values of each column. A tensor, an array, or a column of a Hugging Face dataset is + indexed with a list of indices at once. Any other sequence is indexed item by item. + :raises ValueError: If the columns don't all have the same length. + """ + lengths = {name: len(column) for name, column in columns.items()} + if len(set(lengths.values())) > 1: + raise ValueError(f"All columns must have the same length, got {lengths}.") + self.columns: dict[str, Any] = columns + + def __len__(self) -> int: + """Return the number of rows.""" + return len(next(iter(self.columns.values()))) + + def __getitem__(self, indices: list[int]) -> dict[str, Any]: + """Fetch the rows at the given indices, as a mapping from column names to the values of those rows.""" + return { + name: column[indices] + if isinstance(column, (np.ndarray, torch.Tensor, Column)) + else [column[index] for index in indices] + for name, column in self.columns.items() + } + + +class _Batches(Dataset, ABC): + def __init__(self, rows: ColumnRows, indices: np.ndarray | None, pad_id: int) -> None: + """A dataset that fetches rows and turns them into items per batch. + + :param rows: The rows to draw items from. + :param indices: The indices of the rows that belong to this dataset. If None, all rows belong to it. :param pad_id: The id used to pad batches. Must match the `pad_id` of the model being trained. - :raises ValueError: If the number of targets does not match the number of texts. """ - if len(targets) != len(tokenized_texts): - raise ValueError("Number of targets does not match number of texts.") - self.tokenized_texts = tokenized_texts - self.targets = targets + self.rows = rows + self.indices = np.arange(len(rows)) if indices is None else indices self.pad_id = pad_id def __len__(self) -> int: """Return the length of the dataset.""" - return len(self.tokenized_texts) + return len(self.indices) - def __getitem__(self, index: int) -> tuple[list[int], torch.Tensor]: + def __getitem__(self, index: int) -> Any: """Gets an item.""" - return self.tokenized_texts[index], self.targets[index] + return self.__getitems__([index])[0] + + def __getitems__(self, indices: list[int]) -> list[Any]: + """Fetch and convert a batch of items at once.""" + return self._to_items(self.rows[self.indices[indices].tolist()]) + + @abstractmethod + def _to_items(self, rows: Mapping[str, Any]) -> list[Any]: + """Turn a batch of rows into items.""" + + @abstractmethod + def collate_fn(self, batch: list[Any]) -> tuple[torch.Tensor, torch.Tensor]: + """Collate a batch of items into model inputs and targets.""" + + def _drop_last(self, batch_size: int) -> bool: + """Whether to drop the last batch if it is smaller than `batch_size`.""" + return False + + def to_dataloader(self, shuffle: bool, batch_size: int = 32) -> DataLoader: + """Convert the dataset to a DataLoader.""" + sampler = RandomSampler(self) if shuffle else SequentialSampler(self) + return DataLoader( + self, + collate_fn=self.collate_fn, + batch_sampler=BatchSampler(sampler, batch_size=batch_size, drop_last=self._drop_last(batch_size)), + ) + + +class TextDataset(_Batches): + def __init__( + self, + rows: ColumnRows, + tokenize: Callable[[list[str]], list[list[int]]], + to_targets: Callable[[Any], torch.Tensor], + indices: np.ndarray | None = None, + pad_id: int = 0, + ) -> None: + """A dataset of labeled texts, which are tokenized per batch. - def collate_fn(self, batch: list[tuple[list[list[int]], int]]) -> tuple[torch.Tensor, torch.Tensor]: + :param rows: The labeled texts, in a `text` and a `label` column. + :param tokenize: Turns a batch of texts into lists of token ids. + :param to_targets: Turns a batch of labels into a tensor of targets. + :param indices: The indices of the rows that belong to this dataset. If None, all rows belong to it. + :param pad_id: The id used to pad batches. Must match the `pad_id` of the model being trained. + """ + super().__init__(rows, indices, pad_id) + self.tokenize = tokenize + self.to_targets = to_targets + + def _to_items(self, rows: Mapping[str, Any]) -> list[tuple[list[int], torch.Tensor]]: + """Tokenize the texts and turn the labels into targets.""" + return list(zip(self.tokenize(rows[TEXT_COLUMN]), self.to_targets(rows[LABEL_COLUMN]))) + + def collate_fn(self, batch: list[tuple[list[int], torch.Tensor]]) -> tuple[torch.Tensor, torch.Tensor]: """Collate function.""" texts, targets = zip(*batch) @@ -35,38 +295,28 @@ def collate_fn(self, batch: list[tuple[list[list[int]], int]]) -> tuple[torch.Te return padded, torch.stack(targets) - def to_dataloader(self, shuffle: bool, batch_size: int = 32) -> DataLoader: - """Convert the dataset to a DataLoader.""" - return DataLoader(self, collate_fn=self.collate_fn, shuffle=shuffle, batch_size=batch_size) - -class PairDataset(Dataset): +class PairDataset(_Batches): def __init__( self, - tokenized_texts_a: list[list[int]], - tokenized_texts_b: list[list[int]], + rows: ColumnRows, + tokenize: Callable[[list[str]], list[list[int]]], + indices: np.ndarray | None = None, pad_id: int = 0, ) -> None: - """A dataset of aligned text pairs. + """A dataset of aligned text pairs, which are tokenized per batch. - :param tokenized_texts_a: The tokenized first half of each pair. Each text is a list of token ids. - :param tokenized_texts_b: The tokenized second half of each pair. Each text is a list of token ids. + :param rows: The pairs, in a `text_a` and a `text_b` column. + :param tokenize: Turns a batch of texts into lists of token ids. + :param indices: The indices of the rows that belong to this dataset. If None, all rows belong to it. :param pad_id: The id used to pad batches. Must match the `pad_id` of the model being trained. - :raises ValueError: If the two halves don't have the same number of texts. """ - if len(tokenized_texts_a) != len(tokenized_texts_b): - raise ValueError("The two halves of a pair dataset must have the same number of texts.") - self.tokenized_texts_a = tokenized_texts_a - self.tokenized_texts_b = tokenized_texts_b - self.pad_id = pad_id - - def __len__(self) -> int: - """Return the length of the dataset.""" - return len(self.tokenized_texts_a) + super().__init__(rows, indices, pad_id) + self.tokenize = tokenize - def __getitem__(self, index: int) -> tuple[list[int], list[int]]: - """Gets an item.""" - return self.tokenized_texts_a[index], self.tokenized_texts_b[index] + def _to_items(self, rows: Mapping[str, Any]) -> list[tuple[list[int], list[int]]]: + """Tokenize both halves of each pair.""" + return list(zip(self.tokenize(rows[TEXT_A_COLUMN]), self.tokenize(rows[TEXT_B_COLUMN]))) def collate_fn(self, batch: list[tuple[list[int], list[int]]]) -> tuple[torch.Tensor, torch.Tensor]: """Collate function. @@ -83,7 +333,6 @@ def collate_fn(self, batch: list[tuple[list[int], list[int]]]) -> tuple[torch.Te return torch.stack([padded_a, padded_b]), torch.arange(len(texts_a)) - def to_dataloader(self, shuffle: bool, batch_size: int = 32) -> DataLoader: - """Convert the dataset to a DataLoader. A final batch with a single pair is dropped, unless it is the only pair.""" - drop_last = len(self) > 1 and len(self) % batch_size == 1 - return DataLoader(self, collate_fn=self.collate_fn, shuffle=shuffle, batch_size=batch_size, drop_last=drop_last) + def _drop_last(self, batch_size: int) -> bool: + """Drop a final batch with a single pair, unless it is the only pair.""" + return len(self) > 1 and len(self) % batch_size == 1 diff --git a/model2vec/train/pairs.py b/model2vec/train/pairs.py index 9cd8020..cec1ac5 100644 --- a/model2vec/train/pairs.py +++ b/model2vec/train/pairs.py @@ -1,16 +1,18 @@ from __future__ import annotations import logging +from collections.abc import Sequence from typing import TypeVar +import numpy as np import torch from tokenizers import Tokenizer from torch import nn from model2vec.model import DEFAULT_MAX_LENGTH from model2vec.train.base import BaseFinetuneable -from model2vec.train.dataset import PairDataset -from model2vec.train.utils import DEFAULT_RANDOM_SEED, seed_everything, train_test_split +from model2vec.train.dataset import ColumnRows, PairDataset +from model2vec.train.utils import DEFAULT_RANDOM_SEED, MAX_VALIDATION_SIZE, seed_everything, split_indices logger = logging.getLogger(__name__) @@ -117,30 +119,6 @@ def forward( # type: ignore[override] ids_b = torch.unique(input_ids[1], dim=0, return_inverse=True)[1] return out_a, out_b, ids_a, ids_b - def _check_pair_val_split( - self, - text_a: list[str], - text_b: list[str], - text_a_val: list[str] | None, - text_b_val: list[str] | None, - test_size: float, - ) -> tuple[list[str], list[str], list[str], list[str]]: - if len(text_a) != len(text_b): - raise ValueError("text_a and text_b must have the same length.") - if (text_a_val is not None) != (text_b_val is not None): - raise ValueError("Both text_a_val and text_b_val must be provided together, or neither.") - - if text_a_val is not None and text_b_val is not None: - if len(text_a_val) != len(text_b_val): - raise ValueError("text_a_val and text_b_val must have the same length.") - return text_a, text_a_val, text_b, text_b_val - - pairs = list(zip(text_a, text_b)) - train_pairs, val_pairs, _, _ = train_test_split(pairs, pairs, test_size=test_size) - train_a, train_b = map(list, zip(*train_pairs)) if train_pairs else ([], []) - val_a, val_b = map(list, zip(*val_pairs)) if val_pairs else ([], []) - return train_a, val_a, train_b, val_b - @staticmethod def _check_pair_splits(n_train: int, n_val: int) -> None: """Check that the training and validation sets each have at least two pairs. @@ -156,33 +134,61 @@ def _check_pair_splits(n_train: int, n_val: int) -> None: "a different test_size, or an explicit validation set." ) - def _prepare_pair_dataset(self, text_a: list[str], text_b: list[str], max_length: int | None) -> PairDataset: - """Tokenize both halves of a pair dataset. + def _pair_dataset(self, rows: ColumnRows, indices: np.ndarray | None = None) -> PairDataset: + """Create a dataset of pairs that are tokenized per batch. - :param text_a: The first half of each pair. - :param text_b: The second half of each pair. - :param max_length: The maximum length of the input in tokens. If this is None, no truncation is done. - :return: A PairDataset. + :param rows: The pairs, in a `text_a` and a `text_b` column. + :param indices: The indices of the rows that belong to the dataset. If None, all rows belong to it. + :return: The dataset. """ - return PairDataset( - self._tokenize_texts(text_a, max_length), - self._tokenize_texts(text_b, max_length), - pad_id=self.pad_id, - ) + return PairDataset(rows, self._tokenize_ids, indices, pad_id=self.pad_id) + + def _create_pair_datasets( + self, + text_a: Sequence[str], + text_b: Sequence[str], + text_a_val: Sequence[str] | None, + text_b_val: Sequence[str] | None, + test_size: float | int, + ) -> tuple[PairDataset, PairDataset]: + """Create the training and validation datasets of pairs. + + :param text_a: The first half of each training pair. + :param text_b: The second half of each training pair. + :param text_a_val: The first half of each validation pair. If None, the validation pairs are split off + from the training pairs. + :param text_b_val: The second half of each validation pair. + :param test_size: The size of the validation split if `text_a_val` is None: a fraction of the pairs, + capped at `MAX_VALIDATION_SIZE` rows, or a number of pairs if it is an int. + :return: The train and validation datasets. + :raises ValueError: If only one of `text_a_val` and `text_b_val` is given, or if the halves of the pairs have + different lengths. + """ + if (text_a_val is None) != (text_b_val is None): + raise ValueError("Both text_a_val and text_b_val must be provided together, or neither.") + self._check_aligned(text_a=text_a, text_b=text_b) + self._check_texts(text_a=text_a, text_b=text_b, text_a_val=text_a_val, text_b_val=text_b_val) + rows = ColumnRows(text_a=text_a, text_b=text_b) + if text_a_val is not None and text_b_val is not None: + self._check_aligned(text_a_val=text_a_val, text_b_val=text_b_val) + return self._pair_dataset(rows), self._pair_dataset(ColumnRows(text_a=text_a_val, text_b=text_b_val)) + + train_indices, val_indices = split_indices(len(rows), test_size, max_test_size=MAX_VALIDATION_SIZE) + return self._pair_dataset(rows, train_indices), self._pair_dataset(rows, val_indices) def fit( self: T, - text_a: list[str], - text_b: list[str], + text_a: Sequence[str], + text_b: Sequence[str], learning_rate: float = 1e-3, batch_size: int | None = None, min_epochs: int | None = None, max_epochs: int | None = -1, early_stopping_patience: int | None = 5, - test_size: float = 0.1, + test_size: float | int = 0.1, device: str = "auto", - text_a_val: list[str] | None = None, - text_b_val: list[str] | None = None, + text_a_val: Sequence[str] | None = None, + text_b_val: Sequence[str] | None = None, validation_steps: int | None = None, random_seed: int = DEFAULT_RANDOM_SEED, temperature: float = 0.05, @@ -200,6 +206,9 @@ def fit( If `text_a_val` and `text_b_val` are not provided, the function will automatically split the training data into a train and validation set using `test_size`. + The pairs are read and tokenized per batch. The halves can be lists, or columns of a Hugging Face dataset, + such as `dataset["query"]`, which are not loaded into memory. The dataset must not have a transform. + :param text_a: The first half of each training pair. :param text_b: The second half of each training pair. :param learning_rate: The learning rate. @@ -209,7 +218,8 @@ def fit( If this is -1, the model trains until early stopping is triggered. :param early_stopping_patience: The patience for early stopping. If this is None, early stopping is disabled. - :param test_size: The test size for the train-test split. + :param test_size: The size of the validation split if `text_a_val` is None: a fraction of the pairs, capped + at 10,000 pairs, or a number of pairs if it is an int. :param device: The device to train on. If this is "auto", the device is chosen automatically. :param text_a_val: The first half of each validation pair. :param text_b_val: The second half of each validation pair. @@ -221,21 +231,16 @@ def fit( """ seed_everything(random_seed) logger.info("Re-initializing model.") + self._check_inputs(text_a=text_a, text_b=text_b, text_a_val=text_a_val, text_b_val=text_b_val) loss_function = PairInfoNCELoss(temperature=temperature) - train_a, val_a, train_b, val_b = self._check_pair_val_split(text_a, text_b, text_a_val, text_b_val, test_size) - self._check_pair_splits(len(train_a), len(val_a)) - self._initialize() - - logger.info("Preparing train dataset.") - train_dataset = self._prepare_pair_dataset(train_a, train_b, self.max_length) - logger.info("Preparing validation dataset.") - val_dataset = self._prepare_pair_dataset(val_a, val_b, self.max_length) - + train_dataset, val_dataset = self._create_pair_datasets(text_a, text_b, text_a_val, text_b_val, test_size) + self._check_pair_splits(len(train_dataset), len(val_dataset)) batch_size = self._determine_batch_size(batch_size, len(train_dataset)) if batch_size < 2: raise ValueError(f"batch_size must be at least 2, got {batch_size}.") + self._initialize() self._train( loss_function=loss_function, learning_rate=learning_rate, diff --git a/model2vec/train/similarity.py b/model2vec/train/similarity.py index 367f075..74ff0a0 100644 --- a/model2vec/train/similarity.py +++ b/model2vec/train/similarity.py @@ -1,19 +1,53 @@ from __future__ import annotations import logging -from typing import TypeVar +from collections.abc import Sequence +from typing import Any, TypeVar +import numpy as np import torch +from datasets import Column from tokenizers import Tokenizer from torch import nn from model2vec.model import DEFAULT_MAX_LENGTH from model2vec.train.base import BaseFinetuneable +from model2vec.train.dataset import get_vector_dims_from_column from model2vec.train.utils import DEFAULT_RANDOM_SEED, seed_everything logger = logging.getLogger(__name__) +def _vector_dim(vectors: Any, name: str) -> int: + """Get the dimension of a set of vectors, checking that every vector is present and has the same dimension. + + :param vectors: The vectors: a 2D tensor or array, a sequence of sequences of numbers, or a column of a Hugging + Face dataset that holds lists of numbers. + :param name: The name of the vectors, used in error messages. + :return: The dimension of the vectors. + :raises ValueError: If there are no vectors, if a vector is missing, if the vectors have different dimensions, + or if a column doesn't hold lists of numbers. + """ + if isinstance(vectors, (torch.Tensor, np.ndarray)): + if vectors.ndim != 2: + raise ValueError(f"{name} must be 2-dimensional, got {vectors.ndim} dimensions.") + return vectors.shape[1] + + if isinstance(vectors, Column): + dims = get_vector_dims_from_column(vectors, name) + else: + try: + dims = {len(vector) for vector in vectors} + except TypeError: + raise ValueError(f"Vectors in {name} must be sequences of numbers.") from None + + if not dims: + raise ValueError(f"{name} must not be empty.") + if len(dims) > 1: + raise ValueError(f"All vectors in {name} must have the same dimension, got {sorted(dims)}.") + return dims.pop() + + class CosineLoss(nn.Module): def __call__(self, x: torch.Tensor, y: torch.Tensor) -> torch.Tensor: """Returns the cosine distance loss function.""" @@ -65,17 +99,17 @@ def __init__( def fit( self: T, - X: list[str], - y: torch.Tensor, + X: Sequence[str], + y: torch.Tensor | Sequence[Sequence[float]], learning_rate: float = 1e-3, batch_size: int | None = None, min_epochs: int | None = None, max_epochs: int | None = -1, early_stopping_patience: int | None = 5, - test_size: float = 0.1, + test_size: float | int = 0.1, device: str = "auto", - X_val: list[str] | None = None, - y_val: torch.Tensor | None = None, + X_val: Sequence[str] | None = None, + y_val: torch.Tensor | Sequence[Sequence[float]] | None = None, validation_steps: int | None = None, random_seed: int = DEFAULT_RANDOM_SEED, token_dropout: float = 0.0, @@ -91,6 +125,10 @@ def fit( If `X_val` and `y_val` are not provided, the function will automatically split the training data into a train and validation set using `test_size`. + The texts and vectors are read and tokenized per batch. They can be lists or tensors, or columns of a + Hugging Face dataset, such as `dataset["text"]`, which are not loaded into memory. The dataset must not have a + transform. + :param X: The texts to train on. :param y: The vectors to train on. :param learning_rate: The learning rate. @@ -100,7 +138,8 @@ def fit( If this is -1, the model trains until early stopping is triggered. :param early_stopping_patience: The patience for early stopping. If this is None, early stopping is disabled. - :param test_size: The test size for the train-test split. + :param test_size: The size of the validation split if `X_val` is None: a fraction of the data, capped at + 10,000 rows, or a number of rows if it is an int. :param device: The device to train on. If this is "auto", the device is chosen automatically. :param X_val: The texts to be used for validation. :param y_val: The vectors to be used for validation. @@ -109,22 +148,24 @@ def fit( :param token_dropout: The fraction of tokens to randomly drop from each training sample. Has no effect during validation. Must be in the range [0, 1). :return: The fitted model. + :raises ValueError: If the vectors in `y_val` have a different dimension than those in `y`. """ seed_everything(random_seed) logger.info("Re-initializing model.") + self._check_inputs(X=X, y=y, X_val=X_val, y_val=y_val) + out_dim = _vector_dim(y, "y") + if y_val is not None and (val_dim := _vector_dim(y_val, "y_val")) != out_dim: + raise ValueError(f"The vectors in y_val have dimension {val_dim}, but those in y have dimension {out_dim}.") train_dataset, val_dataset = self._create_datasets(X, y, X_val, y_val, test_size) - batch_size = self._determine_batch_size(batch_size, len(train_dataset)) - - self.out_dim = train_dataset.targets.shape[1] + self.out_dim = out_dim self._initialize() - self._train( loss_function=self._build_loss_function(), learning_rate=learning_rate, train_dataset=train_dataset, val_dataset=val_dataset, - batch_size=batch_size, + batch_size=self._determine_batch_size(batch_size, len(train_dataset)), early_stopping_patience=early_stopping_patience, min_epochs=min_epochs, max_epochs=max_epochs, diff --git a/model2vec/train/utils.py b/model2vec/train/utils.py index a0e31fa..ad38c83 100644 --- a/model2vec/train/utils.py +++ b/model2vec/train/utils.py @@ -1,9 +1,10 @@ from __future__ import annotations import logging +import numbers import random -from collections import Counter, defaultdict -from typing import TYPE_CHECKING, Any, TypeVar +from collections.abc import Sequence +from typing import TYPE_CHECKING, Any import numpy as np import torch @@ -12,6 +13,7 @@ from model2vec.inference import StaticModelPipeline from model2vec.inference.mlp import Activation, Layer, MLPHead +from model2vec.train.dataset import stratify_indices if TYPE_CHECKING: from model2vec.train.base import BaseFinetuneable @@ -21,6 +23,7 @@ logger = logging.getLogger(__name__) DEFAULT_RANDOM_SEED = 42 +MAX_VALIDATION_SIZE = 10_000 _KNOWN_PAD_TOKENS = ("[PAD]", "") @@ -62,66 +65,75 @@ def to_pipeline(model: "BaseFinetuneable | StaticModelForClassification") -> Sta return StaticModelPipeline(static_model, head) -def _index(sequence: Any, indices: list[int]) -> Any: - """Index a sequence with a list of indices, preserving the container type.""" - if isinstance(sequence, list): - return [sequence[index] for index in indices] - return sequence[indices] +def _stratum_test_sizes(sizes: np.ndarray, n_test: int) -> np.ndarray: + """Divide `n_test` test items over strata in proportion to their sizes. + Every stratum gets at least one test item and keeps at least one train item. The total only differs from + `n_test` if this can't be met otherwise. -X_co = TypeVar("X_co") - - -def train_test_split( - X: list[X_co], - y: list, - test_size: float, -) -> tuple[list[X_co], list[X_co], list, list]: - """Split the data. - - For single-label classification, stratification is attempted (if possible). - For multilabel classification, a random split is performed. + :param sizes: The number of items in each stratum. Each stratum has at least two items. + :param n_test: The total number of test items. + :return: The number of test items in each stratum. + """ + quotas = sizes * n_test / sizes.sum() + counts = np.clip(np.round(quotas), 1, sizes - 1).astype(int) + while (excess := int(counts.sum()) - n_test) != 0: + step = 1 if excess > 0 else -1 + candidates = np.flatnonzero(counts > 1 if excess > 0 else counts < sizes - 1) + if not len(candidates): + break + index = candidates[np.argmax((counts - quotas)[candidates] * step)] + counts[index] -= step + return counts + + +def split_indices( + n: int, + test_size: float | int, + max_test_size: int | None = None, + stratify_by: Sequence[Any] | None = None, +) -> tuple[np.ndarray, np.ndarray]: + """Randomly split the indices `0..n-1` into sorted train and test indices. + + :param n: The number of items. + :param test_size: The size of the test split: a fraction of the items if it is a float, or a number of items + if it is an int. At least one item goes into the test split, and at least one into the train split if + `n > 1`. + :param max_test_size: The maximum number of items in the test split if `test_size` is a fraction. + If None, the test split is not capped. + :param stratify_by: The single label of each item, as strings or integers that have been validated. If every + label occurs at least twice, each label is split separately, in the same proportion. If None, the split is + not stratified. + :return: The train indices and the test indices. + :raises ValueError: If `test_size` is a bool. """ - rng = random.Random(DEFAULT_RANDOM_SEED) - n = len(X) - - stratify = isinstance(y, list) and len(y) > 0 and isinstance(y[0], (str, int)) - if stratify: - label_counts = Counter(y) - if min(label_counts.values()) < 2: - logger.info("Some classes have fewer than 2 samples. Stratification is disabled.") - stratify = False - - train_indices: list[int] - test_indices: list[int] - if stratify: - indices_by_label: dict[Any, list[int]] = defaultdict(list) - for index, label in enumerate(y): - indices_by_label[label].append(index) - - train_indices = [] - test_indices = [] - for indices in indices_by_label.values(): - indices = indices[:] - rng.shuffle(indices) - n_test = min(max(1, round(len(indices) * test_size)), len(indices) - 1) - test_indices.extend(indices[:n_test]) - train_indices.extend(indices[n_test:]) - rng.shuffle(train_indices) - rng.shuffle(test_indices) + if isinstance(test_size, bool): + raise ValueError("test_size must be a float or an int, not a bool.") + rng = np.random.default_rng(DEFAULT_RANDOM_SEED) + if isinstance(test_size, numbers.Integral): + n_test = int(test_size) else: - indices = list(range(n)) - rng.shuffle(indices) - n_test = min(max(1, round(n * test_size)), max(n - 1, 0)) - test_indices = indices[:n_test] - train_indices = indices[n_test:] - - X_train = _index(X, train_indices) - X_test = _index(X, test_indices) - y_train = _index(y, train_indices) - y_test = _index(y, test_indices) - - return X_train, X_test, y_train, y_test + n_test = round(n * test_size) + if max_test_size is not None: + n_test = min(n_test, max_test_size) + n_test = min(max(1, n_test), max(n - 1, 0)) + + strata = None if stratify_by is None else stratify_indices(stratify_by) + if strata is not None and len(strata) > n_test: + logger.info("There are more classes than validation samples. Stratification is disabled.") + strata = None + if strata is None: + indices = rng.permutation(n) + return np.sort(indices[n_test:]), np.sort(indices[:n_test]) + + train: list[np.ndarray] = [] + test: list[np.ndarray] = [] + test_sizes = _stratum_test_sizes(np.array([len(members) for members in strata]), n_test) + for members, n_members_test in zip(strata, test_sizes): + members = rng.permutation(members) + test.append(members[:n_members_test]) + train.append(members[n_members_test:]) + return np.sort(np.concatenate(train)), np.sort(np.concatenate(test)) def seed_everything(seed: int) -> None: @@ -131,8 +143,3 @@ def seed_everything(seed: int) -> None: torch.manual_seed(seed) if torch.cuda.is_available(): torch.cuda.manual_seed_all(seed) - - -def logit(x: torch.Tensor) -> torch.Tensor: - """Invert a sigmoid.""" - return -torch.log((1 / x) - 1) diff --git a/pyproject.toml b/pyproject.toml index 45f5b04..5e1ee82 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -61,7 +61,7 @@ dev = [ distill = ["torch", "transformers<5.4.0", "skeletoken>=0.6.0,<0.7.0"] onnx = ["onnx", "torch", "onnxruntime", "onnxscript", "skeletoken>=0.6.0,<0.7.0"] -train = ["torch"] +train = ["torch", "datasets>=4.0.0"] quantization = ["scikit-learn"] integration = ["mteb"] tests = ["scikit-learn", "skops"] diff --git a/tests/test_inference.py b/tests/test_inference.py index ff62ac5..48c4d07 100644 --- a/tests/test_inference.py +++ b/tests/test_inference.py @@ -7,16 +7,36 @@ import numpy as np import pytest import skops.io +from datasets import Dataset from huggingface_hub.errors import EntryNotFoundError from sklearn.neural_network import MLPClassifier, MLPRegressor from sklearn.pipeline import make_pipeline from sklearn.preprocessing import StandardScaler +from model2vec.inference.evaluation import evaluate_single_or_multi_label from model2vec.inference.mlp import Activation, Layer, MLPHead from model2vec.inference.model import StaticModelPipeline, convert_legacy_pipeline from model2vec.model import StaticModel +@pytest.mark.parametrize( + ("labels", "predictions"), + [ + ([["a"], ["a", "b"]], np.array([np.array(["a"]), np.array(["b"])], dtype=object)), + (["a", "b"], np.array(["a", "a"])), + ([[0], [0, 1]], np.array([np.array([0]), np.array([1])], dtype=object)), + ([0, 1], np.array([0, 0])), + ], +) +def test_evaluate_on_dataset_column(labels: list, predictions: np.ndarray) -> None: + """Labels in a column of a Hugging Face dataset are evaluated like a list, also with a numpy or torch format.""" + dataset = Dataset.from_dict({"labels": labels}) + expected = evaluate_single_or_multi_label(predictions, labels) + assert evaluate_single_or_multi_label(predictions, dataset["labels"]) == expected + assert evaluate_single_or_multi_label(predictions, dataset.with_format("numpy")["labels"]) == expected + assert evaluate_single_or_multi_label(predictions, dataset.with_format("torch")["labels"]) == expected + + def test_init_predict(mock_inference_pipeline: StaticModelPipeline) -> None: """Test successful init and predict with StaticModelPipeline.""" target: list[str] | list[list[str]] diff --git a/tests/test_trainable.py b/tests/test_trainable.py index 0bdafa4..3e78d59 100644 --- a/tests/test_trainable.py +++ b/tests/test_trainable.py @@ -1,10 +1,13 @@ import logging +from collections import Counter, UserList from tempfile import TemporaryDirectory from typing import Any +from unittest.mock import patch import numpy as np import pytest import torch +from datasets import Dataset, DatasetDict, Features, Sequence, Value from skeletoken import TokenizerModel from tokenizers import Tokenizer from tokenizers.models import BPE @@ -16,12 +19,23 @@ from model2vec.model import StaticModel from model2vec.train import StaticModelForClassification from model2vec.train.base import BaseFinetuneable -from model2vec.train.dataset import PairDataset, TextDataset +from model2vec.train.classifier import _read_labels +from model2vec.train.dataset import ( + ColumnRows, + PairDataset, + TextDataset, + _column_strata, + iter_column, +) from model2vec.train.pairs import PairInfoNCELoss, StaticModelForPairSimilarity from model2vec.train.regression import StaticModelForRegression -from model2vec.train.similarity import StaticModelForSimilarity +from model2vec.train.similarity import StaticModelForSimilarity, _vector_dim from model2vec.train.trainer import _resolve_max_epochs, resolve_device, run_training_loop -from model2vec.train.utils import get_probable_pad_token_id, logit, seed_everything, train_test_split +from model2vec.train.utils import ( + get_probable_pad_token_id, + seed_everything, + split_indices, +) @pytest.mark.parametrize("n_layers", [0, 1, 2, 3]) @@ -58,10 +72,7 @@ def test_trainable_tokenizer_does_not_pad(mock_trained_pair_similarity_pipeline: """The tokenizer of a trainable model keeps its pad token, but doesn't pad.""" model = mock_trained_pair_similarity_pipeline assert model.tokenizer.padding is not None - assert ( - model._tokenize_texts(["word1 word2", "word2"], max_length=None)[1] - == model._tokenize_texts(["word2"], max_length=None)[0] - ) + assert model._tokenize_ids(["word1 word2", "word2"])[1] == model._tokenize_ids(["word2"])[0] def test_empty_texts_have_finite_gradients(mock_vectors: np.ndarray, mock_tokenizer: Tokenizer) -> None: @@ -70,7 +81,7 @@ 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 ) - dataset = model._prepare_dataset(["word1 word2", ""], torch.tensor([0, 1]), max_length=None) + dataset = model._text_dataset(ColumnRows(text=["word1 word2", ""], label=["0", "1"])) batch, y = next(iter(dataset.to_dataloader(shuffle=False, batch_size=2))) nn.functional.cross_entropy(model(batch), y).backward() @@ -265,16 +276,29 @@ def test_fit_sets_token_dropout_and_disables_it_after_training() -> None: assert torch.allclose(first, second) +def _pretokenized(texts: list[Any]) -> list[list[int]]: + return list(texts) + + def test_textdataset_init() -> None: - """Test the textdataset init.""" - dataset = TextDataset([[0], [1]], torch.arange(2)) + """A text dataset has one item per row, or per index if indices are given.""" + rows = ColumnRows(text=[[1], [2], [3]], label=torch.arange(3)) + assert len(TextDataset(rows, _pretokenized, torch.as_tensor)) == 3 + dataset = TextDataset(rows, _pretokenized, torch.as_tensor, indices=np.array([2, 0])) assert len(dataset) == 2 + assert dataset[0][0] == [3] + assert dataset[0][1].item() == 2 -def test_textdataset_init_incorrect() -> None: - """Test the textdataset init.""" +def test_column_rows() -> None: + """Rows are fetched per column, and tensors are indexed along their first dimension.""" + rows = ColumnRows(text=["a", "b", "c"], label=torch.arange(6).reshape(3, 2)) + assert len(rows) == 3 + fetched = rows[[2, 0]] + assert fetched["text"] == ["c", "a"] + assert fetched["label"].tolist() == [[4, 5], [0, 1]] with pytest.raises(ValueError): - TextDataset([[0]], torch.arange(2)) + ColumnRows(text=["a"], label=torch.arange(2)) def test_training_batch_padding_is_masked(mock_vectors: np.ndarray, mock_tokenizer: Tokenizer) -> None: @@ -282,7 +306,7 @@ def test_training_batch_padding_is_masked(mock_vectors: np.ndarray, mock_tokeniz s = StaticModelForClassification(vectors=torch.from_numpy(mock_vectors).float(), tokenizer=mock_tokenizer, pad_id=1) texts = ["word2", "word2 word3"] - dataset = s._prepare_dataset(texts, torch.arange(2), max_length=None) + dataset = s._text_dataset(ColumnRows(text=texts, label=["0", "1"])) batch, _ = next(iter(dataset.to_dataloader(shuffle=False, batch_size=2))) assert torch.equal(batch, s.tokenize(texts)) @@ -298,7 +322,7 @@ def test_unknown_tokens_are_dropped(mock_vectors: np.ndarray, mock_tokenizer: To expected = static.tokenize(texts) assert [row[row != s.pad_id].tolist() for row in s.tokenize(texts)] == expected - assert s._prepare_dataset(texts, torch.arange(2), max_length=None).tokenized_texts == expected + assert s._tokenize_ids(texts) == expected def test_tokenize_without_unk_token(mock_vectors: np.ndarray) -> None: @@ -316,7 +340,7 @@ def test_tokenize_without_unk_token(mock_vectors: np.ndarray) -> None: expected = static.tokenize(texts) assert [row[row != s.pad_id].tolist() for row in s.tokenize(texts)] == expected - assert s._prepare_dataset(texts, torch.arange(2), max_length=None).tokenized_texts == expected + assert s._tokenize_ids(texts) == expected def test_max_length_is_not_capped_by_the_static_model(mock_vectors: np.ndarray, mock_tokenizer: Tokenizer) -> None: @@ -327,7 +351,7 @@ def test_max_length_is_not_capped_by_the_static_model(mock_vectors: np.ndarray, s = StaticModelForClassification.from_static_model(model=static, max_length=4) assert s.tokenize(texts).shape[1] == 4 - assert [len(row) for row in s._prepare_dataset(texts, torch.arange(1), max_length=None).tokenized_texts] == [6] + assert [len(row) for row in s._tokenize_ids(texts)] == [4] assert len(s.to_static_model().tokenize(texts)[0]) == 4 # The static model keeps its own setting, and the trainer keeps its own once the static model changes. @@ -406,19 +430,19 @@ def test_convert_to_pipeline_regression(mock_trained_regression_pipeline: Static def test_pairdataset_init() -> None: """Test the pair dataset init.""" - dataset = PairDataset([[0], [1]], [[2], [3]]) + dataset = PairDataset(ColumnRows(text_a=[[0], [1]], text_b=[[2], [3]]), _pretokenized) assert len(dataset) == 2 def test_pairdataset_init_incorrect() -> None: """Test the pair dataset init with mismatched lengths.""" with pytest.raises(ValueError): - PairDataset([[0]], [[2], [3]]) + PairDataset(ColumnRows(text_a=[[0]], text_b=[[2], [3]]), _pretokenized) def test_pairdataset_collate() -> None: """Batches should stack the two padded halves into a single (2, batch, seq_len) tensor.""" - dataset = PairDataset([[1], [1, 2]], [[1, 2, 3], [1]], pad_id=0) + dataset = PairDataset(ColumnRows(text_a=[[1], [1, 2]], text_b=[[1, 2, 3], [1]]), _pretokenized, pad_id=0) batch, y = next(iter(dataset.to_dataloader(shuffle=False, batch_size=2))) assert batch.shape == (2, 2, 3) assert torch.equal(y, torch.tensor([0, 1])) @@ -452,6 +476,19 @@ def test_pair_infonce_loss_masks_duplicate_positives() -> None: assert loss.item() == pytest.approx(0.0, abs=1e-6) +def test_pair_infonce_loss_prefers_aligned_pairs() -> None: + """InfoNCE is low when each anchor matches its own positive and high when it matches another.""" + loss_fn = PairInfoNCELoss(temperature=0.05) + out_a = torch.tensor([[1.0, 0.0], [0.0, 1.0]]) + aligned = torch.tensor([[1.0, 0.0], [0.0, 1.0]]) + swapped = torch.tensor([[0.0, 1.0], [1.0, 0.0]]) + + assert loss_fn((out_a, aligned, _distinct(2), _distinct(2)), torch.arange(2)).item() == pytest.approx(0.0, abs=1e-6) + assert loss_fn((out_a, swapped, _distinct(2), _distinct(2)), torch.arange(2)).item() == pytest.approx( + 20.0, abs=1e-4 + ) + + def test_pair_infonce_loss_masks_other_positives_of_the_same_anchor() -> None: """Pairs with an identical first text are all positives for it, so they don't compete.""" out_a = torch.tensor([[1.0, 0.0], [1.0, 0.0]]) @@ -550,7 +587,7 @@ def test_classifier_keeps_head_when_dimensions_match(mock_vectors: np.ndarray, m def test_pair_similarity_forward(mock_trained_pair_similarity_pipeline: StaticModelForPairSimilarity) -> None: """The forward pass should return one head output per half of the pair batch.""" model = mock_trained_pair_similarity_pipeline - dataset = model._prepare_pair_dataset(["dog cat", "dog"], ["puppy", "kitten cat"], max_length=None) + dataset = model._pair_dataset(ColumnRows(text_a=["dog cat", "dog"], text_b=["puppy", "kitten cat"])) batch, _ = next(iter(dataset.to_dataloader(shuffle=False, batch_size=2))) with torch.no_grad(): @@ -621,9 +658,9 @@ def test_pair_similarity_fit_rejects_invalid_temperature_and_batch_size( def test_pairdataset_drops_single_pair_batches() -> None: """A final batch with a single pair is dropped, unless it is the only pair.""" - dataset = PairDataset([[1], [2], [3]], [[1], [2], [3]]) + dataset = PairDataset(ColumnRows(text_a=[[1], [2], [3]], text_b=[[1], [2], [3]]), _pretokenized) assert [len(y) for _, y in dataset.to_dataloader(shuffle=False, batch_size=2)] == [2] - single = PairDataset([[1]], [[1]]) + single = PairDataset(ColumnRows(text_a=[[1]], text_b=[[1]]), _pretokenized) assert [len(y) for _, y in single.to_dataloader(shuffle=False, batch_size=2)] == [1] @@ -632,7 +669,7 @@ def test_pair_similarity_forward_ids_identical_texts( ) -> None: """Identical texts in a batch get the same id.""" model = mock_trained_pair_similarity_pipeline - dataset = model._prepare_pair_dataset(["dog", "cat", "dog"], ["puppy", "puppy", "kitten"], None) + dataset = model._pair_dataset(ColumnRows(text_a=["dog", "cat", "dog"], text_b=["puppy", "puppy", "kitten"])) batch, _ = next(iter(dataset.to_dataloader(shuffle=False, batch_size=3))) with torch.no_grad(): _, _, ids_a, ids_b = model(batch) @@ -660,15 +697,6 @@ def test_convert_to_pipeline_pair_similarity( assert np.allclose(p1, p2, rtol=1e-5, atol=1e-4) -def test_train_test_split() -> None: - """Test the train test split function.""" - a, b, c, d = train_test_split(["0", "1", "2", "3"], ["1", "1", "0", "0"], 0.5) - assert len(a) == 2 - assert len(b) == 2 - assert len(c) == len(a) - assert len(d) == len(b) - - def test_y_val_none() -> None: """Test the y_val function.""" tokenizer = AutoTokenizer.from_pretrained("tests/data/test_tokenizer").backend_tokenizer @@ -793,27 +821,21 @@ def test_get_probable_pad_token_id_through_static_model(mock_vectors: np.ndarray assert get_probable_pad_token_id(model.tokenizer) == pad_id -def test_determine_class_weight(mock_trained_pipeline: StaticModelForClassification) -> None: +def test_resolve_class_weight(mock_trained_pipeline: StaticModelForClassification) -> None: """Test what the class weights are.""" w_dict = dict(zip(mock_trained_pipeline.classes, [0.5, 3])) c1, c2 = mock_trained_pipeline.classes_ - y: list[str] | list[list[str]] - if mock_trained_pipeline.multilabel: - y = [*[[c1]] * 100, *[[c2]] * 50] - else: - y = [*[c1] * 100, *[c2] * 50] - w = mock_trained_pipeline._determine_class_weight(w_dict, y) + counts = Counter({c1: 100, c2: 50}) + w = mock_trained_pipeline._resolve_class_weight(w_dict, counts) assert isinstance(w, torch.Tensor) assert w.tolist() == [0.5, 3] - w = mock_trained_pipeline._determine_class_weight(w_dict, y) - assert isinstance(w, torch.Tensor) - assert w.tolist() == [0.5, 3] - - w = mock_trained_pipeline._determine_class_weight("balanced", y) + w = mock_trained_pipeline._resolve_class_weight("balanced", counts) assert isinstance(w, torch.Tensor) assert w.tolist() == [0.75, 1.5] + assert mock_trained_pipeline._resolve_class_weight(None, counts) is None + def test_determine_interval() -> None: """Test the training interval and check_val_every_epoch are determined correctly.""" @@ -852,12 +874,6 @@ def test_determine_interval() -> None: assert check_val_every_epoch is None -def test_logit() -> None: - """Test on random data.""" - x = torch.arange(10).float() / 10 - assert torch.allclose(logit(torch.sigmoid(x)), x, atol=1e-6) - - def test_seed_everything_cuda(monkeypatch: pytest.MonkeyPatch) -> None: """seed_everything also seeds CUDA RNGs when CUDA is available.""" seeded_with: list[int] = [] @@ -969,3 +985,571 @@ def counting_step(self: torch.optim.lr_scheduler.ReduceLROnPlateau, metrics: flo check_val_every_epoch=None, ) assert len(step_calls) == 3 + + +def test_split_indices() -> None: + """The split is disjoint, sorted, complete, and has the requested test size.""" + train, test = split_indices(10, 0.3) + assert len(test) == 3 + assert sorted([*train, *test]) == list(range(10)) + assert list(train) == sorted(train) + assert list(test) == sorted(test) + + +def test_split_indices_absolute_and_capped_sizes() -> None: + """An int test size is a number of items, and a fractional test size can be capped.""" + assert len(split_indices(100, 7)[1]) == 7 + assert len(split_indices(100, 0.5, max_test_size=10)[1]) == 10 + assert len(split_indices(100, 0.05, max_test_size=10)[1]) == 5 + assert len(split_indices(100, 30, max_test_size=10)[1]) == 30 + + +def test_split_indices_stratified() -> None: + """A list of single labels is split per label, unless a label occurs only once.""" + labels = ["a"] * 6 + ["b"] * 4 + train, test = split_indices(len(labels), 0.5, stratify_by=labels) + assert sorted(labels[i] for i in test) == ["a"] * 3 + ["b"] * 2 + assert sorted([*train, *test]) == list(range(10)) + + labels = ["a"] * 9 + ["b"] + assert len(split_indices(len(labels), 0.5, stratify_by=labels)[1]) == 5 + + +def test_split_indices_numpy_int_and_bool() -> None: + """A numpy int test size is a number of items, and a bool test size is rejected.""" + assert len(split_indices(100, np.int64(7), max_test_size=50)[1]) == 7 + with pytest.raises(ValueError): + split_indices(100, True) + + +def test_split_indices_stratified_respects_test_size() -> None: + """A stratified split holds out exactly the requested number of items, even if it is capped.""" + labels = [str(i % 20) for i in range(1000)] + ["rare"] * 2 + train, test = split_indices(len(labels), 0.5, max_test_size=30, stratify_by=labels) + assert len(test) == 30 + assert {labels[i] for i in test} == set(labels) + assert sorted([*train, *test]) == list(range(len(labels))) + + labels = ["a", "a", "b", "b"] + train, test = split_indices(len(labels), 3, stratify_by=labels) + assert sorted(labels[i] for i in train) == sorted(labels[i] for i in test) == ["a", "b"] + + labels = [str(i % 20) for i in range(1000)] + assert len(split_indices(len(labels), 10, stratify_by=labels)[1]) == 10 + + +@pytest.mark.parametrize("labels", [["a"] * 30 + ["b"] * 10, [0] * 30 + [1] * 10]) +def test_split_indices_stratifies_columns_like_lists(labels: list[Any]) -> None: + """Columns of single labels are stratified exactly like lists, also after a selection or when nested.""" + dataset = Dataset.from_dict({"label": labels, "meta": [{"label": label} for label in labels]}) + selected = dataset.shuffle(seed=0).select(range(30)) + for column in (dataset["label"], dataset["meta"]["label"], selected["label"]): + expected = split_indices(len(column), 0.25, stratify_by=list(column)) + actual = split_indices(len(column), 0.25, stratify_by=column) + assert all(np.array_equal(a, b) for a, b in zip(actual, expected)) + test = split_indices(len(labels), 0.5, stratify_by=dataset["label"])[1] + assert sorted(dataset["label"][test.tolist()]) == labels[:1] * 15 + labels[-1:] * 5 + + +@pytest.mark.parametrize("labels", [["b"] * 30 + ["a"] * 10, [1] * 30 + [0] * 10]) +def test_split_indices_stratifies_arrays_like_lists(labels: list[Any]) -> None: + """Arrays and tensors of single labels are stratified exactly like lists.""" + expected = split_indices(len(labels), 0.25, stratify_by=labels) + arrays: list[Any] = [np.array(labels)] + if isinstance(labels[0], int): + arrays.append(torch.tensor(labels)) + for array in arrays: + actual = split_indices(len(array), 0.25, stratify_by=array) + assert all(np.array_equal(a, b) for a, b in zip(actual, expected)) + + +def test_split_indices_does_not_stratify_singleton_classes() -> None: + """Labels with a class that occurs once are split at random.""" + labels = ["a"] * 19 + ["b"] + column = Dataset.from_dict({"label": labels})["label"] + expected = split_indices(len(labels), 0.25) + for stratify_by in (labels, column): + actual = split_indices(len(labels), 0.25, stratify_by=stratify_by) + assert all(np.array_equal(a, b) for a, b in zip(actual, expected)) + + +@pytest.mark.parametrize( + ("model_class", "labels", "stratified"), + [ + (StaticModelForClassification, ["a", "b"] * 4, True), + (StaticModelForClassification, [["a"], ["b"]] * 4, False), + (StaticModelForRegression, [[0.5, 1.0]] * 8, False), + (StaticModelForSimilarity, [[0.5, 1.0]] * 8, False), + ], + ids=["single-label", "multilabel", "regression", "similarity"], +) +def test_fit_only_stratifies_single_labels( + model_class: type[BaseFinetuneable], + labels: list[Any], + stratified: bool, + mock_vectors: np.ndarray, + mock_tokenizer: Tokenizer, + monkeypatch: pytest.MonkeyPatch, +) -> None: + """Only single-label classifiers stratify the validation split.""" + monkeypatch.setattr("model2vec.train.base.run_training_loop", lambda **kwargs: kwargs["model"].state_dict()) + dataset = Dataset.from_dict({"text": _TRAIN_TEXTS, "label": labels}) + y = dataset["label"] + model = model_class(vectors=torch.from_numpy(mock_vectors).float(), tokenizer=mock_tokenizer) + with patch("model2vec.train.base.split_indices", wraps=split_indices) as split_mock: + model.fit(dataset["text"], y, test_size=0.5) # type: ignore[attr-defined] + stratify_by = split_mock.call_args.kwargs["stratify_by"] + assert (stratify_by is y) if stratified else (stratify_by is None) + + +def test_column_strata_match_list_strata_across_batches() -> None: + """Columns are grouped by label in batches, matching lists, in order of first occurrence.""" + labels = ["b", "a", "c", "a", "b", "c", "b", "a", "d", "d"] + column = Dataset.from_dict({"label": labels})["label"] + with patch("model2vec.train.dataset.iter_column", lambda column: iter_column(column, batch_size=3)): + strata = _column_strata(column) + assert [indices.tolist() for indices in strata] == [[0, 4, 6], [1, 3, 7], [2, 5], [8, 9]] + + +def test_classifier_counts_column_labels_across_batches() -> None: + """Labels in a column are counted and checked in batches.""" + dataset = Dataset.from_dict({"label": ["a", "b", "a"] * 3, "labels": [["a", "b"], ["a"], []] * 3}) + with patch("model2vec.train.dataset.iter_column", lambda column: iter_column(column, batch_size=2)): + assert _read_labels(dataset["label"], "y") == (False, Counter({"a": 6, "b": 3})) + assert _read_labels(dataset["labels"], "y") == (True, Counter({"a": 6, "b": 3})) + with pytest.raises(ValueError, match="must not be missing"): + _read_labels(Dataset.from_dict({"label": ["a"] * 5 + [None]})["label"], "y") + + +def test_column_rows_reads_other_sequences() -> None: + """Generic sequences are read item by item.""" + rows = ColumnRows(text=UserList(["a", "b", "c"]), label=np.array([0, 1, 2])) + fetched = rows[[2, 0]] + assert fetched["text"] == ["c", "a"] + assert fetched["label"].tolist() == [2, 0] + + +def test_column_rows_reads_hf_columns() -> None: + """Rows can be read from the columns of a Hugging Face dataset.""" + dataset = Dataset.from_dict({"a": ["x", "y", "z"], "b": [[1.0], [2.0], [3.0]]}) + rows = ColumnRows(text=dataset["a"], label=dataset["b"]) + assert len(rows) == 3 + assert rows[[2, 0]] == {"text": ["z", "x"], "label": [[3.0], [1.0]]} + + +def _assert_same_weights(expected: nn.Module, actual: nn.Module) -> None: + for (name, expected_tensor), actual_tensor in zip(expected.state_dict().items(), actual.state_dict().values()): + assert torch.equal(expected_tensor, actual_tensor), name + + +_TRAIN_TEXTS = ["word1", "word2", "word3", "word1 word2", "word2 word3", "word3 word1", "word1 word3", "word2"] +_VAL_TEXTS = ["word1 word2", "word3"] + + +def test_pair_similarity_fit_on_columns_matches_lists(mock_vectors: np.ndarray, mock_tokenizer: Tokenizer) -> None: + """Training on the columns of a dataset gives the same model as training on the same pairs as lists.""" + text_a = ["word1", "word2", "word3", "word1 word2", "word2 word3", "word3 word1"] + text_b = ["word2", "word3", "word1", "word3 word1", "word1", "word2 word2"] + val_a, val_b = ["word1", "word3"], ["word2", "word1"] + + from_lists = StaticModelForPairSimilarity(vectors=torch.from_numpy(mock_vectors).float(), tokenizer=mock_tokenizer) + from_lists.fit(text_a, text_b, text_a_val=val_a, text_b_val=val_b, max_epochs=3, batch_size=2, device="cpu") + + dataset = Dataset.from_dict({"a": text_a, "b": text_b}) + val_dataset = Dataset.from_dict({"a": val_a, "b": val_b}) + from_columns = StaticModelForPairSimilarity( + vectors=torch.from_numpy(mock_vectors).float(), tokenizer=mock_tokenizer + ) + from_columns.fit( + dataset["a"], + dataset["b"], + text_a_val=val_dataset["a"], + text_b_val=val_dataset["b"], + max_epochs=3, + batch_size=2, + device="cpu", + ) + + _assert_same_weights(from_lists, from_columns) + + +def test_pair_similarity_fit_on_columns_with_split(mock_vectors: np.ndarray, mock_tokenizer: Tokenizer) -> None: + """Without validation pairs, the validation pairs are split off from the columns.""" + model = StaticModelForPairSimilarity(vectors=torch.from_numpy(mock_vectors).float(), tokenizer=mock_tokenizer) + dataset = Dataset.from_dict({"a": ["word1", "word2", "word3", "word1 word2"], "b": ["word2", "word3", "word1", ""]}) + model.fit(dataset["a"], dataset["b"], test_size=2, max_epochs=1, device="cpu") + + +@pytest.mark.parametrize( + "labels, val_labels", + [ + (["a", "b", "a", "c", "b", "a", "c", "a"], ["a", "b"]), + ([["a"], ["b", "c"], ["a", "b"], [], ["c"], ["a"], ["b"], ["a", "c"]], [["a"], ["b"]]), + ], +) +def test_classifier_fit_on_columns_matches_lists( + mock_vectors: np.ndarray, mock_tokenizer: Tokenizer, labels: list, val_labels: list +) -> None: + """Training a classifier on the columns of a dataset gives the same model as training on lists.""" + from_lists = StaticModelForClassification(vectors=torch.from_numpy(mock_vectors).float(), tokenizer=mock_tokenizer) + from_lists.fit( + _TRAIN_TEXTS, labels, X_val=_VAL_TEXTS, y_val=val_labels, class_weight="balanced", max_epochs=3, batch_size=2 + ) + + dataset = Dataset.from_dict({"text": _TRAIN_TEXTS, "label": labels}) + val_dataset = Dataset.from_dict({"text": _VAL_TEXTS, "label": val_labels}) + from_columns = StaticModelForClassification( + vectors=torch.from_numpy(mock_vectors).float(), tokenizer=mock_tokenizer + ) + from_columns.fit( + dataset["text"], + dataset["label"], + X_val=val_dataset["text"], + y_val=val_dataset["label"], + class_weight="balanced", + max_epochs=3, + batch_size=2, + ) + + assert from_columns.classes_ == from_lists.classes_ + assert from_columns.multilabel == from_lists.multilabel + _assert_same_weights(from_lists, from_columns) + + +def test_regressor_fit_on_columns_matches_lists(mock_vectors: np.ndarray, mock_tokenizer: Tokenizer) -> None: + """Training a regressor on the columns of a dataset gives the same model as training on tensors.""" + y = torch.randn(len(_TRAIN_TEXTS), 3) + y_val = torch.randn(len(_VAL_TEXTS), 3) + + from_tensors = StaticModelForRegression(vectors=torch.from_numpy(mock_vectors).float(), tokenizer=mock_tokenizer) + from_tensors.fit(_TRAIN_TEXTS, y, X_val=_VAL_TEXTS, y_val=y_val, max_epochs=3, batch_size=2) + + dataset = Dataset.from_dict({"text": _TRAIN_TEXTS, "label": y.tolist()}) + val_dataset = Dataset.from_dict({"text": _VAL_TEXTS, "label": y_val.tolist()}) + from_columns = StaticModelForRegression(vectors=torch.from_numpy(mock_vectors).float(), tokenizer=mock_tokenizer) + from_columns.fit( + dataset["text"], + dataset["label"], + X_val=val_dataset["text"], + y_val=val_dataset["label"], + max_epochs=3, + batch_size=2, + ) + + assert from_columns.out_dim == 3 + _assert_same_weights(from_tensors, from_columns) + + +def test_fit_on_columns_with_split(mock_vectors: np.ndarray, mock_tokenizer: Tokenizer) -> None: + """Without validation data, the validation data is split off from the columns.""" + dataset = Dataset.from_dict({"text": _TRAIN_TEXTS, "label": [0, 1] * 4, "vector": [[0.5, 1.0]] * 8}) + classifier = StaticModelForClassification(vectors=torch.from_numpy(mock_vectors).float(), tokenizer=mock_tokenizer) + classifier.fit(dataset["text"], dataset["label"], max_epochs=1) + assert classifier.classes_ == [0, 1] + + regressor = StaticModelForRegression(vectors=torch.from_numpy(mock_vectors).float(), tokenizer=mock_tokenizer) + regressor.fit(dataset["text"], dataset["vector"], max_epochs=1) + assert regressor.out_dim == 2 + + +def test_classifier_rejects_float_labels_in_column(mock_vectors: np.ndarray, mock_tokenizer: Tokenizer) -> None: + """Labels in a column must be strings, integers, or lists of those.""" + model = StaticModelForClassification(vectors=torch.from_numpy(mock_vectors).float(), tokenizer=mock_tokenizer) + dataset = Dataset.from_dict({"text": _TRAIN_TEXTS, "label": [0.5] * 8}) + with pytest.raises(ValueError): + model.fit(dataset["text"], dataset["label"]) + + +@pytest.mark.parametrize("model_class", [StaticModelForClassification, StaticModelForRegression]) +def test_fit_rejects_datasets(model_class: type, mock_vectors: np.ndarray, mock_tokenizer: Tokenizer) -> None: + """`fit` takes columns, not whole datasets.""" + model = model_class(vectors=torch.from_numpy(mock_vectors).float(), tokenizer=mock_tokenizer) + labels: list = [[0.5, 1.0]] * 8 if model_class is StaticModelForRegression else ["a", "b"] * 4 + dataset = Dataset.from_dict({"text": _TRAIN_TEXTS, "label": labels}) + with pytest.raises(ValueError, match="columns"): + model.fit(dataset, labels) + with pytest.raises(ValueError, match="columns"): + model.fit(_TRAIN_TEXTS, dataset) + with pytest.raises(ValueError, match="columns"): + model.fit(_TRAIN_TEXTS, labels, X_val=DatasetDict({"train": dataset}), y_val=labels) + + +def test_pair_similarity_fit_rejects_datasets( + mock_trained_pair_similarity_pipeline: StaticModelForPairSimilarity, +) -> None: + """`fit` takes columns, not whole datasets.""" + model = mock_trained_pair_similarity_pipeline + dataset = Dataset.from_dict({"a": ["word1", "word2"], "b": ["word2", "word3"]}) + with pytest.raises(ValueError, match="columns"): + model.fit(dataset, ["word1", "word2"]) # type: ignore[arg-type] + with pytest.raises(ValueError, match="columns"): + model.fit(["word1", "word2"], ["word2", "word3"], text_a_val=dataset, text_b_val=["x"]) # type: ignore[arg-type] + + +def test_fit_caps_validation_split( + monkeypatch: pytest.MonkeyPatch, mock_vectors: np.ndarray, mock_tokenizer: Tokenizer +) -> None: + """With a fractional test size, fit holds out at most MAX_VALIDATION_SIZE rows.""" + monkeypatch.setattr("model2vec.train.base.MAX_VALIDATION_SIZE", 2) + monkeypatch.setattr("model2vec.train.pairs.MAX_VALIDATION_SIZE", 2) + sizes: list[tuple[int, int]] = [] + + def fake_run_training_loop(**kwargs: Any) -> dict[str, torch.Tensor]: + sizes.append((len(kwargs["train_loader"].dataset), len(kwargs["val_loader"].dataset))) + return kwargs["model"].state_dict() + + monkeypatch.setattr("model2vec.train.base.run_training_loop", fake_run_training_loop) + dataset = Dataset.from_dict({"text": _TRAIN_TEXTS, "label": ["a", "b"] * 4}) + classifier = StaticModelForClassification(vectors=torch.from_numpy(mock_vectors).float(), tokenizer=mock_tokenizer) + classifier.fit(_TRAIN_TEXTS, ["a", "b"] * 4, test_size=0.5) + classifier.fit(dataset["text"], dataset["label"], test_size=0.5) + classifier.fit(dataset["text"], dataset["label"], test_size=3) + pairs = StaticModelForPairSimilarity(vectors=torch.from_numpy(mock_vectors).float(), tokenizer=mock_tokenizer) + pairs.fit(_TRAIN_TEXTS, _TRAIN_TEXTS, test_size=0.5) + assert sizes == [(6, 2), (6, 2), (5, 3), (6, 2)] + + +def test_classifier_checks_validation_labels(mock_vectors: np.ndarray, mock_tokenizer: Tokenizer) -> None: + """Validation labels that don't match the training labels are rejected before training.""" + model = StaticModelForClassification(vectors=torch.from_numpy(mock_vectors).float(), tokenizer=mock_tokenizer) + with pytest.raises(ValueError, match="not in y"): + model.fit(_TRAIN_TEXTS, ["a", "b"] * 4, X_val=_VAL_TEXTS, y_val=["a", "c"]) + with pytest.raises(ValueError, match="multi-label"): + model.fit(_TRAIN_TEXTS, ["a", "b"] * 4, X_val=_VAL_TEXTS, y_val=[["a"], ["b"]]) + dataset = Dataset.from_dict({"text": _VAL_TEXTS, "label": ["a", "c"]}) + with pytest.raises(ValueError, match="not in y"): + model.fit(_TRAIN_TEXTS, ["a", "b"] * 4, X_val=dataset["text"], y_val=dataset["label"]) + + +@pytest.mark.parametrize("labels", [["a", None] * 4, [["a"], None] * 4, [["a"], ["b", None]] * 4]) +def test_classifier_rejects_missing_labels( + labels: list[Any], mock_vectors: np.ndarray, mock_tokenizer: Tokenizer +) -> None: + """Missing labels are rejected before training, in lists and in columns.""" + model = StaticModelForClassification(vectors=torch.from_numpy(mock_vectors).float(), tokenizer=mock_tokenizer) + dataset = Dataset.from_dict({"text": _TRAIN_TEXTS, "label": labels}) + with pytest.raises(ValueError): + model.fit(dataset["text"], dataset["label"]) + with pytest.raises(ValueError): + model.fit(_TRAIN_TEXTS, labels) + + +def test_fit_rejects_missing_texts(mock_vectors: np.ndarray, mock_tokenizer: Tokenizer) -> None: + """Missing texts are rejected before training, in lists and in columns.""" + texts = [*_TRAIN_TEXTS[:-1], None] + model = StaticModelForClassification(vectors=torch.from_numpy(mock_vectors).float(), tokenizer=mock_tokenizer) + dataset = Dataset.from_dict({"text": texts, "label": ["a", "b"] * 4}) + with pytest.raises(ValueError, match="X must be strings"): + model.fit(dataset["text"], dataset["label"]) + with pytest.raises(ValueError, match="X_val must be strings"): + model.fit(_TRAIN_TEXTS, ["a", "b"] * 4, X_val=["word1", None], y_val=["a", "b"]) # type: ignore[list-item] + with pytest.raises(ValueError, match="X must be strings"): + model.fit(dataset.shuffle(seed=0)["text"], dataset.shuffle(seed=0)["label"]) + filtered = dataset.filter(lambda row: row["text"] is not None) + model.fit(filtered["text"], filtered["label"], max_epochs=1) + pairs = StaticModelForPairSimilarity(vectors=torch.from_numpy(mock_vectors).float(), tokenizer=mock_tokenizer) + with pytest.raises(ValueError, match="text_b must be strings"): + pairs.fit(_TRAIN_TEXTS, texts) # type: ignore[arg-type] + + +def test_fit_rejects_non_string_column(mock_vectors: np.ndarray, mock_tokenizer: Tokenizer) -> None: + """A column of non-string values is rejected as texts.""" + model = StaticModelForClassification(vectors=torch.from_numpy(mock_vectors).float(), tokenizer=mock_tokenizer) + dataset = Dataset.from_dict({"text": list(range(8)), "label": ["a", "b"] * 4}) + with pytest.raises(ValueError, match="X must be strings"): + model.fit(dataset["text"], dataset["label"]) + + +def test_fit_names_mismatched_lengths(mock_vectors: np.ndarray, mock_tokenizer: Tokenizer) -> None: + """Columns of different lengths are reported by their argument names.""" + model = StaticModelForClassification(vectors=torch.from_numpy(mock_vectors).float(), tokenizer=mock_tokenizer) + with pytest.raises(ValueError, match="X and y must have the same length"): + model.fit(_TRAIN_TEXTS, ["a", "b"] * 3) + with pytest.raises(ValueError, match="X_val and y_val must have the same length"): + model.fit(_TRAIN_TEXTS, ["a", "b"] * 4, X_val=_VAL_TEXTS, y_val=["a"]) + pairs = StaticModelForPairSimilarity(vectors=torch.from_numpy(mock_vectors).float(), tokenizer=mock_tokenizer) + with pytest.raises(ValueError, match="text_a_val and text_b_val must have the same length"): + pairs.fit(_TRAIN_TEXTS, _TRAIN_TEXTS, text_a_val=_VAL_TEXTS, text_b_val=["word1"]) + + +def test_classifier_fit_on_nested_columns(mock_vectors: np.ndarray, mock_tokenizer: Tokenizer) -> None: + """Labels can be read from nested columns.""" + model = StaticModelForClassification(vectors=torch.from_numpy(mock_vectors).float(), tokenizer=mock_tokenizer) + dataset = Dataset.from_dict({"text": _TRAIN_TEXTS, "meta": [{"label": "a"}, {"label": "b"}] * 4}) + model.fit(dataset["text"], dataset["meta"]["label"], max_epochs=1) + assert model.classes_ == ["a", "b"] + + +@pytest.mark.parametrize("data_format", [None, "torch", "numpy"]) +def test_classifier_fit_on_formatted_columns( + data_format: str | None, mock_vectors: np.ndarray, mock_tokenizer: Tokenizer +) -> None: + """Columns of a dataset with a torch or numpy format can be used for training and validation.""" + dataset = Dataset.from_dict({"text": _TRAIN_TEXTS, "label": [0, 1] * 4}).with_format(data_format) + model = StaticModelForClassification(vectors=torch.from_numpy(mock_vectors).float(), tokenizer=mock_tokenizer) + model.fit(dataset["text"], dataset["label"], X_val=_VAL_TEXTS, y_val=[0, 1], max_epochs=1) # type: ignore[arg-type] + assert model.classes_ == [0, 1] + + +@pytest.mark.parametrize("data_format", [None, "torch", "numpy"]) +def test_classifier_fit_on_formatted_multilabel_columns( + data_format: str | None, mock_vectors: np.ndarray, mock_tokenizer: Tokenizer +) -> None: + """Multi-label columns of a dataset with a torch or numpy format can be used for training.""" + dataset = Dataset.from_dict({"text": _TRAIN_TEXTS, "label": [[1, 2], [1]] * 4}).with_format(data_format) + model = StaticModelForClassification(vectors=torch.from_numpy(mock_vectors).float(), tokenizer=mock_tokenizer) + model.fit(dataset["text"], dataset["label"], max_epochs=1) # type: ignore[arg-type] + assert model.multilabel + assert model.classes_ == [1, 2] + + +def test_classifier_fit_on_fixed_size_list_column(mock_vectors: np.ndarray, mock_tokenizer: Tokenizer) -> None: + """Multi-label columns with a fixed number of labels per row are multi-label.""" + features = Features({"text": Value("string"), "label": Sequence(Value("string"), length=1)}) + dataset = Dataset.from_dict({"text": _TRAIN_TEXTS, "label": [["a"], ["b"]] * 4}, features=features) + model = StaticModelForClassification(vectors=torch.from_numpy(mock_vectors).float(), tokenizer=mock_tokenizer) + model.fit(dataset["text"], dataset["label"], max_epochs=1) + assert model.multilabel + assert model.classes_ == ["a", "b"] + + +@pytest.mark.parametrize( + ("y", "y_val", "message"), + [ + ([[0.5, 1.0]] * 7 + [[0.5]], None, "same dimension"), + ([[0.5, 1.0]] * 7 + [None], None, "sequences of numbers"), + ([[0.5, 1.0]] * 8, [[0.5, 1.0, 1.5]] * 2, "y_val have dimension 3"), + (torch.ones(8), None, "2-dimensional"), + ([], None, "must not be empty"), + ], +) +def test_regressor_checks_vectors( + y: Any, y_val: Any, message: str, mock_vectors: np.ndarray, mock_tokenizer: Tokenizer +) -> None: + """Vectors that are missing or have different dimensions are rejected before training.""" + model = StaticModelForRegression(vectors=torch.from_numpy(mock_vectors).float(), tokenizer=mock_tokenizer) + X_val = None if y_val is None else _VAL_TEXTS + with pytest.raises(ValueError, match=message): + model.fit(_TRAIN_TEXTS, y, X_val=X_val, y_val=y_val) + + +@pytest.mark.parametrize( + ("vectors", "message"), + [ + ([[0.5, 1.0]] * 7 + [[0.5]], "same dimension"), + ([[0.5, 1.0]] * 7 + [None], "must not be missing"), + ([[0.5, None]] + [[0.5, 1.0]] * 7, "must not be missing"), + ([["a", "b"]] * 8, "lists of numbers"), + ([0.5] * 8, "lists of numbers"), + ], +) +def test_regressor_checks_vector_columns( + vectors: list, message: str, mock_vectors: np.ndarray, mock_tokenizer: Tokenizer +) -> None: + """Vectors in a column that are missing, not numbers, or have different dimensions are rejected before training.""" + dataset = Dataset.from_dict({"text": _TRAIN_TEXTS, "vector": vectors}) + model = StaticModelForRegression(vectors=torch.from_numpy(mock_vectors).float(), tokenizer=mock_tokenizer) + with pytest.raises(ValueError, match=message): + model.fit(dataset["text"], dataset["vector"]) + + +def test_iter_column_follows_selection() -> None: + """A column is read in batches of the selected rows.""" + dataset = Dataset.from_dict( + {"vector": [[float(i)] for i in range(10)], "meta": [{"vector": [float(i)]} for i in range(10)]} + ) + selected = dataset.shuffle(seed=0).select(range(7)) + for column in (selected["vector"], selected["meta"]["vector"]): + batches = list(iter_column(column, batch_size=3)) + assert [len(batch) for batch in batches] == [3, 3, 1] + assert [row for batch in batches for row in batch.to_pylist()] == list(selected["vector"]) + + +def test_vector_dim_checks_every_batch() -> None: + """Vectors in a column are checked across batches.""" + dataset = Dataset.from_dict({"vector": [[0.5, 1.0]] * 10_000 + [[0.5]]}) + for column in (dataset["vector"], dataset.shuffle(seed=0)["vector"]): + with pytest.raises(ValueError, match="same dimension"): + _vector_dim(column, "y") + assert _vector_dim(dataset.select(range(10_000))["vector"], "y") == 2 + + +def test_regressor_fit_on_fixed_size_vector_column(mock_vectors: np.ndarray, mock_tokenizer: Tokenizer) -> None: + """Vectors in a column with a fixed dimension can be used for training.""" + features = Features({"text": Value("string"), "vector": Sequence(Value("float32"), length=2)}) + dataset = Dataset.from_dict({"text": _TRAIN_TEXTS, "vector": [[0.5, 1.0]] * 8}, features=features) + model = StaticModelForRegression(vectors=torch.from_numpy(mock_vectors).float(), tokenizer=mock_tokenizer) + model.fit(dataset["text"], dataset["vector"], max_epochs=1) + assert model.out_dim == 2 + + +def test_fit_rejects_iterable_datasets(mock_vectors: np.ndarray, mock_tokenizer: Tokenizer) -> None: + """Iterable datasets and their columns have no length, and are rejected.""" + iterable = Dataset.from_dict({"text": _TRAIN_TEXTS, "label": ["a", "b"] * 4}).to_iterable_dataset() + model = StaticModelForClassification(vectors=torch.from_numpy(mock_vectors).float(), tokenizer=mock_tokenizer) + with pytest.raises(ValueError, match="X comes from an iterable"): + model.fit(iterable["text"], ["a", "b"] * 4) + with pytest.raises(ValueError, match="y comes from an iterable"): + model.fit(_TRAIN_TEXTS, iterable) + + +def test_fit_rejects_non_sequences(mock_vectors: np.ndarray, mock_tokenizer: Tokenizer) -> None: + """Objects that are not sequences, arrays, or tensors, such as columns of an Arrow-formatted dataset, are rejected.""" + dataset = Dataset.from_dict({"text": _TRAIN_TEXTS, "label": ["a", "b"] * 4}).with_format("arrow") + model = StaticModelForClassification(vectors=torch.from_numpy(mock_vectors).float(), tokenizer=mock_tokenizer) + with pytest.raises(ValueError, match="y must be a list, a tuple, an array, a tensor, or a column"): + model.fit(_TRAIN_TEXTS, dataset["label"]) + with pytest.raises(ValueError, match="X must be a list"): + model.fit(dataset["text"], ["a", "b"] * 4) + with pytest.raises(ValueError, match="X must be a list.*got str"): + model.fit("word1 word2", ["a", "b"] * 5) # type: ignore[arg-type] + + +def test_classifier_accepts_tuple_and_list_labels(mock_vectors: np.ndarray, mock_tokenizer: Tokenizer) -> None: + """Multi-label training and validation labels can be tuples and lists.""" + model = StaticModelForClassification(vectors=torch.from_numpy(mock_vectors).float(), tokenizer=mock_tokenizer) + model.fit(_TRAIN_TEXTS, [("a",), ("b",)] * 4, X_val=_VAL_TEXTS, y_val=[["a"], ["b"]], max_epochs=1) # type: ignore[arg-type] + assert model.multilabel + + +@pytest.mark.parametrize("y", [np.array(["a", "b"] * 4), np.array([0, 1] * 4), torch.tensor([0, 1] * 4)]) +def test_classifier_accepts_array_labels(mock_vectors: np.ndarray, mock_tokenizer: Tokenizer, y: Any) -> None: + """Single-label training and validation labels can be arrays and tensors.""" + model = StaticModelForClassification(vectors=torch.from_numpy(mock_vectors).float(), tokenizer=mock_tokenizer) + model.fit(_TRAIN_TEXTS, y, X_val=_VAL_TEXTS, y_val=y[:2], max_epochs=1) + assert not model.multilabel + assert model.classes_ == sorted(y.tolist()[:2]) + + +def test_fit_rejects_transformed_columns(mock_vectors: np.ndarray, mock_tokenizer: Tokenizer) -> None: + """Columns of a dataset with a transform are rejected, also when nested or after a selection.""" + dataset = Dataset.from_dict( + {"text": _TRAIN_TEXTS, "label": ["a", "b"] * 4, "meta": [{"label": "a"}, {"label": "b"}] * 4} + ) + transformed = dataset.with_transform(lambda batch: batch) + model = StaticModelForClassification(vectors=torch.from_numpy(mock_vectors).float(), tokenizer=mock_tokenizer) + with pytest.raises(ValueError, match="y is a column of a Hugging Face dataset with a transform"): + model.fit(dataset["text"], transformed["label"]) + with pytest.raises(ValueError, match="y is a column"): + model.fit(dataset["text"], transformed["meta"]["label"]) + with pytest.raises(ValueError, match="X is a column"): + model.fit(transformed.shuffle(seed=0)["text"], dataset["label"]) + pairs = StaticModelForPairSimilarity(vectors=torch.from_numpy(mock_vectors).float(), tokenizer=mock_tokenizer) + with pytest.raises(ValueError, match="text_a_val is a column"): + pairs.fit(_TRAIN_TEXTS, _TRAIN_TEXTS, text_a_val=transformed["text"], text_b_val=dataset["text"]) + model.fit(dataset["text"], transformed.with_format(None)["label"], max_epochs=1) + + +def test_classifier_to_targets(mock_vectors: np.ndarray, mock_tokenizer: Tokenizer) -> None: + """Class labels become indices or multi-hot vectors, and unknown labels are rejected.""" + model = StaticModelForClassification(vectors=torch.from_numpy(mock_vectors).float(), tokenizer=mock_tokenizer) + model.classes_, model.multilabel = ["a", "b"], False + assert model._to_targets(["b", "a"]).tolist() == [1, 0] + with pytest.raises(ValueError): + model._to_targets(["c"]) + model.classes_, model.multilabel = [0, 1], False # type: ignore[list-item] + assert model._to_targets(torch.tensor([1, 0])).tolist() == [1, 0] + model.classes_, model.multilabel = ["a", "b", "c"], True + assert model._to_targets([["a", "c"], []]).tolist() == [[1, 0, 1], [0, 0, 0]] + model.classes_, model.multilabel = [0, 1, 2], True # type: ignore[list-item] + assert model._to_targets([torch.tensor([0, 2]), np.array([1])]).tolist() == [[1, 0, 1], [0, 1, 0]] diff --git a/uv.lock b/uv.lock index 092c7f3..eedafb8 100644 --- a/uv.lock +++ b/uv.lock @@ -1284,11 +1284,13 @@ tests = [ { name = "skops" }, ] train = [ + { name = "datasets" }, { name = "torch" }, ] [package.metadata] requires-dist = [ + { name = "datasets", marker = "extra == 'train'", specifier = ">=4.0.0" }, { name = "huggingface-hub", specifier = ">=1.0.0" }, { name = "ipython", marker = "extra == 'dev'" }, { name = "jinja2" },