Classification

This documentation describes the current classification API used with STFormer-style models. It covers:

The high-level Classifier orchestration class (Classifier.py)
Data utilities and collators (classifier_utils.py)
Evaluation helpers (evaluation_utils.py)

The goal is to give a single place to understand how to prepare datasets, train sequence- or gene-level classifiers, and evaluate their performance.

Classifier

from Classifier import Classifier

class Classifier:
    def __init__(
        self,
        metadata_column,
        mode: str = "spot",
        classifier_type: str = "sequence",
        gene_class_dict=None,
        label_mapping=None,
        token_dictionary_file=None,
        filter_data=None,
        rare_threshold: float = 0.0,
        max_examples=None,
        max_examples_per_class=None,
        training_args=None,
        ray_config=None,
        freeze_layers: int = 0,
        forward_batch_size: int = 32,
        nproc: int = 4,
    )

A flexible orchestration class that:

Loads and filters a Hugging Face dataset.
Builds label columns for either sequence-level or gene-level classification.
Configures and launches Hugging Face Trainer (or Ray Tune for hyperparameter search).
Saves labeled datasets, trained models, and predictions.

Parameters (init)

metadata_column (str):
Name of the column in the dataset that contains the original labels (e.g. "cell_type", "treatment").
mode (str, default "spot"):
Tokenization mode, relevant for gene-level classification:
- "spot" – sequences are treated as single spots.
- "extended" – sequences are spot + neighbor concatenations. For gene-level classification, labels in the neighbor half are masked out with -100 so that only the first half is supervised.
classifier_type (str, default "sequence"):
Type of classifier to train:
- "sequence" – one label per example (cell/spot classification).
- "gene" – per-token labels (gene-level classification, implemented as a token classification head).
gene_class_dict (dict | None):
Required when classifier_type="gene". A mapping from class name to a list of gene IDs or tokens belonging to that class. Used by label_classes to build token-level labels.
label_mapping (dict | None):
Optional mapping from label value → class ID. If None, the mapping is inferred automatically from the data during prepare_data and stored on self.label_mapping.
token_dictionary_file (str | None):
Path to a pickle file containing a token vocabulary dictionary. Required for gene-level classification so gene identifiers in gene_class_dict can be mapped to token IDs.
filter_data (dict | None):
Optional metadata filters applied before labeling. The dict should map column_name → allowed_values.
rare_threshold (float, default 0.0):
Minimum fraction for a label to be kept. Any label whose frequency is below rare_threshold is dropped via remove_rare.
max_examples (int | None):
Global cap on the number of examples after filtering. If None, all remaining examples are kept.
max_examples_per_class (int | None):
Per-class cap when downsampling. If provided, at most this many examples per class are retained after shuffling.
training_args (dict | None):
Keyword arguments merged into Classifier.default_training_args to construct Hugging Face TrainingArguments.
ray_config (dict | None):
Hyperparameter search space for Ray Tune. If set and n_trials > 0 is passed to train, the training will be delegated to _ray_tune.
freeze_layers (int, default 0):
Number of initial transformer encoder layers to freeze.
forward_batch_size (int, default 32):
Batch size used during evaluation and for scaling the effective training batch size.
nproc (int, default 4):
Number of worker processes used in dataset filtering and mapping operations.

build_custom_tokenizer

from Classifier import build_custom_tokenizer

tokenizer = build_custom_tokenizer(
    token_dict_path: str,
    pad_token: str = "<pad>",
    mask_token: str = "<mask>",
)

Build a PreTrainedTokenizerFast from a custom token vocabulary stored as a pickle.

Parameters

token_dict_path (str):
Path to a pickle file containing a dict[str, int] vocabulary.
pad_token (str, default "<pad>"):
Token used for padding sequences.
mask_token (str, default "<mask>"):
Token used to indicate masked positions.

Returns

PreTrainedTokenizerFast – ready-to-use tokenizer instance.

