Pretrain and Create Model for Classification Based Tasks

from stFormer.classifier.Classifier import Classifier

1.1 Classify From Pretrained Model

We take out Subtype based information to evaluate classification fine-tuning and evaluation

  1. Data Loading & Splitting

    • Load train_ds from dataset_path.

    • If eval_dataset_path provided, load eval_ds;
      otherwise do a train_test_split(test_size, seed=42).

  2. model_init Function

    • Loads base model & config from model_checkpoint.

    • Overrides num_labels to match self.label_mapping.

    • Optionally freezes the first self.freeze_layers encoder layers.

    • Adds a classification head onto BERT pretreained model if loading from masked learning objective

  3. Tokenizer & Data Collator

    • AutoTokenizer.from_pretrained(...) with padding="max_length"

    • DataCollatorWithPadding to pad to tokenizer.model_max_length.

  4. Classification

    • Evaluation metrics compute metrics to determine training/test loss and accuracy

    • training args takes dictionary of BERT training arguments for hyperparameter selection and model updating

  5. Best Checkpoint Selection and Saving

    • Saves model checkpoints to output directory based upon ``eval strategy`

    • Returns final trainer model and saves final model to output_directory

classifier = Classifier(
    metadata_column = 'Tissue',
    mode='spot',
    classifier_type = 'sequence', #for class predictions
    token_dictionary_file='output/spot/token_dictionary.pickle',
    rare_threshold=0.1, #remove rare data types (less than 10% of samples)
    max_examples_per_class=10000, #option to downsample
    nproc=24,
)
ds_path, map_path = classifier.prepare_data(
    input_data = 'annotated.dataset/',
    output_directory = 'tmp/clasifier',
    output_prefix = 'Tissue_Classifier',
    )

In this example we utilize the model that was trained with a masked learning objective. While this is definitely possible, we suggest utilizing another Bert model that was trained using a classification task and then fine-tune on specific task

trainer = classifier.train(
    model_checkpoint='output/spot/spot_model', # pretrained model path
    dataset_path = ds_path, # dataset path from prepare data
    output_directory = 'output/models/classification', #output evaluation 
    test_size=0.2, # splits dataset into test/train splits
)

Using model checkpoint: output/spot/spot_model
Number of labels from data: 3
Label mapping: {'Brain': 0, 'Breast': 1, 'Skin': 2}
Linear(in_features=256, out_features=3, bias=True)
Max label: tensor(2)
Label mapping size: 3
[1500/1500 22:31, Epoch 2/2]
Step Training Loss Validation Loss Accuracy Precision Recall F1
100 1.083700 1.001198 0.539333 0.416546 0.536904 0.438580
200 0.786700 0.569715 0.796833 0.823148 0.796179 0.780779
300 0.689500 0.648199 0.687167 0.729824 0.686503 0.641001
400 0.586700 0.423288 0.868833 0.878322 0.868952 0.868393
500 0.490300 0.976297 0.581667 0.611228 0.579323 0.547443
600 0.505300 0.524849 0.801667 0.823287 0.800729 0.798036
700 0.365200 0.249833 0.940833 0.942768 0.940911 0.940833
800 0.330000 0.908484 0.667667 0.694001 0.667420 0.620839
900 0.291800 0.830796 0.706167 0.740024 0.704866 0.680967
1000 0.248000 0.150817 0.967167 0.967391 0.967184 0.967184
1100 0.185900 0.732362 0.754667 0.783978 0.754984 0.746594
1200 0.193700 0.144109 0.968833 0.969112 0.968875 0.968854
1300 0.165200 0.136830 0.967167 0.967243 0.967176 0.967176
1400 0.138800 0.169318 0.958500 0.959216 0.958587 0.958527
1500 0.161700 0.140191 0.969833 0.970116 0.969880 0.969856

1.2 Train Gene Classifier

from stFormer.classifier.Classifier import Classifier
import pandas as pd
import pickle
import numpy as np
import os

We are replicating publication analysis by training a classifier to predict responsive genes in TNBC

  1. load genes upregulated in response to neoadjuvent care in TNBC

  2. load list of random shuffled genes as background

  3. load ensembl to gene_name mapping dictionary

  4. create dictionary for respnder and random genes

  5. Run gene classification for predictions of responsive genes in dataset

os.chdir('analyses/models.to.test/Extended.model/')

file1 = "upregulated.top300"
file2 = "gene.shuffled.upregulated"
genes_responder = list(np.loadtxt(file1,dtype=str))
genes_random = list(np.loadtxt(file2, dtype=str))

training_args = {"num_train_epochs": 30.0, "weight_decay": 0.25, "learning_rate": 3e-6, "warmup_steps":1500, "lr_scheduler_type": "polynomial"}

GeneClassifier Token Classification Overview

We take out per-cell Subtype labels and instead classify individual genes (tokens) within each sequence.

  1. Data Loading & Splitting

    • Load train_ds from dataset_path.

    • If eval_dataset_path is provided, load eval_ds.
      Otherwise, perform train_test_split(test_size, seed=42).

  2. Label Mapping (Gene Classes)

    • Use classifier_utils.label_classes("gene", ...) to map each input token (gene) to a class label.

    • Generates a per-token labels field matching input_ids shape.

  3. Tokenizer & Data Collator

    • Uses DataCollatorForGeneClassification to pad both input_ids and labels in sync.

  4. Model Initialization

    • Loads base model & config from model_checkpoint.

    • Creates a TokenClassification head on the pretrained model.

  5. Classification Training

    • Computes token-level metrics (e.g., F1 score, accuracy).

  6. Best Checkpoint Selection and Saving

    • Saves model checkpoints to output directory based on evaluation_strategy.

    • Final model and tokenizer are saved to output_directory.

    • Predictions and evaluation metrics are returned for downstream analysis.

ray_config = {"num_train_epochs": [1.0,],
"learning_rate": (1e-3, 1e-2),
"weight_decay": (0.01, 0.05),
"lr_scheduler_type": ["linear", "cosine", "polynomial"],
"warmup_steps": (5, 50),
"seed": (100, 1000),
"per_device_train_batch_size": [10,],
}

gene_class_dict = {'Responder': genes_responder,'Random.genes': genes_random}

# 2) Instantiate for token-classification
gene_classifier = Classifier(
    metadata_column=None,             # no sequence-level label
    mode='extended',
    gene_class_dict=gene_class_dict, #specify this dictionary to use Gene Classifier
    classifier_type = 'gene',
    freeze_layers=4,                  # freeze first two BERT layers (optional)
    forward_batch_size=50,
    max_examples=10_000,
    nproc=16,
    token_dictionary_file='SpatialModel/new_token_dictionary.pickle'
)

# 3) Prepare your dataset (must already contain `input_ids` for each cell)
ds_path, map_path = gene_classifier.prepare_data(
    input_data="STFormer_TNBC_neighbor.dataset",
    output_directory="tmp/gene_classifier",
    output_prefix="gene_classifier"
)

trainer = gene_classifier.train(
    model_checkpoint="run-8eb93bdf/checkpoint-1000",
    dataset_path=ds_path,
    output_directory="tmp/gene_classifier",
)

metrics = gene_classifier.evaluate(
    model_directory="models/visium_gene_classifier/final_model",
    eval_dataset_path=ds_path,
    id_class_dict_file=map_path,
    output_directory="output/gene_classifier",
)

1.3 Plot Predictions using Evaluation Utils

Utilize seaborn, truth, and predicted values to create a confusion matrix and plot results

from stFormer.classifier.Classifier import Classifier
from datasets import load_from_disk
from sklearn.metrics import confusion_matrix
import pickle
import os

#Produce & save raw predictions
eval_ds = load_from_disk(ds_path).shuffle(seed=42).select(range(1000))
preds = trainer.predict(eval_ds)
y_true = preds.label_ids
y_pred = preds.predictions.argmax(-1)

with open("output/models/classification/predictions.pkl", "wb") as f:
    pickle.dump({"y_true": y_true, "y_pred": y_pred}, f)

import numpy as np
with open('output/eval_tissue_nohyperopt/predictions.pkl','rb') as f:
    preds = pickle.load(f)
y_true = preds.label_ids
y_pred = preds.predictions.argmax(-1)

map_path = 'tmp_eval/train_tissue_id_class_dict.pkl'

# Load the id→class mapping you dumped in prepare_data()
with open(map_path, "rb") as f:
    id_map = pickle.load(f)       

# We need a list of class names in label‐index order:
inv_map = {v:k for k,v in id_map.items()}
class_order = [inv_map[i] for i in range(len(inv_map))]

cm = confusion_matrix(y_true, y_pred, labels=list(id_map.values()))

import matplotlib.pyplot as plt 
import seaborn as sns
plt.figure(figsize=(8, 8))
sns.heatmap(cm, annot=True, fmt="d", xticklabels=class_order, yticklabels=class_order, cmap="Blues")
plt.xlabel("Predicted")
plt.ylabel("True")
plt.show()

#save heatmap with inbuilt plotting functionality
classifier.plot_predictions(
    predictions_file="output/models/classification/predictions.pkl",
    id_class_dict_file=map_path,
    title="Visium Spot Subtype Predictions",
    output_directory="output/models/classification",
    output_prefix="visium_spot",
    class_order=class_order
)