llm-cancer-survival

Survival analysis pipeline that combines clinical baseline features with LLM-extracted features and LLM embeddings from medical reports. Models are trained using Random Survival Forests with k-fold cross-validation, evaluated via concordance index and Integrated Brier Score.

Project Structure

src/
  utils.py          - Shared utilities: subdataset creation, PCA, RSF training,
                      SHAP calculation, Brier score / calibration helpers
  train.py          - Training pipeline configuration (feature combinations,
                      entity definitions) and CLI entry point

notebooks/
  train.ipynb       - Interactive notebook to train RSF models and view SHAP plots
  llm_extract.ipynb - LLM-based feature extraction and embedding generation
                      from free-text medical reports

mock_data/          - Mock datasets for NSCLC and colon cancer
results/            - Training outputs (c-index, risk scores, SHAP, calibration)

Setup

This project requires two separate Python environments due to incompatible numpy version requirements between scikit-survival/shap and transformers.

Training Environment

Used for RSF model training, SHAP explanations, and calibration (notebooks/train.ipynb, src/train.py).

conda create -n rsfshap python=3.9
conda activate rsfshap
pip install -r requirements.txt

Extraction Environment

Used for LLM-based feature extraction and embedding generation (notebooks/llm_extract.ipynb).

conda create -n llm-extract python=3.11
conda activate llm-extract
pip install -r requirements_extraction.txt

Usage

Feature Extraction

Use notebooks/llm_extract.ipynb to extract structured clinical features and embeddings from medical reports using an LLM.

Model Training

Either run the notebook interactively:

notebooks/train.ipynb

Or run from the command line:

cd src
python train.py

Results (c-index per fold, risk scores, SHAP values, calibration data) are saved to results/<entity>/.