_load_tokenizer

tokenizer = classifier._load_tokenizer(
    name_or_path: str,
)

Internal helper to load a tokenizer from:

A Hugging Face model name (e.g. "bert-base-uncased"),
A local tokenizer directory,
Or a pickled token dictionary (delegating to build_custom_tokenizer).

You usually do not call this directly; it is used inside train and evaluate.

prepare_data

dataset_path, label_map_path = classifier.prepare_data(
    input_data,
    output_directory: str,
    output_prefix: str,
)

Filter, downsample, and label a dataset for classification, then save it to disk.

The method:

Loads the dataset via classifier_utils.load_and_filter.
Drops rare labels using classifier_utils.remove_rare.
Optionally down-samples using downsample_and_shuffle.
Builds label columns:
- "sequence" classifier: one label column.
- "gene" classifier: per-token labels sequences via label_classes. In "extended" mode, only the first half is supervised (second half set to -100).
Saves the dataset (DatasetDict) with save_to_disk.
Saves the label mapping to a pickle.

Parameters

input_data (str | Dataset | DatasetDict):
Either a path to a dataset saved with load_from_disk or an in-memory dataset.
output_directory (str):
Directory to write the labeled dataset and label map.
output_prefix (str):
Prefix used for the saved dataset and mapping filenames.

Returns

dataset_path (str):
Path to the labeled dataset directory.
label_map_path (str):
Path to the pickled id_class_dict mapping.

train

trainer_or_result = classifier.train(
    model_checkpoint: str,
    dataset_path: str,
    output_directory: str,
    eval_dataset=None,
    n_trials: int = 0,
    test_size: float = 0.2,
    tokenizer_name_or_path: str | None = None,
)

Train a classifier (sequence or gene-level) or run Ray Tune hyperparameter search, depending on ray_config and n_trials.

Behavior

If ray_config is set and n_trials > 0:
- Delegates to _ray_tune, which runs Ray Tune over the provided hyperparameter space and returns the tuning result.
Otherwise:
1. Loads the labeled dataset from dataset_path.
2. If eval_dataset is None, creates a train/eval split using test_size.
3. Loads a tokenizer (model checkpoint or tokenizer_name_or_path).
4. Chooses collator:
  - DataCollatorForCellClassification when classifier_type="sequence".
  - DataCollatorForGeneClassification when classifier_type="gene".
5. Builds a classification model using AutoModelForSequenceClassification or AutoModelForTokenClassification.
6. Instantiates a Trainer with:
  - training args (TrainingArguments),
  - compute_metrics from evaluation_utils as compute_metrics,
  - the chosen collator and datasets.
7. Calls trainer.train().
8. Saves:
  - final model and tokenizer in {output_directory}/final_model,
  - predictions on eval set to {output_directory}/predictions.pkl.

Parameters

model_checkpoint (str):
Hugging Face model name or local path to the backbone to fine-tune.
dataset_path (str):
Path to the labeled dataset created by prepare_data.
output_directory (str):
Directory where checkpoints, final models, and prediction files are stored.
eval_dataset (str | Dataset | DatasetDict | None, default None):
Optional separate evaluation dataset. If None, the dataset at dataset_path is split into train/eval.
n_trials (int, default 0):
Number of Ray Tune trials to run. If > 0 and ray_config is set, triggers hyperparameter search.
test_size (float, default 0.2):
Fraction of data reserved for evaluation when eval_dataset is not provided.
tokenizer_name_or_path (str | None):
Optional explicit tokenizer source. If None, model_checkpoint is used.

Returns

Trainer – trained Hugging Face Trainer instance, or
Ray Tune result object when _ray_tune is used.

evaluate

metrics = classifier.evaluate(
    model_directory: str,
    eval_dataset_path: str,
    id_class_dict_file: str,
    output_directory: str,
    tokenizer_name_or_path: str | None = None,
)

