{ "cells": [ { "cell_type": "markdown", "id": "300f96f9", "metadata": {}, "source": [ "# Pretrain and Create Model for Classification Based Tasks" ] }, { "cell_type": "code", "execution_count": 1, "id": "b2580ba1", "metadata": {}, "outputs": [], "source": [ "from stFormer.classifier.Classifier import Classifier" ] }, { "cell_type": "markdown", "id": "f0b1ec1f", "metadata": {}, "source": [ "## 1.1 Classify From Pretrained Model" ] }, { "cell_type": "markdown", "id": "31d4f877", "metadata": {}, "source": [ "We take out Subtype based information to evaluate classification fine-tuning and evaluation\n", "\n", "1. **Data Loading & Splitting** \n", " - Load `train_ds` from `dataset_path`. \n", " - If `eval_dataset_path` provided, load `eval_ds`; \n", " otherwise do a `train_test_split(test_size, seed=42)`.\n", "\n", "2. **`model_init` Function** \n", " - Loads base model & config from `model_checkpoint`. \n", " - Overrides `num_labels` to match `self.label_mapping`. \n", " - Optionally freezes the first `self.freeze_layers` encoder layers.\n", " - Adds a classification head onto BERT pretreained model if loading from masked learning objective\n", "\n", "3. **Tokenizer & Data Collator** \n", " - `AutoTokenizer.from_pretrained(...)` with `padding=\"max_length\"` \n", " - `DataCollatorWithPadding` to pad to `tokenizer.model_max_length`.\n", "\n", "4. **Classification**\n", " - `Evaluation metrics` compute metrics to determine training/test loss and accuracy\n", " - `training args` takes dictionary of BERT training arguments for hyperparameter selection and model updating\n", "\n", "5. **Best Checkpoint Selection and Saving**\n", " - Saves model checkpoints to output directory based upon ``eval strategy` \n", " - Returns final `trainer` model and saves final model to `output_directory`\n" ] }, { "cell_type": "code", "execution_count": null, "id": "1508f6cd", "metadata": {}, "outputs": [], "source": [ "classifier = Classifier(\n", " metadata_column = 'Tissue',\n", " mode='spot',\n", " classifier_type = 'sequence', #for class predictions\n", " token_dictionary_file='output/spot/token_dictionary.pickle',\n", " rare_threshold=0.1, #remove rare data types (less than 10% of samples)\n", " max_examples_per_class=10000, #option to downsample\n", " nproc=24,\n", ")\n", "ds_path, map_path = classifier.prepare_data(\n", " input_data = 'annotated.dataset/',\n", " output_directory = 'tmp/clasifier',\n", " output_prefix = 'Tissue_Classifier',\n", " )" ] }, { "cell_type": "markdown", "id": "9e4c9e22", "metadata": {}, "source": [ "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" ] }, { "cell_type": "code", "execution_count": null, "id": "07976680", "metadata": {}, "outputs": [ { "name": "stdout", "output_type": "stream", "text": [ "Using model checkpoint: output/spot/spot_model\n", "Number of labels from data: 3\n", "Label mapping: {'Brain': 0, 'Breast': 1, 'Skin': 2}\n", "Linear(in_features=256, out_features=3, bias=True)\n", "Max label: tensor(2)\n", "Label mapping size: 3\n" ] }, { "data": { "text/html": [ "\n", "
\n", " \n", " \n", " [1500/1500 22:31, Epoch 2/2]\n", "
\n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", "
StepTraining LossValidation LossAccuracyPrecisionRecallF1
1001.0837001.0011980.5393330.4165460.5369040.438580
2000.7867000.5697150.7968330.8231480.7961790.780779
3000.6895000.6481990.6871670.7298240.6865030.641001
4000.5867000.4232880.8688330.8783220.8689520.868393
5000.4903000.9762970.5816670.6112280.5793230.547443
6000.5053000.5248490.8016670.8232870.8007290.798036
7000.3652000.2498330.9408330.9427680.9409110.940833
8000.3300000.9084840.6676670.6940010.6674200.620839
9000.2918000.8307960.7061670.7400240.7048660.680967
10000.2480000.1508170.9671670.9673910.9671840.967184
11000.1859000.7323620.7546670.7839780.7549840.746594
12000.1937000.1441090.9688330.9691120.9688750.968854
13000.1652000.1368300.9671670.9672430.9671760.967176
14000.1388000.1693180.9585000.9592160.9585870.958527
15000.1617000.1401910.9698330.9701160.9698800.969856

" ], "text/plain": [ "" ] }, "metadata": {}, "output_type": "display_data" }, { "data": { "text/html": [], "text/plain": [ "" ] }, "metadata": {}, "output_type": "display_data" } ], "source": [ "trainer = classifier.train(\n", " model_checkpoint='output/spot/spot_model', # pretrained model path\n", " dataset_path = ds_path, # dataset path from prepare data\n", " output_directory = 'output/models/classification', #output evaluation \n", " test_size=0.2, # splits dataset into test/train splits\n", ")" ] }, { "cell_type": "markdown", "id": "3af33fda", "metadata": {}, "source": [ "## 1.2 Train Gene Classifier" ] }, { "cell_type": "code", "execution_count": 2, "id": "fa7b7680", "metadata": {}, "outputs": [], "source": [ "from stFormer.classifier.Classifier import Classifier\n", "import pandas as pd\n", "import pickle\n", "import numpy as np\n", "import os" ] }, { "cell_type": "markdown", "id": "b2c9cb72", "metadata": {}, "source": [ "We are replicating publication analysis by training a classifier to predict responsive genes in TNBC\n", "1. load genes upregulated in response to neoadjuvent care in TNBC\n", "2. load list of random shuffled genes as background\n", "2. load ensembl to gene_name mapping dictionary\n", "3. create dictionary for respnder and random genes\n", "4. Run gene classification for predictions of responsive genes in dataset" ] }, { "cell_type": "code", "execution_count": 4, "id": "d607b8bd", "metadata": {}, "outputs": [], "source": [ "os.chdir('analyses/models.to.test/Extended.model/')" ] }, { "cell_type": "code", "execution_count": 5, "id": "da264104", "metadata": {}, "outputs": [], "source": [ "file1 = \"upregulated.top300\"\n", "file2 = \"gene.shuffled.upregulated\"\n", "genes_responder = list(np.loadtxt(file1,dtype=str))\n", "genes_random = list(np.loadtxt(file2, dtype=str))\n", "\n", "training_args = {\"num_train_epochs\": 30.0, \"weight_decay\": 0.25, \"learning_rate\": 3e-6, \"warmup_steps\":1500, \"lr_scheduler_type\": \"polynomial\"}" ] }, { "cell_type": "markdown", "id": "fffd523a", "metadata": {}, "source": [ "### GeneClassifier Token Classification Overview\n", "\n", "We take out per-cell Subtype labels and instead classify **individual genes (tokens)** within each sequence.\n", "\n", "1. **Data Loading & Splitting** \n", " - Load `train_ds` from `dataset_path`. \n", " - If `eval_dataset_path` is provided, load `eval_ds`. \n", " Otherwise, perform `train_test_split(test_size, seed=42)`.\n", "\n", "2. **Label Mapping (Gene Classes)** \n", " - Use `classifier_utils.label_classes(\"gene\", ...)` to map each **input token (gene)** to a class label. \n", " - Generates a per-token `labels` field matching `input_ids` shape.\n", "\n", "3. **Tokenizer & Data Collator** \n", " - Uses `DataCollatorForGeneClassification` to pad both `input_ids` and `labels` in sync.\n", "\n", "4. **Model Initialization** \n", " - Loads base model & config from `model_checkpoint`. \n", " - Creates a `TokenClassification` head on the pretrained model.\n", "\n", "5. **Classification Training** \n", " - Computes token-level metrics (e.g., F1 score, accuracy). \n", "\n", "6. **Best Checkpoint Selection and Saving** \n", " - Saves model checkpoints to output directory based on `evaluation_strategy`. \n", " - Final model and tokenizer are saved to `output_directory`. \n", " - Predictions and evaluation metrics are returned for downstream analysis." ] }, { "cell_type": "code", "execution_count": 20, "id": "059a3181", "metadata": {}, "outputs": [], "source": [ "\n", "ray_config = {\"num_train_epochs\": [1.0,],\n", "\"learning_rate\": (1e-3, 1e-2),\n", "\"weight_decay\": (0.01, 0.05),\n", "\"lr_scheduler_type\": [\"linear\", \"cosine\", \"polynomial\"],\n", "\"warmup_steps\": (5, 50),\n", "\"seed\": (100, 1000),\n", "\"per_device_train_batch_size\": [10,],\n", "}\n" ] }, { "cell_type": "code", "execution_count": null, "id": "6760dd9b", "metadata": {}, "outputs": [], "source": [ "gene_class_dict = {'Responder': genes_responder,'Random.genes': genes_random}\n", "\n", "# 2) Instantiate for token-classification\n", "gene_classifier = Classifier(\n", " metadata_column=None, # no sequence-level label\n", " mode='extended',\n", " gene_class_dict=gene_class_dict, #specify this dictionary to use Gene Classifier\n", " classifier_type = 'gene',\n", " freeze_layers=4, # freeze first two BERT layers (optional)\n", " forward_batch_size=50,\n", " max_examples=10_000,\n", " nproc=16,\n", " token_dictionary_file='SpatialModel/new_token_dictionary.pickle'\n", ")\n", "\n", "# 3) Prepare your dataset (must already contain `input_ids` for each cell)\n", "ds_path, map_path = gene_classifier.prepare_data(\n", " input_data=\"STFormer_TNBC_neighbor.dataset\",\n", " output_directory=\"tmp/gene_classifier\",\n", " output_prefix=\"gene_classifier\"\n", ")\n" ] }, { "cell_type": "code", "execution_count": null, "id": "10fafed3", "metadata": {}, "outputs": [], "source": [ "trainer = gene_classifier.train(\n", " model_checkpoint=\"run-8eb93bdf/checkpoint-1000\",\n", " dataset_path=ds_path,\n", " output_directory=\"tmp/gene_classifier\",\n", ")" ] }, { "cell_type": "code", "execution_count": null, "id": "00b1bf8c", "metadata": {}, "outputs": [], "source": [ "metrics = gene_classifier.evaluate(\n", " model_directory=\"models/visium_gene_classifier/final_model\",\n", " eval_dataset_path=ds_path,\n", " id_class_dict_file=map_path,\n", " output_directory=\"output/gene_classifier\",\n", ")" ] }, { "cell_type": "markdown", "id": "d518c2b0", "metadata": {}, "source": [ "## 1.3 Train and Evaluate Model with Hyperparameter search" ] }, { "cell_type": "markdown", "id": "32b835a3", "metadata": {}, "source": [ "In this example we utilize ray configuration to loop through a list of hyperparameters to search for the best configuration of arguments for a classification task.\n", "\n", "\n", "Performs end-to-end hyperparameter search for a sequence-classification head using Ray Tune and Hugging Face Trainer.\n", "1. **Define Hyperparameter Search Space** \n", " - Pull ranges/choices from `self.ray_config` for \n", " `learning_rate`, `num_train_epochs`, `weight_decay`, etc. \n", "\n", "2. **CLI Reporter** \n", " - `CLIReporter` shows per-trial metrics (`eval_loss`, `eval_accuracy`) \n", " and hyperparameter values in the console.\n", "\n", "3. **Trainer & Hyperparameter Search** \n", " - Instantiate `Trainer` with `model_init`, datasets, collator, and `compute_metrics`. \n", " - Run `trainer.hyperparameter_search(...)` with Ray backend and `HyperOptSearch`.\n", "\n", "4. **Best Checkpoint Selection & Saving** \n", " - Use `ExperimentAnalysis` to find best trial/checkpoint by `eval_loss`. \n", " - Load that checkpoint into a fresh `BertForSequenceClassification`. \n", " - Save model & tokenizer under `output_directory/best_model`." ] }, { "cell_type": "markdown", "id": "15a14104", "metadata": {}, "source": [ "Load Datasets and Test/Train Split" ] }, { "cell_type": "code", "execution_count": 1, "id": "6166998d", "metadata": {}, "outputs": [], "source": [ "from stFormer.classifier.Classifier import Classifier\n", "import stFormer.classifier.classifier_utils as cu\n", "from datasets import load_from_disk\n", "import pandas as pd\n", "import numpy as np\n", "import random\n", "\n", "random.seed(123)" ] }, { "cell_type": "code", "execution_count": null, "id": "f8de24b0", "metadata": {}, "outputs": [], "source": [ "\n", "ds = load_from_disk('annotated.dataset')\n", "#ds_filt = cu.remove_rare(ds,rare_threshold=0.05,nproc=24,state_key='Tissue')\n", "\n", "train1 = pd.read_csv('data/train1.csv').dropna()\n", "test1 = pd.read_csv('data/test1.csv').dropna()\n", "train2 = pd.read_csv('data/train2.csv').dropna()\n", "test2 = pd.read_csv('data/test2.csv').dropna()\n", "\n", "train_samples = np.unique(train1['Sample'].tolist())\n", "test_samples = np.unique(test1['Sample'].tolist())\n", "ds_train = ds.filter(lambda ex: ex['Sample ID'] in train_samples,num_proc = 24)\n", "ds_test = ds.filter(lambda ex: ex['Sample ID'] in test_samples,num_proc=24)\n", "\n" ] }, { "cell_type": "markdown", "id": "a12b4bf6", "metadata": {}, "source": [ "Set up hyperparameters, classification information, and prepare dataset for classification\n", "1. For more hyperparameter options, please visit our docs: " ] }, { "cell_type": "code", "execution_count": null, "id": "a99095fa", "metadata": {}, "outputs": [], "source": [ "\n", "hyperparameters ={\n", " \"learning_rate\":[1e-5,1e-3],\n", " \"weight_decay\": [0.0, 0.3],\n", " \"warmup_ratio\": [0,0.3]\n", " #'lr_scheduler_type': [\"linear\",\"cosine\",\"polynomial\"], \n", " #'per_device_train_batch_size': [32]\n", " }\n", "\n", "classifier = Classifier(\n", " metadata_column = 'Tissue',\n", " mode='spot',\n", " ray_config = hyperparameters,\n", " token_dictionary_file='output/spot/token_dictionary.pickle',\n", " nproc=24,\n", ")" ] }, { "cell_type": "code", "execution_count": null, "id": "08f12111", "metadata": {}, "outputs": [], "source": [ "ds_path, map_path = classifier.prepare_data(\n", " input_data = ds_train, #takes Dataset Object or dataset file path\n", " output_directory = 'tmp_eval', #filtered dataset out location\n", " output_prefix = 'train_tissue' \n", ")\n", "\n", "# 2) Prepare the data exactly the same way\n", "eval_ds_path, eval_map_path = classifier.prepare_data(\n", " input_data = ds_test,\n", " output_directory = 'tmp_eval',\n", " output_prefix = 'eval_tissue')" ] }, { "cell_type": "markdown", "id": "8165e75a", "metadata": {}, "source": [ "Train the model" ] }, { "cell_type": "code", "execution_count": null, "id": "17bf5a7b", "metadata": {}, "outputs": [], "source": [ "trainer = classifier.train(\n", " model_checkpoint='output/models/tissue_classification/best_model', # pretrained model path\n", " dataset_path = ds_path, # dataset path from prepare data\n", " output_directory = 'output/eval_tissue_nohyperopt', #output evaluation \n", " eval_dataset = eval_ds_path,\n", " n_trials = 4\n", " )\n" ] }, { "cell_type": "markdown", "id": "8bab0f33", "metadata": {}, "source": [ "## 1.3 Plot Predictions using Evaluation Utils" ] }, { "cell_type": "markdown", "id": "1ba65199", "metadata": {}, "source": [ "Utilize seaborn, truth, and predicted values to create a confusion matrix and plot results" ] }, { "cell_type": "code", "execution_count": 1, "id": "d1f7744c", "metadata": {}, "outputs": [], "source": [ "from stFormer.classifier.Classifier import Classifier\n", "from datasets import load_from_disk\n", "from sklearn.metrics import confusion_matrix\n", "import pickle\n", "import os" ] }, { "cell_type": "code", "execution_count": null, "id": "2f3cd261", "metadata": {}, "outputs": [], "source": [ "\n", "#Produce & save raw predictions\n", "eval_ds = load_from_disk(ds_path).shuffle(seed=42).select(range(1000))\n", "preds = trainer.predict(eval_ds)\n", "y_true = preds.label_ids\n", "y_pred = preds.predictions.argmax(-1)\n", "\n", "with open(\"output/models/classification/predictions.pkl\", \"wb\") as f:\n", " pickle.dump({\"y_true\": y_true, \"y_pred\": y_pred}, f)" ] }, { "cell_type": "code", "execution_count": 13, "id": "ba574aa1", "metadata": {}, "outputs": [], "source": [ "import numpy as np\n", "with open('output/eval_tissue_nohyperopt/predictions.pkl','rb') as f:\n", " preds = pickle.load(f)\n", "y_true = preds.label_ids\n", "y_pred = preds.predictions.argmax(-1)\n", "\n", "map_path = 'tmp_eval/train_tissue_id_class_dict.pkl'" ] }, { "cell_type": "code", "execution_count": 14, "id": "99e3f11d", "metadata": {}, "outputs": [], "source": [ "\n", "# Load the id→class mapping you dumped in prepare_data()\n", "with open(map_path, \"rb\") as f:\n", " id_map = pickle.load(f) \n", "\n", "# We need a list of class names in label‐index order:\n", "inv_map = {v:k for k,v in id_map.items()}\n", "class_order = [inv_map[i] for i in range(len(inv_map))]\n", "\n", "cm = confusion_matrix(y_true, y_pred, labels=list(id_map.values()))" ] }, { "cell_type": "code", "execution_count": null, "id": "5d7a6169", "metadata": {}, "outputs": [], "source": [ "import matplotlib.pyplot as plt \n", "import seaborn as sns\n", "plt.figure(figsize=(8, 8))\n", "sns.heatmap(cm, annot=True, fmt=\"d\", xticklabels=class_order, yticklabels=class_order, cmap=\"Blues\")\n", "plt.xlabel(\"Predicted\")\n", "plt.ylabel(\"True\")\n", "plt.show()" ] }, { "cell_type": "code", "execution_count": null, "id": "c1bc9ee1", "metadata": {}, "outputs": [], "source": [ "#save heatmap with inbuilt plotting functionality\n", "classifier.plot_predictions(\n", " predictions_file=\"output/models/classification/predictions.pkl\",\n", " id_class_dict_file=map_path,\n", " title=\"Visium Spot Subtype Predictions\",\n", " output_directory=\"output/models/classification\",\n", " output_prefix=\"visium_spot\",\n", " class_order=class_order\n", ")" ] } ], "metadata": { "kernelspec": { "display_name": "stFormer", "language": "python", "name": "stformer" }, "language_info": { "codemirror_mode": { "name": "ipython", "version": 3 }, "file_extension": ".py", "mimetype": "text/x-python", "name": "python", "nbconvert_exporter": "python", "pygments_lexer": "ipython3", "version": "3.10.12" } }, "nbformat": 4, "nbformat_minor": 5 }