FTPM: Fine-Tuning Pre-trained Model to advance precision prognosis

A Cox survival analysis framework built on Geneformer for cancer prognosis prediction using gene expression data

Status
Activity
Community
Downloads
License
Platform
Contribute

---

FTPM fine-tunes Geneformer to predict cancer survival from gene-expression tokens.

The project provides:

• Geneformer Cox models for single-cancer and pan-cancer analysis;
• 2048-token and 4096-token model variants;
• five-fold out-of-fold (OOF) risk prediction;
• forest plots, Kaplan-Meier plots, and C-index summaries combining FTPM risk, clinical variables, and tumor mutational burden (TMB);
• an MCP server for running the workflows with an AI assistant.

Scripts

Model training

Script	Scope	Tokens	Pooling
FTPM_2048_brca.py	BRCA	2048	Mean
FTPM_4096_brca.py	BRCA	4096	Attention
FTPM_2048_pan-caner.py	Pan-cancer	2048	Mean
FTPM_4096_pan-caner.py	Pan-cancer	4096	Attention
FTPM_cancer_4096_brca.py	BRCA with early stopping	4096	Attention
FTPM_cancer_4096_pan-caner.py	Pan-cancer with early stopping	4096	Attention

Pan-cancer training supports:

BRCA KIRC LUSC LUAD OV GBM LAML PRAD THCA HNSC
COAD LGG STAD LIHC SKCM BLCA CESC SARC ESCA PAAD

Downstream survival analysis

Script	Cox models per cancer	Output
forest_CI_p_FTPM_TMB_clinic.py	3	48 forest plots, 16 KM plots, 3 CSV tables
forest_KM_ROC_CI_p_FTPM_TMB_clinic.py	4	64 forest plots, 16 KM plots, 3 CSV tables

Both scripts analyze 16 cancers:

BRCA KIRC LUSC LUAD OV GBM LAML PRAD
THCA HNSC COAD LGG STAD LIHC SKCM BLCA

The three common Cox models are:

baseline + TMB + FTPM
TMB + FTPM
FTPM only

The four-model script additionally fits:

baseline + FTPM, without TMB

FTPM is the categorical High/Low group derived from the OOF risk score. The historical second filename contains ROC, but that script does not calculate ROC curves or AUC.

Installation

Training requires a CUDA GPU. The 2048-token models typically require at least 40 GB VRAM, while the 4096-token models typically require at least 80 GB.

pip install torch transformers datasets scikit-learn scikit-survival geneformer
pip install numpy pandas scipy lifelines matplotlib

Input data

Training data

Gene-expression token files have no header. The training scripts transpose the loaded matrix, so input columns correspond to samples and rows correspond to token positions. Clinical files contain the event indicator, sample ID, and survival time in days.

Pan-cancer data are arranged as:

data_root/
├── TCGA-BRCA/
│   ├── BRCA_expr_token_4096_without_colnames.txt
│   └── BRCA_sur_clinic_lnc_gene_without_colnames.txt
├── TCGA-LUAD/
│   └── ...
└── ...

Forest/KM analysis data

risk/TCGA-BRCA/OOF_risk.txt
risk/TCGA-BRCA/OOF_sample_ids.txt
tmb/TCGA-BRCA.tmb.csv

Requirements:

• OOF_risk.txt: one numeric risk score per row;
• OOF_sample_ids.txt: one matching sample ID per row, in the same order;
• clinical/TMB CSV columns: Tumor_Sample_Barcode, time, event, age_group, gender, and TMB;
• time is expected in months by the plotting scripts.

Usage

Run commands from the repository root.

Train a 4096-token BRCA model

python scripts/FTPM_4096_brca.py \
  --expr_file /path/to/BRCA_expr_token_4096_without_colnames.txt \
  --clinic_file /path/to/BRCA_sur_clinic.txt \
  --pretrained_model_dir /path/to/geneformer \
  --out results/brca

Train selected pan-cancer models

python scripts/FTPM_cancer_4096_pan-caner.py \
  --data_root /path/to/data_root \
  --pretrained_model_dir /path/to/geneformer \
  --out results/pan_cancer \
  --cancer_types BRCA,LUAD,KIRC

Run the three-model forest/KM analysis

python scripts/forest_CI_p_FTPM_TMB_clinic.py \
  --score-template 'risk/TCGA-{cancer}/OOF_risk.txt' \
  --sid-template 'risk/TCGA-{cancer}/OOF_sample_ids.txt' \
  --meta-template 'tmb/{cancer}.tmb.csv' \
  --fig-dir 'results/forest_ci/forest' \
  --km-dir 'results/forest_ci/km' \
  --tbl-dir 'results/forest_ci/tables'

Run the four-model forest/KM analysis

python scripts/forest_KM_ROC_CI_p_FTPM_TMB_clinic.py \
  --score-template 'risk/TCGA-{cancer}/OOF_risk.txt' \
  --sid-template 'risk/TCGA-{cancer}/OOF_sample_ids.txt' \
  --meta-template 'tmb/{cancer}.tmb.csv' \
  --fig-dir 'results/forest_km/forest' \
  --km-dir 'results/forest_km/km' \
  --tbl-dir 'results/forest_km/tables'

Use different output directories because the two analysis scripts share some output filenames.

The CSV tables are:

table1.CI_p_FTPM.csv
table1.CI_p_TMB_FTPM.csv
table1.CI_p_baseline_TMB_FTPM.csv

MCP server

Install and register the MCP server:

cd ftpm_vibe
pip install -e .
claude mcp add ftmp -- python -m ftmp_mcp

Main MCP tools:

Tool	Purpose
check_ftmp_installation	Check Python, GPU, dependencies, and scripts
run_ftmp_cox_2048_brca	Train the 2048-token BRCA model
run_ftmp_cox_4096_brca	Train the 4096-token BRCA model
run_ftmp_cox_2048_pan_cancer	Train 2048-token pan-cancer models
run_ftmp_cox_4096_pan_cancer	Train 4096-token pan-cancer models
run_ftmp_cox_cancer_4096_brca	Train BRCA with early stopping
run_ftmp_cox_cancer_4096_pan_cancer	Train pan-cancer models with early stopping
run_ftmp_forest_ci_analysis	Run the three-model downstream analysis
run_ftmp_forest_km_clinic_analysis	Run the four-model downstream analysis
search_ftmp_docs	Search this README

Existing ftmp_* package and tool names are retained for backward compatibility.

Results

Always report OOF metrics as model performance. Metrics calculated from a final model evaluated on its own training data are apparent performance and may be optimistic.

Citation

If you use FTPM, please cite the associated publication when available.

License

MIT