Evaluate a trained model on a labeled dataset.

Behavior

Loads model + tokenizer from model_directory.
Loads eval dataset from eval_dataset_path.
Loads the class mapping from id_class_dict_file.
Recreates the correct collator (sequence vs gene).
Instantiates a Trainer with compute_metrics.
Calls trainer.evaluate() and returns the metrics.

Parameters

model_directory (str):
Directory containing the trained model and tokenizer (e.g. final_model).
eval_dataset_path (str):
Path to the evaluation dataset (same format as produced by prepare_data).
id_class_dict_file (str):
Pickled mapping between class names and class IDs.
output_directory (str):
Directory used for any logs or plots (if produced).
tokenizer_name_or_path (str | None):
Optional tokenizer override; if None, the tokenizer in model_directory is used.

Returns

dict – metrics dictionary as returned by Trainer.evaluate() plus the extra statistics computed by compute_metrics.

plot_confusion_matrix

classifier.plot_confusion_matrix(
    conf_mat,
    output_directory: str,
    output_prefix: str,
    class_order: list[str],
)

Thin wrapper around evaluation_utils.plot_confusion_matrix that saves a heatmap of the confusion matrix.

Parameters

conf_mat (np.ndarray or array-like):
Confusion matrix of shape (num_classes, num_classes).
output_directory (str):
Directory where the figure is saved.
output_prefix (str):
Prefix for the output filename.
class_order (list[str]):
Ordered list of class names for axes labeling.

plot_predictions

classifier.plot_predictions(
    predictions_file: str,
    id_class_dict_file: str,
    title: str,
    output_directory: str,
    output_prefix: str,
    class_order: list[str],
)

Wrapper around evaluation_utils.plot_predictions to visualize prediction outputs (saved from Trainer.predict).

Parameters

predictions_file (str):
Path to the pickled predictions dict.
id_class_dict_file (str):
Pickled mapping between class IDs and names.
title (str):
Title for the plot.
output_directory (str):
Directory where the prediction plot is saved.
output_prefix (str):
Prefix for the output filename.
class_order (list[str]):
Ordered list of class names for x/y axes.

Dataset Utilities (`classifier_utils.py`)

The classifier_utils.py helper module provides a set of functions and collators used by Classifier.

load_and_filter

from classifier_utils import load_and_filter

data = load_and_filter(
    filter_data,
    nproc: int,
    input_data_file,
)

Load a dataset (from disk or in-memory) and apply metadata filters.

filter_data (dict | None):
Mapping from column name to allowed values.
nproc (int):
Number of worker processes used for filtering.
input_data_file (str | Dataset | DatasetDict):
Path to a dataset saved with load_from_disk or a Dataset/DatasetDict.

Returns a filtered Dataset or DatasetDict.

remove_rare

from classifier_utils import remove_rare

data = remove_rare(
    data,
    rare_threshold: float,
    state_key: str,
    nproc: int,
)

Drop examples whose label frequency is below a threshold.

data (Dataset | DatasetDict): Input dataset.
rare_threshold (float): Minimum frequency (0–1) for label retention.
state_key (str): Column containing label values.
nproc (int): Number of worker processes.

downsample_and_shuffle

from classifier_utils import downsample_and_shuffle

data = downsample_and_shuffle(
    data,
    max_ncells: int | None,
    max_ncells_per_class: int | None,
    cell_state_dict: dict,
)

Shuffle and optionally downsample the dataset, globally and per class.

data (Dataset | DatasetDict): Dataset to downsample.
max_ncells (int | None): Global maximum number of examples.
max_ncells_per_class (int | None): Max per-class examples after shuffling.
cell_state_dict (dict): Mapping from class ID to list of example indices.

subsample_by_class

from classifier_utils import subsample_by_class

data = subsample_by_class(
    data,
    cell_state_dict: dict,
    max_ncells_per_class: int,
)

Subsample examples per class according to max_ncells_per_class.

rename_cols

from classifier_utils import rename_cols

data = rename_cols(
    data,
    mapping: dict[str, str],
)

Rename dataset columns according to a mapping.

flatten_list

from classifier_utils import flatten_list

flat = flatten_list(nested_list)

Flatten a nested list of lists into a single list.

label_classes

from classifier_utils import label_classes

labeled_data, id_class_dict = label_classes(
    data,
    gene_class_dict: dict,
    label_mapping: dict | None,
    token_dictionary_file: str | None,
    mode: str = "spot",
)

Create per-token class labels for gene-level classification.

data (Dataset | DatasetDict):
Input tokenized dataset.
gene_class_dict (dict):
Mapping from class name → list of gene IDs/tokens.
label_mapping (dict | None):
Optional mapping from class name → integer ID. If None, it is created.
token_dictionary_file (str | None):
Path to token dictionary used to map gene IDs to token IDs.
mode (str, default "spot"):
"spot" for single sequences; "extended" for spot + neighbor, where only the first half is supervised.

Returns the labeled dataset and the id_class_dict mapping.

DataCollatorForCellClassification

from classifier_utils import DataCollatorForCellClassification

A custom data collator for sequence-level classification. It:

Pads input IDs and attention masks,
Collects labels as integer class IDs,
Returns a batch suitable for AutoModelForSequenceClassification.

You normally do not instantiate this directly; it is constructed inside Classifier.train.

DataCollatorForGeneClassification

from classifier_utils import DataCollatorForGeneClassification

A custom data collator for gene-level (token) classification. It:

Pads input_ids and labels,
Preserves -100 values to mark ignored positions,
Returns a batch suitable for AutoModelForTokenClassification.

Evaluation Utilities (`evaluation_utils.py`)

py_softmax

from evaluation_utils import py_softmax

probs = py_softmax(x, axis=-1)

Pure-Python softmax utility over a NumPy array.

compute_metrics

from evaluation_utils import compute_metrics

metrics = compute_metrics(eval_pred)

Metric function intended for use with Trainer. It typically computes:

overall accuracy,
macro/micro F1,
precision and recall,
confusion matrix and per-class stats (if labels are available).

It accepts the standard (predictions, labels) tuple used by Trainer.

evaluate_model

from evaluation_utils import evaluate_model

metrics = evaluate_model(
    model,
    dataloader,
    device,
)

Standalone evaluation helper that runs a model over a dataloader and aggregates metrics using compute_metrics.

plot_confusion_matrix

from evaluation_utils import plot_confusion_matrix

plot_confusion_matrix(
    conf_mat,
    class_order: list[str],
    output_directory: str,
    output_prefix: str,
)

Plot and save a confusion matrix heatmap.

plot_predictions

from evaluation_utils import plot_predictions

plot_predictions(
    predictions_file: str,
    id_class_dict_file: str,
    title: str,
    output_directory: str,
    output_prefix: str,
    class_order: list[str],
)

Load saved predictions and class mappings, reconstruct the confusion matrix, and save a labeled heatmap to disk.

Classification

Classifier

Parameters (init)

build_custom_tokenizer

Parameters

Returns

_load_tokenizer

prepare_data

Parameters

Returns

train

Behavior

Parameters

Returns

evaluate

Behavior

Parameters

Returns

plot_confusion_matrix

Parameters

plot_predictions

Parameters

Dataset Utilities (classifier_utils.py)

load_and_filter

remove_rare

downsample_and_shuffle

subsample_by_class

rename_cols

flatten_list

label_classes

DataCollatorForCellClassification

DataCollatorForGeneClassification

Evaluation Utilities (evaluation_utils.py)

py_softmax

compute_metrics

evaluate_model

plot_confusion_matrix

plot_predictions

Dataset Utilities (`classifier_utils.py`)

Evaluation Utilities (`evaluation_utils.py`)