Prepare PRISM backend for Hugging Face Spaces

.dockerignore

@@ -0,0 +1,25 @@
+.git
+.venv
+__pycache__
+**/__pycache__
+*.pyc
+
+data/*.sqlite3
+reports
+evaluation_results
+notebooks
+
+frontend/node_modules
+frontend/dist
+
+models/saved/*.pth
+models/saved/*.pt
+models/saved/*.ckpt
+models/saved/*.bin
+models/saved/*.safetensors
+models/saved/*.onnx
+models/saved/training_runs
+
+docs
+*.log
+.cache

.gitattributes

@@ -33,3 +33,9 @@ saved_model/**/* filter=lfs diff=lfs merge=lfs -text
 *.zip filter=lfs diff=lfs merge=lfs -text
 *.zst filter=lfs diff=lfs merge=lfs -text
 *tfevents* filter=lfs diff=lfs merge=lfs -text
+medical_docs/mds-clinical-diagnostic-criteria-for-parkinson's-disease.pdf filter=lfs diff=lfs merge=lfs -text
+medical_docs/mds-research-criteria-for-prodromal-parkinson's-disease.pdf filter=lfs diff=lfs merge=lfs -text
+medical_docs/MDS-UPDRS.pdf filter=lfs diff=lfs merge=lfs -text
+medical_docs/Movement[[:space:]]Disorders[[:space:]]-[[:space:]]2019[[:space:]]-[[:space:]]Heinzel[[:space:]]-[[:space:]]Update[[:space:]]of[[:space:]]the[[:space:]]MDS[[:space:]]research[[:space:]]criteria[[:space:]]for[[:space:]]prodromal[[:space:]]Parkinson[[:space:]]s[[:space:]]disease.pdf filter=lfs diff=lfs merge=lfs -text
+medical_docs/parkinsons-disease-in-adults-pdf-1837629189061.pdf filter=lfs diff=lfs merge=lfs -text
+medical_docs/PPD.pdf filter=lfs diff=lfs merge=lfs -text

Dockerfile

@@ -0,0 +1,24 @@
+FROM python:3.11-slim
+
+ENV PYTHONDONTWRITEBYTECODE=1 \
+    PYTHONUNBUFFERED=1 \
+    PORT=7860 \
+    PD_ALLOWED_ORIGINS=* \
+    PD_LOAD_TRANSFORMERS=0 \
+    MPLBACKEND=Agg
+
+WORKDIR /app
+
+RUN apt-get update \
+    && apt-get install -y --no-install-recommends libgomp1 \
+    && rm -rf /var/lib/apt/lists/*
+
+COPY requirements-hf.txt ./requirements-hf.txt
+RUN pip install --no-cache-dir --upgrade pip \
+    && pip install --no-cache-dir -r requirements-hf.txt
+
+COPY . .
+
+EXPOSE 7860
+
+CMD ["python", "start_server.py", "--skip-init"]

README.md

@@ -1,10 +1,177 @@
 ---
-title: Prism Backend
-emoji: 📊
-colorFrom: yellow
-colorTo: yellow
+title: PRISM Parkinsons Assessment API
 sdk: docker
-pinned: false
+app_port: 7860
 ---
 
-Check out the configuration reference at https://huggingface.co/docs/hub/spaces-config-reference
+# Parkinson's Disease Assessment Portal
+
+A comprehensive medical assessment system for Parkinson's disease diagnosis using the PPMI (Parkinson's Progression Markers Initiative) curated dataset. The system combines traditional machine learning, medical transformer models, and RAG-based report generation.
+
+## Dataset
+
+Uses the **PPMI Curated Data Cut** CSV files containing patient data with the following key features:
+
+| Category | Features |
+|----------|----------|
+| Demographics | Age, Sex, Education Years, Race, BMI |
+| Family History | Family PD history (categorical + binary) |
+| Motor Symptoms | Tremor, Rigidity, Bradykinesia, Postural Instability |
+| Non-Motor | REM sleep, Epworth Sleepiness, Depression (GDS), Anxiety (STAI) |
+| Cognitive | MoCA, Clock Drawing, Benton JLO |
+
+**Target classes (COHORT):** HC (Healthy Control) · PD (Parkinson's Disease) · SWEDD · PRODROMAL
+
+## Architecture
+
+```
+├── src/
+│   ├── data_preprocessing.py      # Patient-level leak-free data pipeline
+│   ├── web_interface.py            # Flask web app
+│   ├── rag_system.py               # Medical knowledge base + report generation
+│   ├── document_manager.py         # PDF/text document indexing (TF-IDF)
+│   ├── feature_mapping.py          # Patient questionnaire ↔ PPMI feature mapping
+│   ├── analyze_data.py             # Dataset EDA script
+│   ├── train_traditional_models.py # Train LightGBM, XGBoost, SVM
+│   ├── train_transformer_models.py # Train PubMedBERT, BioGPT, Clinical-T5
+│   ├── train_multimodal.py         # Train multimodal ensemble
+│   ├── evaluate_traditional_models.py
+│   └── models/
+│       ├── traditional_ml.py       # LightGBM, XGBoost, SVM wrappers
+│       ├── transformer_models.py   # DistilBERT, BioBERT, PubMedBERT for tabular
+│       ├── medical_transformers.py # PubMedBERT, BioGPT, Clinical-T5 classifiers
+│       └── multimodal_ml.py        # Stacking ensemble
+├── templates/                      # Flask HTML templates
+├── static/                         # CSS, JS assets
+├── medical_docs/                   # Medical literature for RAG
+├── models/saved/                   # Trained model weights
+├── start_server.py                 # Entry point for web app
+└── requirements.txt
+```
+
+## Features
+
+- **Leak-Free Preprocessing**: Patient-level train/test split ensures no patient appears in both sets
+- **Traditional ML**: LightGBM, XGBoost, SVM with class weight balancing
+- **Medical Transformers**: PubMedBERT (encoder), BioGPT (decoder), Clinical-T5 (encoder-decoder)
+- **Multimodal Ensemble**: Stacking ensemble combining all model predictions
+- **RAG-Enhanced Reports**: Retrieves medical literature to generate comprehensive diagnostic reports
+- **Web Interface**: Patient assessment form with automated report generation and PDF export
+
+## Installation
+
+```bash
+# Create virtual environment
+python -m venv venv
+source venv/bin/activate    # Linux/Mac
+venv\Scripts\activate       # Windows
+
+# Install dependencies
+pip install -r requirements.txt
+
+# Install PyTorch with CUDA (recommended for GPU training)
+pip install torch --index-url https://download.pytorch.org/whl/cu124
+```
+
+`requirements.txt` now includes `sacremoses`, which BioGPT needs for tokenization. If the A4000 preflight fails on `sacremoses`, rerun `pip install -r requirements.txt` inside the training venv.
+
+## RTX A4000 Setup
+
+For the RTX A4000 training machine, use this flow from the project root.
+
+Ubuntu:
+
+```bash
+source venv/bin/activate
+bash check_a4000_ready.sh
+bash train_a4000_models.sh
+```
+
+Windows:
+
+```bat
+venv\Scripts\activate
+python check_a4000_ready.py
+train_a4000_models.bat
+```
+
+What the preflight checks:
+
+- CUDA-enabled PyTorch import and `torch.cuda.is_available()`
+- detected GPU name, CUDA version, and VRAM
+- BioGPT tokenizer dependency (`sacremoses`)
+- required PPMI CSV files
+- `medical_docs/` availability for RAG training
+- free disk space and output path write access
+
+Helper scripts:
+
+- `check_a4000_ready.sh` / `check_a4000_ready.bat` run the GPU/data preflight
+- `train_a4000_models.sh` / `train_a4000_models.bat` run preflight, then start training with `--gpu-profile rtx-a4000`
+- `resume_a4000_training.sh` / `resume_a4000_training.bat` resume the same run if the session is interrupted
+
+The A4000 training recipe now defaults to class-weighted focal loss and keeps the best transformer checkpoint by validation F1, with validation loss used only as a tie-breaker.
+
+Recommended direct commands:
+
+```bash
+python src/train_model_suite.py train --run-name a4000_full --gpu-profile rtx-a4000 --epochs 30 --patience 8 --traditional-trials 6 --transformer-trials 6 --transformer-loss focal --focal-gamma 1.5
+python src/train_model_suite.py resume --run-name a4000_full --gpu-profile rtx-a4000 --epochs 30 --patience 8 --traditional-trials 6 --transformer-trials 6 --transformer-loss focal --focal-gamma 1.5
+python src/train_model_suite.py status --run-name a4000_full
+```
+
+## Usage
+
+### Train Models
+
+```bash
+cd src
+
+# Train traditional ML models
+python train_traditional_models.py
+
+# Train transformer models (requires GPU recommended)
+python train_transformer_models.py
+
+# Train multimodal ensemble
+python train_multimodal.py
+```
+
+For the full resumable training pipeline with the A4000 profile, run from the project root instead of `src/`:
+
+```bash
+bash train_a4000_models.sh
+```
+
+### Run Web App
+
+```bash
+# From project root
+python start_server.py
+# Access at http://localhost:5000
+```
+
+### Evaluate Models
+
+```bash
+cd src
+python evaluate_traditional_models.py
+```
+
+## Model Performance
+
+Models are evaluated on a held-out test set using patient-level splitting:
+
+| Model | Type |
+|-------|------|
+| LightGBM | Gradient Boosting |
+| XGBoost | Gradient Boosting |
+| SVM (RBF) | Support Vector Machine |
+| PubMedBERT | Encoder-only Transformer |
+| BioGPT | Decoder-only Transformer |
+| Clinical-T5 | Encoder-Decoder Transformer |
+| Multimodal Ensemble | Stacking (all above) |
+
+## License
+
+This project uses data from the [Parkinson's Progression Markers Initiative (PPMI)](https://www.ppmi-info.org/).

medical_docs/MDS-UPDRS.pdf

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:9da07b4d53f785bc89dab47d74b24cfb434ead17bfbfb1efe7988a169710bcc1
+size 1386246

medical_docs/Movement Disorders - 2019 - Heinzel - Update of the MDS research criteria for prodromal Parkinson s disease.pdf

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:170de7732faa3e982862248c0f47d9a51996de3eaeb3c579c6e0f7f93f408efb
+size 137288

medical_docs/PPD.pdf

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:9cc5a74b061065ab475701df7c2316f4b192e3018e258db5f854f42ff578c831
+size 1791069

medical_docs/differential_diagnosis.txt

@@ -0,0 +1,33 @@
+Title: Differential Diagnosis of Parkinsonian Syndromes
+Author: Dr. Sarah Chen
+Date: 2023-08-10
+
+Abstract:
+This paper discusses the differential diagnosis of Parkinson's Disease (PD) and other parkinsonian syndromes, providing clinical guidelines for accurate diagnosis.
+
+Key Points:
+1. Parkinson's Disease vs. Essential Tremor:
+   - PD tremor: Typically resting, asymmetric, 4-6 Hz frequency
+   - Essential tremor: Primarily postural/action tremor, often symmetric, 8-12 Hz frequency
+   - Essential tremor patients typically lack bradykinesia and rigidity
+   - Family history more common in essential tremor
+
+2. Parkinson's Disease vs. SWEDD:
+   - SWEDD (Scans Without Evidence of Dopaminergic Deficit) patients present with parkinsonian symptoms
+   - Normal dopaminergic imaging distinguishes SWEDD from PD
+   - SWEDD patients often show dystonic features and lack true bradykinesia
+   - SWEDD patients typically do not respond well to levodopa therapy
+
+3. Parkinson's Disease vs. Atypical Parkinsonism:
+   - Multiple System Atrophy (MSA): Early autonomic failure, cerebellar signs, poor levodopa response
+   - Progressive Supranuclear Palsy (PSP): Early falls, vertical gaze palsy, axial rigidity
+   - Corticobasal Degeneration (CBD): Asymmetric apraxia, alien limb phenomenon, cortical sensory loss
+
+4. Prodromal Parkinson's Disease:
+   - REM sleep behavior disorder: 80% convert to synucleinopathy within 10-15 years
+   - Hyposmia: Reduced sense of smell often precedes motor symptoms by years
+   - Constipation: Common non-motor symptom that may appear decades before diagnosis
+   - Subtle motor changes: Reduced arm swing, changes in handwriting, voice softening
+
+Conclusion:
+Accurate differential diagnosis requires careful clinical assessment, consideration of response to levodopa, neuroimaging, and monitoring of disease progression. Early non-motor symptoms may help identify prodromal PD cases before classic motor features emerge.

medical_docs/mds-clinical-diagnostic-criteria-for-parkinson's-disease.pdf

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:16fa3730d16168654a1ee5097c21c295cacd4d3e2189f76449ba06676e6ff2ef
+size 201526

medical_docs/mds-research-criteria-for-prodromal-parkinson's-disease.pdf

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:5b5f966f634a4150f5252eecd132963cf44047d7ba41aa32143ac47b56bfaa3e
+size 281649

medical_docs/parkinsons-disease-in-adults-pdf-1837629189061.pdf

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:5679922ecf82e77df4b5c206661b3dfb61082e5aacb00bb43aa0a9926cfbd9c0
+size 167734

medical_docs/parkinsons_research.txt

@@ -0,0 +1,20 @@
+Title: Recent Advances in Parkinson's Disease Research
+Author: Dr. James Wilson
+Date: 2023-05-15
+
+Abstract:
+This paper reviews recent advances in Parkinson's Disease (PD) research, focusing on early detection, biomarkers, and treatment options.
+
+Key Findings:
+1. Early Detection: Tremor, rigidity, bradykinesia, and postural instability remain the cardinal symptoms of PD. However, non-motor symptoms like REM sleep behavior disorder, hyposmia, and constipation often precede motor symptoms by years and may serve as early indicators.
+
+2. Biomarkers: Alpha-synuclein aggregation in peripheral tissues shows promise as a biomarker. Neuroimaging techniques including DaTscan can help differentiate PD from other movement disorders.
+
+3. Genetic Factors: Mutations in LRRK2, GBA, and SNCA genes are associated with increased PD risk. Family history remains a significant risk factor, with first-degree relatives having 2-3 times higher risk.
+
+4. Treatment Approaches: Levodopa remains the gold standard for symptomatic treatment. Deep brain stimulation shows efficacy for motor symptoms in advanced cases. Emerging therapies targeting alpha-synuclein are in clinical trials.
+
+5. Differential Diagnosis: SWEDD (Scans Without Evidence of Dopaminergic Deficit) patients present with parkinsonian symptoms but normal dopaminergic imaging. Essential tremor differs from PD tremor by being primarily postural rather than resting.
+
+Conclusion:
+Early diagnosis and personalized treatment approaches based on symptom profile, genetic factors, and biomarkers represent the future of PD management. Multidisciplinary care remains essential for optimal patient outcomes.

models/saved/lightgbm_model.joblib

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:83d39246aac8638f4df77e1d4b66ebf65470d03710c5faa68aa8aa38398a22b6
+size 27464708

models/saved/multimodal_ensemble.joblib

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:7186412962046f02ecf6b67fdc425291e16101f0c67193c7cd0622febadbb659
+size 661938

models/saved/svm_model.joblib

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:e63ad40596edab8b6afdde418533c6a5737bae475c630b90d8e54ad435ddb9ff
+size 1545535

models/saved/traditional_class_mapping.json

@@ -0,0 +1,6 @@
+{
+  "1": 0,
+  "2": 1,
+  "3": 2,
+  "4": 3
+}

models/saved/traditional_preprocessor.joblib

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:8756a573be06ce68e104fb59d0f9458e0316b5b73f6da122af0cd7fffb3cb0a8
+size 3871426

models/saved/xgboost_model.joblib

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:a4a1e1dd62168b450c05fc4f47f9861bfb664051aff34ab38c35dc35ab096029
+size 2100189

requirements-hf.txt

@@ -0,0 +1,18 @@
+# Hugging Face Spaces runtime dependencies.
+# This keeps the free CPU container lean by using traditional-model inference.
+pandas>=2.1.0
+numpy>=1.24.3
+scipy>=1.11.0
+scikit-learn>=1.3.0
+lightgbm>=4.1.0
+xgboost>=2.0.0
+torch>=2.1.0
+Flask>=2.3.0
+flask-cors>=4.0.0
+Werkzeug>=2.3.0
+PyPDF2>=3.0.0
+matplotlib>=3.7.2
+seaborn>=0.12.2
+reportlab>=4.0.7
+joblib>=1.3.0
+tqdm>=4.66.0

src/__init__.py

@@ -0,0 +1,3 @@
+"""
+Source package for the Parkinson's Disease Assessment System.
+"""

src/analyze_data.py

@@ -0,0 +1,111 @@
+from pathlib import Path
+
+import matplotlib.pyplot as plt
+import pandas as pd
+import seaborn as sns
+
+from data_preprocessing import DataPreprocessor
+
+
+def analyze_dataset():
+    root_dir = Path(__file__).resolve().parents[1]
+    dataset_path = root_dir / "PPMI_Curated_Data_Cut_Public_20250714.csv"
+    notebooks_dir = root_dir / "notebooks"
+    notebooks_dir.mkdir(exist_ok=True)
+
+    if not dataset_path.exists():
+        raise FileNotFoundError(f"Dataset not found: {dataset_path}")
+
+    df = pd.read_csv(dataset_path)
+    print("\nDataset Shape:", df.shape)
+
+    if "COHORT" not in df.columns:
+        raise ValueError("Dataset is missing required column: COHORT")
+
+    print("\nCOHORT Distribution:")
+    cohort_dist = df["COHORT"].value_counts(dropna=False)
+    print(cohort_dist)
+
+    preprocessor = DataPreprocessor()
+    selected_features = list(preprocessor.core.selected_features)
+
+    analysis_features = [
+        col
+        for col in selected_features
+        if col in df.columns and col not in {"COHORT", "PATNO"}
+    ]
+
+    print(f"\nUsing {len(analysis_features)} available analysis features:")
+    print(analysis_features)
+
+    if not analysis_features:
+        raise ValueError("No expected analysis features were found in the dataset.")
+
+    print("\nKey Features Statistics:")
+    print(df[analysis_features].describe(include="all").transpose())
+
+    print("\nMissing Values Analysis:")
+    missing = df[analysis_features].isnull().sum()
+    missing_pct = (missing / len(df)) * 100
+    missing_summary = pd.DataFrame(
+        {
+            "Missing Values": missing,
+            "Percentage": missing_pct,
+        }
+    ).sort_values("Percentage", ascending=False)
+    print(missing_summary[missing_summary["Missing Values"] > 0])
+
+    plt.figure(figsize=(10, 6))
+    sns.countplot(data=df, x="COHORT")
+    plt.title("Distribution of Cohorts")
+    plt.tight_layout()
+    plt.savefig(notebooks_dir / "cohort_distribution.png")
+    plt.close()
+
+    numerical_features = (
+        df[analysis_features].select_dtypes(include=["number"]).columns.tolist()
+    )
+
+    if numerical_features:
+        plt.figure(figsize=(12, 8))
+        correlation_matrix = df[numerical_features].corr(numeric_only=True)
+        sns.heatmap(correlation_matrix, annot=False, cmap="coolwarm", center=0)
+        plt.title("Feature Correlation Matrix")
+        plt.tight_layout()
+        plt.savefig(notebooks_dir / "correlation_matrix.png")
+        plt.close()
+
+        n_features = len(numerical_features)
+        n_cols = 3
+        n_rows = (n_features + n_cols - 1) // n_cols
+
+        plt.figure(figsize=(15, 5 * n_rows))
+        for i, feature in enumerate(numerical_features, 1):
+            plt.subplot(n_rows, n_cols, i)
+            sns.histplot(data=df, x=feature, kde=True)
+            plt.title(feature)
+        plt.tight_layout()
+        plt.savefig(notebooks_dir / "feature_distributions.png")
+        plt.close()
+    else:
+        print("\nNo numerical features available for correlation/distribution plots.")
+
+    print("\nFeature Types:")
+    for feature in analysis_features:
+        dtype = df[feature].dtype
+        n_unique = df[feature].nunique(dropna=True)
+        print(f"- {feature}: {dtype}, {n_unique} unique values")
+
+    missing_expected = [
+        col
+        for col in selected_features
+        if col not in df.columns and col not in {"COHORT", "PATNO"}
+    ]
+    if missing_expected:
+        print("\nExpected features not found in dataset:")
+        for feature in missing_expected:
+            print(f"- {feature}")
+
+
+if __name__ == "__main__":
+    analyze_dataset()

src/data_preprocessing.py

@@ -0,0 +1,294 @@
+import os
+
+import pandas as pd
+import numpy as np
+
+from sklearn.model_selection import GroupShuffleSplit, train_test_split
+from sklearn.preprocessing import StandardScaler, OneHotEncoder
+from sklearn.impute import SimpleImputer, KNNImputer
+from sklearn.compose import ColumnTransformer
+from sklearn.pipeline import Pipeline
+
+
+class PPMIDataPreprocessor:
+    """
+    Clean, leak-proof preprocessing pipeline for PPMI dataset.
+    Handles patient-level splitting, imputation, encoding, scaling.
+    """
+
+    def __init__(self):
+
+        # -------------------------
+        # Selected Features
+        # -------------------------
+        self.selected_features = [
+            # Demographics
+            "age", "SEX", "EDUCYRS", "race", "BMI",
+
+            # Family History
+            "fampd", "fampd_bin",
+
+            # Motor symptoms
+            "sym_tremor", "sym_rigid", "sym_brady", "sym_posins",
+
+            # Non-motor
+            "rem", "ess", "gds", "stai",
+
+            # Cognitive
+            "moca", "clockdraw", "bjlot",
+            
+            # --- EXTENDED MEDICAL FEATURES ---
+            # Olfactory
+            "upsit", "upsit_pctl",
+            
+            # UPDRS (Gold standard for PD)
+            "updrs1_score", "updrs2_score", "updrs3_score", "updrs4_score", "updrs_totscore",
+            
+            # DatScan (Imaging)
+            "mean_caudate", "mean_putamen", 
+            
+            # Biomarkers (CSF)
+            "abeta", "tau", "ptau",
+
+            # Target
+            "COHORT",
+
+            # Patient ID
+            "PATNO"
+        ]
+
+        # Features to impute using KNN (to preserve correlations)
+        self.knn_cols = [
+            "moca", "clockdraw", "bjlot", "rem", "gds", "ess", "stai",
+            "upsit", "upsit_pctl",
+            "updrs1_score", "updrs2_score", "updrs3_score", "updrs4_score", "updrs_totscore",
+            "mean_caudate", "mean_putamen",
+            "abeta", "tau", "ptau"
+        ]
+
+        # Categorical
+        self.categorical_cols = ["SEX", "race", "fampd"]
+
+        # Numeric
+        self.numeric_cols = [
+            "age", "EDUCYRS", "BMI",
+            "fampd_bin",
+            "sym_tremor", "sym_rigid", "sym_brady", "sym_posins",
+            
+            # New Medical Features (added for enhancement)
+            "updrs1_score", "updrs2_score", "updrs3_score", "updrs4_score", "updrs_totscore",
+            "mean_caudate", "mean_putamen",
+            "abeta", "tau", "ptau",
+            "upsit", "upsit_pctl"
+        ]
+
+        self.preprocessor = None
+
+    # ------------------------------------------------------------
+    def _clean_biomarker(self, x):
+        """Clean biomarker strings like '>1700' or '<200' to floats."""
+        if pd.isna(x):
+            return np.nan
+        if isinstance(x, str):
+            x = x.replace(">", "").replace("<", "")
+        try:
+            return float(x)
+        except ValueError:
+            return np.nan
+
+    def load_data(self, file_path):
+        """Load and filter only necessary columns."""
+        df = pd.read_csv(file_path, low_memory=False)
+        
+        # Clean biomarker columns if they exist
+        for col in ["abeta", "tau", "ptau"]:
+            if col in df.columns:
+                df[col] = df[col].apply(self._clean_biomarker)
+
+        # Select available features (intersection with file columns)
+        available = [c for c in self.selected_features if c in df.columns]
+        # Always require COHORT/PATNO
+        if "COHORT" not in available or "PATNO" not in available:
+            raise ValueError("Dataset missing required COHORT or PATNO columns")
+            
+        df = df[available].dropna(subset=["COHORT", "PATNO"])
+        return df
+
+    # ------------------------------------------------------------
+    def patient_split(self, df, test_size=0.2):
+        """Split such that the same patient never appears in both sets."""
+        gss = GroupShuffleSplit(test_size=test_size, n_splits=1, random_state=42)
+        train_idx, test_idx = next(gss.split(df, groups=df["PATNO"]))
+
+        train_df = df.iloc[train_idx].reset_index(drop=True)
+        test_df = df.iloc[test_idx].reset_index(drop=True)
+
+        return train_df, test_df
+
+    # ------------------------------------------------------------
+    def build_preprocessor(self):
+        """Builds leak-proof transformers."""
+
+        numeric_transformer = Pipeline(steps=[
+            ("imputer_mean", SimpleImputer(strategy="mean")),
+            ("scaler", StandardScaler())
+        ])
+
+        knn_transformer = Pipeline(steps=[
+            ("imputer_knn", KNNImputer(n_neighbors=5)),
+            ("scaler", StandardScaler())
+        ])
+
+        categorical_transformer = Pipeline(steps=[
+            ("imputer_mode", SimpleImputer(strategy="most_frequent")),
+            ("encoder", OneHotEncoder(handle_unknown="ignore"))
+        ])
+
+        self.preprocessor = ColumnTransformer(
+            transformers=[
+                ("num", numeric_transformer,
+                 [c for c in self.numeric_cols if c not in self.knn_cols]),
+
+                ("knn", knn_transformer, self.knn_cols),
+
+                ("cat", categorical_transformer, self.categorical_cols),
+            ],
+            remainder="drop"
+        )
+
+        return self.preprocessor
+
+    # ------------------------------------------------------------
+    def prepare(self, file_path):
+        """
+        COMPLETE DATA PREPARATION PIPELINE
+        Returns X_train, X_test, y_train, y_test
+        """
+
+        df = self.load_data(file_path)
+        train_df, test_df = self.patient_split(df)
+
+        X_train = train_df.drop(["COHORT", "PATNO"], axis=1)
+        y_train = train_df["COHORT"]
+
+        X_test = test_df.drop(["COHORT", "PATNO"], axis=1)
+        y_test = test_df["COHORT"]
+
+        # Build processor & fit only on training data
+        pre = self.build_preprocessor()
+        X_train_processed = pre.fit_transform(X_train)
+        X_test_processed = pre.transform(X_test)
+
+        return X_train_processed, X_test_processed, y_train.values, y_test.values
+
+
+class DataPreprocessor:
+    """Backwards-compatible wrapper around the new PPMI pipeline."""
+
+    def __init__(self):
+        self.core = PPMIDataPreprocessor()
+        self.feature_names_ = None
+        self.class_mapping_ = None
+        self.preprocessor_ = None
+        self.train_df_ = None
+        self.test_df_ = None
+
+    @staticmethod
+    def _to_python_scalar(value):
+        if isinstance(value, np.generic):
+            return value.item()
+        return value
+
+    def _load_all_files(self, file_paths):
+        if isinstance(file_paths, str):
+            file_paths = [file_paths]
+
+        dataframes = []
+        for path in file_paths:
+            if not path:
+                continue
+            if not os.path.exists(path):
+                print(f"[WARN] DataPreprocessor: '{path}' not found, skipping.")
+                continue
+            df = self.core.load_data(path)
+            dataframes.append(df)
+
+        if not dataframes:
+            raise FileNotFoundError("No valid CSV files were provided to DataPreprocessor.")
+
+        combined = pd.concat(dataframes, ignore_index=True)
+        combined = combined.drop_duplicates().reset_index(drop=True)
+        return combined
+
+    def prepare_data(self, file_paths, test_size=0.2, use_patient_split=True):
+        """Expose the legacy API expected by the training scripts."""
+
+        df = self._load_all_files(file_paths)
+
+        if use_patient_split:
+            train_df, test_df = self.core.patient_split(df, test_size=test_size)
+        else:
+            train_df, test_df = train_test_split(
+                df,
+                test_size=test_size,
+                random_state=42,
+                stratify=df["COHORT"],
+            )
+
+        self.train_df_ = train_df.reset_index(drop=True)
+        self.test_df_ = test_df.reset_index(drop=True)
+
+        classes_sorted = np.sort(df["COHORT"].unique())
+        self.class_mapping_ = {
+            self._to_python_scalar(original): int(idx)
+            for idx, original in enumerate(classes_sorted)
+        }
+
+        X_train = train_df.drop(["COHORT", "PATNO"], axis=1)
+        y_train = train_df["COHORT"].map(self.class_mapping_).values
+        X_test = test_df.drop(["COHORT", "PATNO"], axis=1)
+        y_test = test_df["COHORT"].map(self.class_mapping_).values
+
+        pre = self.core.build_preprocessor()
+        X_train_processed = pre.fit_transform(X_train)
+        X_test_processed = pre.transform(X_test)
+        self.preprocessor_ = pre
+
+        try:
+            self.feature_names_ = pre.get_feature_names_out(X_train.columns).tolist()
+        except AttributeError:
+            self.feature_names_ = None
+
+        return X_train_processed, X_test_processed, y_train, y_test
+
+    def get_feature_names(self):
+        if self.feature_names_ is None:
+            raise ValueError("Feature names are unavailable. Call prepare_data() first.")
+        return self.feature_names_
+
+    def get_preprocessor(self):
+        if self.preprocessor_ is None:
+            raise ValueError("Preprocessor is unavailable. Call prepare_data() first.")
+        return self.preprocessor_
+
+    def get_class_mapping(self):
+        if self.class_mapping_ is None:
+            raise ValueError("Class mapping is unavailable. Call prepare_data() first.")
+        return self.class_mapping_
+
+    def get_split_frames(self):
+        if self.train_df_ is None or self.test_df_ is None:
+            raise ValueError("Split frames are unavailable. Call prepare_data() first.")
+        return self.train_df_.copy(), self.test_df_.copy()
+
+
+# -------------------------------------------------------------------
+# Script example (not executed when imported)
+# -------------------------------------------------------------------
+if __name__ == "__main__":
+    file_path = "PPMI_Curated_Data.csv"
+    prep = PPMIDataPreprocessor()
+    X_train, X_test, y_train, y_test = prep.prepare(file_path)
+
+    print("Train:", X_train.shape)
+    print("Test:", X_test.shape)

src/document_manager.py

@@ -0,0 +1,455 @@
+"""
+Document Manager for RAG System
+Handles loading, processing, and retrieving medical documents for the RAG system
+"""
+
+import os
+import re
+import json
+import numpy as np
+import pandas as pd
+from typing import Dict, List, Tuple, Optional, Union
+from pathlib import Path
+import pickle
+import shutil
+from sklearn.feature_extraction.text import TfidfVectorizer
+from sklearn.metrics.pairwise import cosine_similarity
+import PyPDF2  # For handling PDF files
+
+class DocumentManager:
+    """Manages medical documents for the RAG system."""
+    
+    def __init__(self, docs_dir: str = "medical_docs"):
+        """Initialize the document manager.
+        
+        Args:
+            docs_dir: Directory where medical documents are stored
+        """
+        self.docs_dir = Path(docs_dir)
+        self.documents = {}
+        self.document_embeddings = {}
+        self.vectorizer = TfidfVectorizer(stop_words='english')
+        
+        # Create docs directory if it doesn't exist
+        os.makedirs(self.docs_dir, exist_ok=True)
+        
+        # For backward compatibility, also check for files directly in the docs_dir
+        self.main_dir = self.docs_dir
+        
+        # Create subdirectories for different document types
+        self.papers_dir = self.docs_dir / "papers"
+        self.guidelines_dir = self.docs_dir / "guidelines"
+        self.textbooks_dir = self.docs_dir / "textbooks"
+        
+        os.makedirs(self.papers_dir, exist_ok=True)
+        os.makedirs(self.guidelines_dir, exist_ok=True)
+        os.makedirs(self.textbooks_dir, exist_ok=True)
+        
+        # Load existing documents if any
+        self.load_documents()
+
+    def _build_document_entry(
+        self,
+        doc_id: str,
+        doc_type: str,
+        content: str,
+        file_path: Path,
+        metadata: Optional[Dict] = None,
+    ) -> Dict:
+        """Create a normalized in-memory representation for a document."""
+        resolved_path = Path(file_path)
+        size_bytes = None
+        try:
+            size_bytes = resolved_path.stat().st_size
+        except OSError:
+            pass
+
+        return {
+            "id": doc_id,
+            "type": doc_type,
+            "content": content,
+            "metadata": metadata or self._extract_metadata(content),
+            "file_path": str(resolved_path),
+            "size_bytes": size_bytes,
+        }
+    
+    def load_documents(self) -> None:
+        """Load all documents from the docs directory."""
+        self.documents = {}
+        
+        # Load documents from each subdirectory
+        for doc_type, directory in [
+            ("paper", self.papers_dir),
+            ("guideline", self.guidelines_dir),
+            ("textbook", self.textbooks_dir)
+        ]:
+            # Load text files
+            for file_path in directory.glob("*.txt"):
+                doc_id = f"{doc_type}_{file_path.stem}"
+                with open(file_path, 'r', encoding='utf-8') as f:
+                    content = f.read()
+
+                self.documents[doc_id] = self._build_document_entry(
+                    doc_id, doc_type, content, file_path
+                )
+            
+            # Load PDF files
+            for file_path in directory.glob("*.pdf"):
+                doc_id = f"{doc_type}_{file_path.stem}"
+                content = self._extract_text_from_pdf(file_path)
+
+                self.documents[doc_id] = self._build_document_entry(
+                    doc_id, doc_type, content, file_path
+                )
+        
+        # Also check for files directly in the main directory
+        # Text files
+        for file_path in self.main_dir.glob("*.txt"):
+            doc_id = f"document_{file_path.stem}"
+            with open(file_path, 'r', encoding='utf-8') as f:
+                content = f.read()
+
+            self.documents[doc_id] = self._build_document_entry(
+                doc_id, "paper", content, file_path
+            )
+            
+        # PDF files
+        for file_path in self.main_dir.glob("*.pdf"):
+            doc_id = f"document_{file_path.stem}"
+            content = self._extract_text_from_pdf(file_path)
+
+            self.documents[doc_id] = self._build_document_entry(
+                doc_id, "paper", content, file_path
+            )
+        
+        # Create document embeddings
+        self._create_embeddings()
+        
+        # Count document types
+        doc_counts = {'paper': 0, 'guideline': 0, 'textbook': 0, 'total': len(self.documents)}
+        for doc in self.documents.values():
+            doc_type = doc.get('type', 'unknown')
+            if doc_type in doc_counts:
+                doc_counts[doc_type] += 1
+        
+        print(f"Loaded {doc_counts} medical documents")
+    
+    def _extract_text_from_pdf(self, file_path: Path) -> str:
+        """Extract text content from a PDF file."""
+        try:
+            text = ""
+            with open(file_path, 'rb') as file:
+                pdf_reader = PyPDF2.PdfReader(file)
+                for page_num in range(len(pdf_reader.pages)):
+                    page = pdf_reader.pages[page_num]
+                    text += page.extract_text() + "\n\n"
+            return text
+        except Exception as e:
+            print(f"Error extracting text from PDF {file_path}: {e}")
+            return f"Error extracting text: {str(e)}"
+    
+    def _extract_metadata(self, content: str) -> Dict:
+        """Extract metadata from document content."""
+        metadata = {
+            "title": "",
+            "authors": "",
+            "year": "",
+            "source": "",
+            "keywords": []
+        }
+        
+        # Try to extract metadata from the first 20 lines
+        lines = content.split('\n')[:20]
+        for line in lines:
+            if line.lower().startswith("title:"):
+                metadata["title"] = line.split(":", 1)[1].strip()
+            elif line.lower().startswith("author") or line.lower().startswith("authors"):
+                metadata["authors"] = line.split(":", 1)[1].strip()
+            elif line.lower().startswith("year:"):
+                metadata["year"] = line.split(":", 1)[1].strip()
+            elif line.lower().startswith("source:"):
+                metadata["source"] = line.split(":", 1)[1].strip()
+            elif line.lower().startswith("keywords:"):
+                keywords = line.split(":", 1)[1].strip()
+                metadata["keywords"] = [k.strip() for k in keywords.split(",")]
+        
+        return metadata
+    
+    def _create_embeddings(self) -> None:
+        """Create TF-IDF embeddings for all documents."""
+        self.document_embeddings = {}
+        if not self.documents:
+            return
+        
+        # Extract document contents
+        doc_ids = list(self.documents.keys())
+        doc_contents = [self.documents[doc_id]["content"] for doc_id in doc_ids]
+        
+        # Create TF-IDF matrix
+        tfidf_matrix = self.vectorizer.fit_transform(doc_contents)
+        
+        # Store embeddings
+        for i, doc_id in enumerate(doc_ids):
+            self.document_embeddings[doc_id] = tfidf_matrix[i]
+    
+    def add_document(self, file_path: str, doc_type: str = "paper", title: str = None, author: str = None) -> str:
+        """Add a new document to the collection.
+
+        Args:
+            file_path: Path to the document file
+            doc_type: Type of document (paper, guideline, textbook)
+            title: Optional title metadata
+            author: Optional author metadata
+
+        Returns:
+            Document ID of the added document
+        """
+        file_path = Path(file_path)
+
+        if not file_path.exists():
+            raise FileNotFoundError(f"Document not found: {file_path}")
+
+        # Determine target directory
+        if doc_type == "paper":
+            target_dir = self.papers_dir
+        elif doc_type == "guideline":
+            target_dir = self.guidelines_dir
+        elif doc_type == "textbook":
+            target_dir = self.textbooks_dir
+        else:
+            doc_type = "paper"
+            target_dir = self.papers_dir
+
+        target_path = target_dir / file_path.name
+
+        # Copy file as binary-safe (works for PDF/text)
+        if file_path.resolve() != target_path.resolve():
+            shutil.copy2(file_path, target_path)
+
+        # Read content based on file extension
+        if target_path.suffix.lower() == '.pdf':
+            content = self._extract_text_from_pdf(target_path)
+        else:
+            with open(target_path, 'r', encoding='utf-8', errors='ignore') as f:
+                content = f.read()
+
+        metadata = self._extract_metadata(content)
+        if title:
+            metadata['title'] = title
+        if author:
+            metadata['authors'] = author
+
+        doc_id = f"{doc_type}_{target_path.stem}"
+
+        self.documents[doc_id] = self._build_document_entry(
+            doc_id,
+            doc_type,
+            content,
+            target_path,
+            metadata=metadata,
+        )
+
+        self._create_embeddings()
+        return doc_id
+
+    def remove_document(self, doc_id: str) -> bool:
+        """Remove a document from the collection.
+        
+        Args:
+            doc_id: ID of the document to remove
+            
+        Returns:
+            True if document was removed, False otherwise
+        """
+        if doc_id not in self.documents:
+            return False
+        
+        # Get file path
+        file_path = self.documents[doc_id]["file_path"]
+        
+        # Remove file
+        try:
+            os.remove(file_path)
+        except Exception as e:
+            print(f"Error removing file {file_path}: {e}")
+        
+        # Remove from documents collection
+        del self.documents[doc_id]
+        
+        # Remove from embeddings
+        if doc_id in self.document_embeddings:
+            del self.document_embeddings[doc_id]
+
+        if self.documents:
+            self._create_embeddings()
+        else:
+            self.document_embeddings = {}
+        
+        return True
+    
+    def search_documents(self, query: str, top_k: int = 5) -> List[Dict]:
+        """Search for documents matching the query.
+        
+        Args:
+            query: Search query
+            top_k: Number of top results to return
+            
+        Returns:
+            List of matching documents with relevance scores
+        """
+        if not self.documents:
+            return []
+        
+        # Create query embedding
+        query_embedding = self.vectorizer.transform([query])
+        
+        # Calculate similarity scores
+        scores = {}
+        for doc_id, doc_embedding in self.document_embeddings.items():
+            similarity = cosine_similarity(query_embedding, doc_embedding)[0][0]
+            scores[doc_id] = similarity
+        
+        # Sort by similarity score
+        sorted_docs = sorted(scores.items(), key=lambda x: x[1], reverse=True)
+        
+        # Return top-k results
+        results = []
+        for doc_id, score in sorted_docs[:top_k]:
+            doc = self.documents[doc_id].copy()
+            doc["relevance"] = float(score)
+            results.append(doc)
+        
+        return results
+    
+    def get_document_by_id(self, doc_id: str) -> Optional[Dict]:
+        """Get a document by its ID.
+        
+        Args:
+            doc_id: ID of the document
+            
+        Returns:
+            Document dict or None if not found
+        """
+        return self.documents.get(doc_id)
+    
+    def _serialize_document(
+        self,
+        doc: Dict,
+        include_content: bool = True,
+        preview_length: int = 180,
+    ) -> Dict:
+        """Return a client-facing representation of a document."""
+        data = doc.copy()
+        metadata = dict(data.get("metadata", {}))
+        if not metadata.get("title"):
+            metadata["title"] = Path(data.get("file_path", "")).stem
+        data["metadata"] = metadata
+        data["title"] = metadata.get("title", "")
+        data["author"] = metadata.get("authors", "")
+
+        content = data.get("content", "") or ""
+        data["preview"] = (
+            f"{content[:preview_length]}..." if len(content) > preview_length else content
+        )
+
+        if not include_content:
+            data.pop("content", None)
+
+        return data
+
+    def get_all_documents(self, include_content: bool = True, preview_length: int = 180) -> List[Dict]:
+        """Compatibility helper expected by web_interface.py."""
+        out = []
+        for doc in self.documents.values():
+            out.append(
+                self._serialize_document(
+                    doc,
+                    include_content=include_content,
+                    preview_length=preview_length,
+                )
+            )
+        return out
+
+    def get_document(self, doc_id: str) -> Optional[Dict]:
+        """Compatibility helper expected by web_interface.py."""
+        doc = self.get_document_by_id(doc_id)
+        if doc is None:
+            return None
+        return self._serialize_document(doc, include_content=True)
+
+    def get_document_summary(self, doc_id: str, preview_length: int = 180) -> Optional[Dict]:
+        """Return a lighter-weight document representation for list views."""
+        doc = self.get_document_by_id(doc_id)
+        if doc is None:
+            return None
+        return self._serialize_document(doc, include_content=False, preview_length=preview_length)
+
+    def get_document_count(self) -> Dict[str, int]:
+        """Get count of documents by type.
+        
+        Returns:
+            Dict with counts by document type
+        """
+        counts = {
+            "paper": 0,
+            "guideline": 0,
+            "textbook": 0,
+            "total": len(self.documents)
+        }
+        
+        for doc in self.documents.values():
+            counts[doc["type"]] += 1
+        
+        return counts
+    
+    def extract_relevant_passages(self, query: str, top_k: int = 3, 
+                                 passage_length: int = 500) -> List[Dict]:
+        """Extract relevant passages from documents for a query.
+        
+        Args:
+            query: Search query
+            top_k: Number of top passages to return
+            passage_length: Approximate length of each passage
+            
+        Returns:
+            List of relevant passages with metadata
+        """
+        # First, get relevant documents
+        relevant_docs = self.search_documents(query, top_k=top_k)
+        
+        passages = []
+        for doc in relevant_docs:
+            content = doc["content"]
+            
+            # Split content into paragraphs
+            paragraphs = re.split(r'\n\s*\n', content)
+            
+            # Create passages by combining paragraphs
+            current_passage = ""
+            for para in paragraphs:
+                if len(current_passage) + len(para) <= passage_length:
+                    current_passage += para + "\n\n"
+                else:
+                    # Add current passage to results
+                    if current_passage:
+                        passages.append({
+                            "text": current_passage.strip(),
+                            "doc_id": doc["id"],
+                            "doc_title": doc["metadata"]["title"],
+                            "relevance": doc["relevance"]
+                        })
+                    current_passage = para + "\n\n"
+            
+            # Add final passage
+            if current_passage:
+                passages.append({
+                    "text": current_passage.strip(),
+                    "doc_id": doc["id"],
+                    "doc_title": doc["metadata"]["title"],
+                    "relevance": doc["relevance"]
+                })
+        
+        # Sort passages by relevance
+        passages.sort(key=lambda x: x["relevance"], reverse=True)
+        
+        return passages[:top_k]

src/dual_report_generator.py

@@ -0,0 +1,112 @@
+"""
+Dual report generator for patient-friendly and clinician-facing outputs.
+"""
+
+from __future__ import annotations
+
+from dataclasses import dataclass
+from typing import Dict, Any
+from datetime import datetime
+
+
+CLASS_NAMES = ['Healthy Control', "Parkinson's Disease", 'SWEDD', 'Prodromal PD']
+
+
+@dataclass
+class PatientReportGenerator:
+    def generate_report(self, prediction_results: Dict[str, Any], patient_data: Dict[str, Any]) -> str:
+        pred_idx = int(prediction_results.get('ensemble_prediction', 0))
+        pred_label = CLASS_NAMES[pred_idx] if 0 <= pred_idx < len(CLASS_NAMES) else 'Unknown'
+        confidence = float(prediction_results.get('confidence', 0.0)) * 100.0
+
+        lines = [
+            'PATIENT SUMMARY REPORT',
+            '=' * 40,
+            f'Date: {datetime.now().strftime("%Y-%m-%d %H:%M:%S")}',
+            '',
+            f'Primary assessment: {pred_label}',
+            f'Confidence: {confidence:.1f}%',
+            '',
+            'What this means:',
+            '- This is an AI-assisted assessment from questionnaire/clinical inputs.',
+            '- It is NOT a final diagnosis.',
+            '- Please consult a neurologist for clinical confirmation.',
+            '',
+            'Next steps:',
+            '- Book specialist follow-up',
+            '- Track symptom progression',
+            '- Bring this report to your clinician',
+            '',
+            'Disclaimer: Educational/support use only. Not medical advice.'
+        ]
+        return '\n'.join(lines)
+
+
+@dataclass
+class DoctorReportGenerator:
+    def generate_report(self, prediction_results: Dict[str, Any], patient_data: Dict[str, Any], literature_insights: str = '') -> str:
+        pred_idx = int(prediction_results.get('ensemble_prediction', 0))
+        pred_label = CLASS_NAMES[pred_idx] if 0 <= pred_idx < len(CLASS_NAMES) else 'Unknown'
+        confidence = float(prediction_results.get('confidence', 0.0)) * 100.0
+        probs = prediction_results.get('ensemble_probabilities', [0, 0, 0, 0])
+
+        lines = [
+            'CLINICAL REPORT',
+            '=' * 40,
+            f'Date: {datetime.now().strftime("%Y-%m-%d %H:%M:%S")}',
+            '',
+            f'Predicted class: {pred_label}',
+            f'Model confidence: {confidence:.2f}%',
+            'Probability distribution:',
+            f"- HC: {float(probs[0]) * 100:.2f}%" if len(probs) > 0 else '- HC: N/A',
+            f"- PD: {float(probs[1]) * 100:.2f}%" if len(probs) > 1 else '- PD: N/A',
+            f"- SWEDD: {float(probs[2]) * 100:.2f}%" if len(probs) > 2 else '- SWEDD: N/A',
+            f"- PRODROMAL: {float(probs[3]) * 100:.2f}%" if len(probs) > 3 else '- PRODROMAL: N/A',
+            '',
+            'Patient features submitted:',
+        ]
+
+        for k, v in patient_data.items():
+            lines.append(f'- {k}: {v}')
+
+        if literature_insights:
+            lines.extend(['', 'Literature insights:', literature_insights])
+
+        lines.extend(['', 'Disclaimer: Decision support only; correlate clinically.'])
+        return '\n'.join(lines)
+
+
+class DualReportManager:
+    def __init__(self):
+        self.patient_generator = PatientReportGenerator()
+        self.doctor_generator = DoctorReportGenerator()
+
+    def generate_both_reports(self, prediction_results: Dict[str, Any], patient_data: Dict[str, Any], literature_insights: str = '') -> Dict[str, str]:
+        return {
+            'patient_report': self.patient_generator.generate_report(prediction_results, patient_data),
+            'doctor_report': self.doctor_generator.generate_report(prediction_results, patient_data, literature_insights),
+        }
+
+    def save_reports(self, reports: Dict[str, str], report_dir: str, patient_id: str) -> Dict[str, str]:
+        import os
+        from pathlib import Path
+        os.makedirs(report_dir, exist_ok=True)
+
+        safe_patient_id = Path(str(patient_id or "patient")).stem
+        safe_patient_id = "".join(
+            ch if ch.isalnum() or ch in "._- " else "_" for ch in safe_patient_id
+        ).strip(" .") or "patient"
+
+        patient_report_path = os.path.join(report_dir, f'patient_report_{safe_patient_id}.txt')
+        doctor_report_path = os.path.join(report_dir, f'clinical_report_{safe_patient_id}.txt')
+
+        with open(patient_report_path, 'w', encoding='utf-8') as f:
+            f.write(reports.get('patient_report', ''))
+
+        with open(doctor_report_path, 'w', encoding='utf-8') as f:
+            f.write(reports.get('doctor_report', ''))
+
+        return {
+            'patient_report_path': patient_report_path,
+            'doctor_report_path': doctor_report_path,
+        }

src/evaluate_traditional_models.py

@@ -0,0 +1,187 @@
+"""Re-train and evaluate traditional models on the leak-free patient split."""
+import os
+import sys
+from pathlib import Path
+import json
+
+import joblib
+import numpy as np
+import pandas as pd
+from sklearn.metrics import (
+    accuracy_score,
+    classification_report,
+    confusion_matrix,
+    precision_recall_fscore_support,
+    roc_curve,
+    auc
+)
+from sklearn.preprocessing import label_binarize
+from sklearn.utils.class_weight import compute_class_weight
+from lightgbm import LGBMClassifier
+from xgboost import XGBClassifier
+from sklearn.svm import SVC
+
+sys.path.append(str(Path(__file__).parent))
+from data_preprocessing import DataPreprocessor  # type: ignore
+
+ROOT = Path(__file__).resolve().parents[1]
+EVAL_DIR = ROOT / "evaluation_results" / "model_metrics"
+CLASS_REPORT_DIR = EVAL_DIR / "classification_reports"
+CONF_MATRIX_DIR = EVAL_DIR / "confusion_matrices"
+PLOTS_DIR = EVAL_DIR / "plots"
+ROC_DIR = EVAL_DIR / "roc_curves"
+
+for path in (CLASS_REPORT_DIR, CONF_MATRIX_DIR, PLOTS_DIR, ROC_DIR):
+    path.mkdir(parents=True, exist_ok=True)
+
+FILE_PATHS = [
+    ROOT / "PPMI_Curated_Data_Cut_Public_20240129.csv",
+    ROOT / "PPMI_Curated_Data_Cut_Public_20241211.csv",
+    ROOT / "PPMI_Curated_Data_Cut_Public_20250321.csv",
+    ROOT / "PPMI_Curated_Data_Cut_Public_20250714.csv",
+]
+CLASS_NAMES = ["HC", "PD", "SWEDD", "PRODROMAL"]
+
+
+def load_or_create_split():
+    split_path = ROOT / "evaluation_results" / "leak_free_split.npz"
+    meta_path = ROOT / "evaluation_results" / "leak_free_split_meta.joblib"
+    if split_path.exists() and meta_path.exists():
+        split = np.load(split_path)
+        meta = joblib.load(meta_path)
+        return split, meta
+
+    preprocessor = DataPreprocessor()
+    X_train, X_test, y_train, y_test = preprocessor.prepare_data(
+        FILE_PATHS,
+        test_size=0.2,
+        use_patient_split=True,
+    )
+    split_path.parent.mkdir(parents=True, exist_ok=True)
+    np.savez(split_path, X_train=X_train, X_test=X_test, y_train=y_train, y_test=y_test)
+    joblib.dump(
+        {
+            "feature_names": preprocessor.get_feature_names(),
+            "class_mapping": preprocessor.get_class_mapping(),
+        },
+        meta_path,
+    )
+    return np.load(split_path), joblib.load(meta_path)
+
+
+def train_models(X_train, y_train, class_weight_dict):
+    models = {}
+
+    lgb_params = dict(
+        random_state=42,
+        objective="multiclass",
+        num_class=len(CLASS_NAMES),
+        n_estimators=400,
+        learning_rate=0.03,
+        max_depth=7,
+        class_weight=class_weight_dict,
+    )
+    models["LightGBM"] = LGBMClassifier(**lgb_params)
+
+    xgb_params = dict(
+        random_state=42,
+        objective="multi:softmax",
+        num_class=len(CLASS_NAMES),
+        n_estimators=300,
+        learning_rate=0.05,
+        max_depth=6,
+        subsample=0.9,
+        colsample_bytree=0.9,
+        eval_metric="mlogloss",
+    )
+    models["XGBoost"] = XGBClassifier(**xgb_params)
+
+    models["SVM"] = SVC(
+        random_state=42,
+        probability=True,
+        kernel="rbf",
+        C=8.0,
+        gamma="scale",
+        class_weight=class_weight_dict,
+    )
+
+    for name, model in models.items():
+        model.fit(X_train, y_train)
+    return models
+
+
+def evaluate_model(name, model, X_test, y_test):
+    y_pred = model.predict(X_test)
+    y_prob = model.predict_proba(X_test)
+
+    accuracy = accuracy_score(y_test, y_pred)
+    precision, recall, f1, _ = precision_recall_fscore_support(
+        y_test, y_pred, average="weighted", zero_division=0
+    )
+    report = classification_report(
+        y_test, y_pred, target_names=CLASS_NAMES, zero_division=0
+    )
+    cm = confusion_matrix(y_test, y_pred)
+
+    # Save classification report
+    (CLASS_REPORT_DIR / f"{name}.txt").write_text(
+        f"{name} Classification Report (leak-free split)\n" + "-" * 60 + "\n" + report
+    )
+
+    # Save confusion matrix csv
+    cm_df = pd.DataFrame(cm, index=CLASS_NAMES, columns=CLASS_NAMES)
+    cm_df.to_csv(CONF_MATRIX_DIR / f"{name}_confusion_matrix.csv")
+
+    # Save ROC curves
+    y_bin = label_binarize(y_test, classes=range(len(CLASS_NAMES)))
+    roc_data = []
+    for idx, class_name in enumerate(CLASS_NAMES):
+        fpr, tpr, _ = roc_curve(y_bin[:, idx], y_prob[:, idx])
+        roc_auc = auc(fpr, tpr)
+        roc_df = pd.DataFrame({"fpr": fpr, "tpr": tpr})
+        roc_df.to_csv(ROC_DIR / f"{name}_class_{class_name}_roc.csv", index=False)
+        roc_data.append({"class": class_name, "auc": roc_auc})
+
+    return {
+        "model": name,
+        "accuracy": accuracy,
+        "precision": precision,
+        "recall": recall,
+        "f1": f1,
+    }
+
+
+def main():
+    split, meta = load_or_create_split()
+    feature_names = meta.get("feature_names") if isinstance(meta, dict) else None
+    X_train = split["X_train"]
+    X_test = split["X_test"]
+    y_train = split["y_train"]
+    y_test = split["y_test"]
+
+    if feature_names is not None:
+        X_train = pd.DataFrame(X_train, columns=feature_names)
+        X_test = pd.DataFrame(X_test, columns=feature_names)
+
+    classes = np.unique(y_train)
+    class_weights = compute_class_weight(class_weight="balanced", classes=classes, y=y_train)
+    class_weight_dict = {cls: weight for cls, weight in zip(classes, class_weights)}
+
+    models = train_models(X_train, y_train, class_weight_dict)
+
+    metrics = []
+    for name, model in models.items():
+        metrics.append(evaluate_model(name, model, X_test, y_test))
+        joblib.dump(model, ROOT / "models" / "saved" / f"{name.lower()}_model.joblib")
+
+    summary_path = EVAL_DIR / "model_metrics_summary_traditional.csv"
+    pd.DataFrame(metrics).to_csv(summary_path, index=False)
+    print(f"Saved traditional summary to {summary_path}")
+
+    (EVAL_DIR / "traditional_metrics_latest.json").write_text(
+        json.dumps(metrics, indent=2)
+    )
+
+
+if __name__ == "__main__":
+    main()

src/evaluate_transformer_models.py

@@ -0,0 +1,274 @@
+"""Evaluate transformer models (BioGPT, PubMedBERT, Clinical-T5) on the leak-free split."""
+from __future__ import annotations
+
+import json
+import sys
+from pathlib import Path
+
+import joblib
+import numpy as np
+import pandas as pd
+import torch
+from sklearn.metrics import (
+    accuracy_score,
+    auc,
+    classification_report,
+    confusion_matrix,
+    precision_recall_fscore_support,
+    roc_curve,
+)
+from sklearn.preprocessing import label_binarize
+from torch.utils.data import DataLoader
+
+sys.path.append(str(Path(__file__).parent))
+from data_preprocessing import DataPreprocessor  # type: ignore
+from document_manager import DocumentManager  # type: ignore
+from models.medical_transformers import (  # type: ignore
+    BioMistralClassifier as BioGPTForTabular,
+    ClinicalT5Classifier as ClinicalT5ForTabular,
+    PubMedBERTClassifier as PubMedBERTForTabular,
+)
+from models.transformer_models import TabularDataset  # type: ignore
+
+ROOT = Path(__file__).resolve().parents[1]
+EVAL_DIR = ROOT / "evaluation_results" / "model_metrics"
+CLASS_REPORT_DIR = EVAL_DIR / "classification_reports"
+CONF_MATRIX_DIR = EVAL_DIR / "confusion_matrices"
+ROC_DIR = EVAL_DIR / "roc_curves"
+SUMMARY_PATH = EVAL_DIR / "model_metrics_summary_transformers.csv"
+LATEST_JSON = EVAL_DIR / "transformer_metrics_latest.json"
+LEAK_FREE_SPLIT_PATH = ROOT / "evaluation_results" / "leak_free_split.npz"
+LEAK_FREE_META_PATH = ROOT / "evaluation_results" / "leak_free_split_meta.joblib"
+CLASS_NAMES = ["HC", "PD", "SWEDD", "PRODROMAL"]
+
+for directory in (CLASS_REPORT_DIR, CONF_MATRIX_DIR, ROC_DIR):
+    directory.mkdir(parents=True, exist_ok=True)
+
+
+def _load_or_create_leak_free_split() -> tuple[np.ndarray, np.ndarray, np.ndarray, np.ndarray, list[str]]:
+    if LEAK_FREE_SPLIT_PATH.exists() and LEAK_FREE_META_PATH.exists():
+        split = np.load(LEAK_FREE_SPLIT_PATH)
+        meta = joblib.load(LEAK_FREE_META_PATH)
+        feature_names = meta.get("feature_names") if isinstance(meta, dict) else None
+        if feature_names is None:
+            raise ValueError("Leak-free metadata missing feature names")
+        print("Loaded cached leak-free split artifacts.")
+        return split["X_train"], split["X_test"], split["y_train"], split["y_test"], feature_names
+
+    preprocessor = DataPreprocessor()
+    file_paths = [
+        ROOT / "PPMI_Curated_Data_Cut_Public_20240129.csv",
+        ROOT / "PPMI_Curated_Data_Cut_Public_20241211.csv",
+        ROOT / "PPMI_Curated_Data_Cut_Public_20250321.csv",
+        ROOT / "PPMI_Curated_Data_Cut_Public_20250714.csv",
+    ]
+    X_train, X_test, y_train, y_test = preprocessor.prepare_data(
+        file_paths,
+        test_size=0.2,
+        use_patient_split=True,
+    )
+    feature_names = preprocessor.get_feature_names()
+    LEAK_FREE_SPLIT_PATH.parent.mkdir(parents=True, exist_ok=True)
+    np.savez(
+        LEAK_FREE_SPLIT_PATH,
+        X_train=X_train,
+        X_test=X_test,
+        y_train=y_train,
+        y_test=y_test,
+    )
+    joblib.dump({"feature_names": feature_names}, LEAK_FREE_META_PATH)
+    print("Saved new leak-free split artifacts.")
+    return X_train, X_test, y_train, y_test, feature_names
+
+
+def _prepare_batch(batch, device):
+    if len(batch) == 3:
+        data, targets, contexts = batch
+        contexts = list(contexts)
+    else:
+        data, targets = batch
+        contexts = None
+    data = data.to(device)
+    targets = targets.to(device)
+    return data, targets, contexts
+
+
+def _build_context_cache(features: np.ndarray, feature_names: list[str], doc_manager: DocumentManager) -> list[str]:
+    cache = []
+    total = len(features)
+    for idx, row in enumerate(features):
+        feature_desc = {name: float(val) for name, val in zip(feature_names, row)}
+        query_parts = []
+        for symptom_key, col in {
+            "tremor": "sym_tremor",
+            "rigidity": "sym_rigid",
+            "bradykinesia": "sym_brady",
+            "postural instability": "sym_posins",
+        }.items():
+            if feature_desc.get(col, 0) > 0:
+                query_parts.append(f"{symptom_key} severity:{feature_desc[col]:.2f}")
+        moca = feature_desc.get("moca", 30)
+        if moca < 26:
+            query_parts.append("cognitive impairment")
+        age = feature_desc.get("age", 0)
+        if age:
+            query_parts.append(f"age {int(age)}")
+        if feature_desc.get("fampd", 0) > 0:
+            query_parts.append("family history Parkinson's disease")
+        query = "Parkinson's disease " + " ".join(query_parts)
+        passages = doc_manager.extract_relevant_passages(query, top_k=2)
+        if not passages:
+            cache.append("")
+        else:
+            combined = []
+            for passage in passages:
+                title = passage.get("doc_title") or passage.get("doc_id") or "document"
+                combined.append(f"From '{title}' {passage['text'][:300]}...")
+            cache.append(" ".join(combined))
+        if (idx + 1) % 250 == 0 or idx + 1 == total:
+            print(f"  Cached RAG context for {idx + 1}/{total} samples")
+    return cache
+
+
+def _save_outputs(model_name: str, report: str, cm: np.ndarray, y_test: np.ndarray, y_prob: np.ndarray) -> None:
+    (CLASS_REPORT_DIR / f"{model_name}.txt").write_text(
+        f"{model_name} Classification Report (leak-free split)\n" + "-" * 60 + "\n" + report
+    )
+    cm_df = pd.DataFrame(cm, index=CLASS_NAMES, columns=CLASS_NAMES)
+    cm_df.to_csv(CONF_MATRIX_DIR / f"{model_name}_confusion_matrix.csv")
+
+    y_bin = label_binarize(y_test, classes=range(len(CLASS_NAMES)))
+    for idx, class_name in enumerate(CLASS_NAMES):
+        fpr, tpr, _ = roc_curve(y_bin[:, idx], y_prob[:, idx])
+        roc_df = pd.DataFrame({"fpr": fpr, "tpr": tpr})
+        roc_df.to_csv(ROC_DIR / f"{model_name}_class_{class_name}_roc.csv", index=False)
+
+
+def main() -> None:
+    device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
+    print(f"Evaluating transformers on {device} (leak-free split)")
+    X_train, X_test, y_train, y_test, feature_names = _load_or_create_leak_free_split()
+    X_train = np.asarray(X_train)
+    X_test = np.asarray(X_test)
+    y_train = np.asarray(y_train)
+    y_test = np.asarray(y_test)
+
+    docs_path = ROOT / "medical_docs"
+    doc_manager = DocumentManager(docs_dir=str(docs_path))
+    print("Building RAG contexts for test set...")
+    test_contexts = _build_context_cache(X_test, feature_names, doc_manager)
+
+    test_dataset = TabularDataset(X_test, y_test, feature_names, contexts=test_contexts)
+    test_loader = DataLoader(test_dataset, batch_size=64, shuffle=False)
+
+    model_dir = ROOT / "models" / "saved"
+    model_configs = {
+        "BioGPT": {
+            "builder": lambda: BioGPTForTabular(
+                input_dim=X_train.shape[1],
+                num_classes=len(CLASS_NAMES),
+                dropout=0.15,
+                train_decoder_layers=6,
+            ),
+            "checkpoints": [
+                model_dir / "biogpt_transformer.pth",
+                model_dir / "biogpt.pth",
+                model_dir / "biomistral.pth",
+            ],
+        },
+        "PubMedBERT": {
+            "builder": lambda: PubMedBERTForTabular(
+                input_dim=X_train.shape[1],
+                num_classes=len(CLASS_NAMES),
+                dropout=0.15,
+                freeze_bert=False,
+            ),
+            "checkpoints": [
+                model_dir / "pubmedbert_transformer.pth",
+                model_dir / "pubmedbert.pth",
+            ],
+        },
+        "Clinical-T5": {
+            "builder": lambda: ClinicalT5ForTabular(
+                input_dim=X_train.shape[1],
+                num_classes=len(CLASS_NAMES),
+                dropout=0.15,
+                freeze_encoder=False,
+            ),
+            "checkpoints": [
+                model_dir / "clinical_t5_transformer.pth",
+                model_dir / "clinicalt5_transformer.pth",
+                model_dir / "clinical_t5.pth",
+            ],
+        },
+    }
+
+    summary_rows = []
+
+    for pretty_name, cfg in model_configs.items():
+        checkpoint_path = next((path for path in cfg["checkpoints"] if path.exists()), None)
+        if checkpoint_path is None:
+            expected = ", ".join(path.name for path in cfg["checkpoints"])
+            print(f"[WARN] Skipping {pretty_name}: no checkpoint found ({expected})")
+            continue
+
+        print(f"\nEvaluating {pretty_name}...")
+        model = cfg["builder"]().to(device)
+        state = torch.load(checkpoint_path, map_location=device)
+        model.load_state_dict(state)
+        model.eval()
+
+        all_targets = []
+        all_preds = []
+        all_prob = []
+
+        with torch.no_grad():
+            for batch in test_loader:
+                data, targets, contexts = _prepare_batch(batch, device)
+                outputs = model(data, contexts)
+                probs = torch.softmax(outputs, dim=1)
+                preds = torch.argmax(probs, dim=1)
+
+                all_targets.extend(targets.cpu().numpy())
+                all_preds.extend(preds.cpu().numpy())
+                all_prob.append(probs.cpu().numpy())
+
+        y_true = np.array(all_targets)
+        y_pred = np.array(all_preds)
+        y_prob = np.vstack(all_prob)
+
+        accuracy = accuracy_score(y_true, y_pred)
+        precision, recall, f1, _ = precision_recall_fscore_support(
+            y_true, y_pred, average="weighted", zero_division=0
+        )
+        report = classification_report(
+            y_true, y_pred, target_names=CLASS_NAMES, zero_division=0
+        )
+        cm = confusion_matrix(y_true, y_pred)
+        _save_outputs(pretty_name, report, cm, y_true, y_prob)
+
+        summary_rows.append(
+            {
+                "model": pretty_name,
+                "accuracy": accuracy,
+                "precision": precision,
+                "recall": recall,
+                "f1": f1,
+            }
+        )
+
+        print(
+            f"{pretty_name} -> accuracy {accuracy:.4f}, precision {precision:.4f}, recall {recall:.4f}, f1 {f1:.4f}"
+        )
+
+    if summary_rows:
+        summary_df = pd.DataFrame(summary_rows)
+        summary_df.to_csv(SUMMARY_PATH, index=False)
+        LATEST_JSON.write_text(json.dumps(summary_rows, indent=2))
+        print(f"Saved transformer summary to {SUMMARY_PATH}")
+    else:
+        print("No transformer metrics were generated; check checkpoints.")
+
+
+if __name__ == "__main__":
+    main()

src/feature_mapping.py

@@ -0,0 +1,189 @@
+class FeatureMapper:
+    def __init__(self):
+        # Mapping between patient-facing questions and the trained inference schema.
+        self.feature_mapping = {
+            "basic_info": {
+                "age": {
+                    "question": "What is your age?",
+                    "type": "numeric",
+                    "dataset_feature": "age",
+                },
+                "sex": {
+                    "question": "What is your biological sex?",
+                    "type": "categorical",
+                    "options": ["Male", "Female"],
+                    "dataset_feature": "SEX",
+                    "mapping": {"Male": 1, "Female": 0},
+                },
+                "education": {
+                    "question": "Years of education completed?",
+                    "type": "numeric",
+                    "dataset_feature": "EDUCYRS",
+                },
+                "race": {
+                    "question": "What is your race?",
+                    "type": "categorical",
+                    "options": ["White", "Black/African American", "Asian", "Other"],
+                    "dataset_feature": "race",
+                    "mapping": {
+                        "White": 1,
+                        "Black/African American": 2,
+                        "Asian": 3,
+                        "Other": 4,
+                    },
+                },
+                "bmi": {
+                    "question": "What is your BMI?",
+                    "type": "numeric",
+                    "dataset_feature": "BMI",
+                },
+            },
+            "family_history": {
+                "family_pd": {
+                    "question": "Do you have any family members with Parkinson's disease?",
+                    "type": "categorical",
+                    "options": [
+                        "No family history",
+                        "First degree relative",
+                        "Other relative",
+                    ],
+                    "dataset_feature": "fampd",
+                    "mapping": {
+                        "No family history": 3,
+                        "First degree relative": 1,
+                        "Other relative": 2,
+                    },
+                }
+            },
+            "motor_symptoms": {
+                "tremor": {
+                    "question": "Tremor severity (0-4)",
+                    "type": "numeric",
+                    "dataset_feature": "sym_tremor",
+                    "scale": "0-4",
+                },
+                "rigidity": {
+                    "question": "Rigidity severity (0-4)",
+                    "type": "numeric",
+                    "dataset_feature": "sym_rigid",
+                    "scale": "0-4",
+                },
+                "bradykinesia": {
+                    "question": "Bradykinesia severity (0-4)",
+                    "type": "numeric",
+                    "dataset_feature": "sym_brady",
+                    "scale": "0-4",
+                },
+                "balance": {
+                    "question": "Postural instability severity (0-4)",
+                    "type": "numeric",
+                    "dataset_feature": "sym_posins",
+                    "scale": "0-4",
+                },
+            },
+            "non_motor_symptoms": {
+                "rem_sleep": {
+                    "question": "Do you act out dreams or have REM sleep behaviour symptoms?",
+                    "type": "categorical",
+                    "options": ["No", "Yes"],
+                    "dataset_feature": "rem",
+                    "mapping": {"No": 0, "Yes": 1},
+                },
+                "daytime_sleepiness": {
+                    "question": "Epworth Sleepiness Scale score?",
+                    "type": "numeric",
+                    "dataset_feature": "ess",
+                    "scale": "0-24",
+                },
+                "depression": {
+                    "question": "Geriatric Depression Scale score?",
+                    "type": "numeric",
+                    "dataset_feature": "gds",
+                    "scale": "0-15",
+                },
+                "anxiety": {
+                    "question": "State-Trait Anxiety Inventory score?",
+                    "type": "numeric",
+                    "dataset_feature": "stai",
+                    "scale": "20-80",
+                },
+            },
+            "cognitive_symptoms": {
+                "memory": {
+                    "question": "MoCA score?",
+                    "type": "numeric",
+                    "dataset_feature": "moca",
+                    "scale": "0-30",
+                },
+                "clock_draw": {
+                    "question": "Clock drawing test score?",
+                    "type": "numeric",
+                    "dataset_feature": "clockdraw",
+                    "scale": "0-4",
+                },
+                "bjlot": {
+                    "question": "Benton line orientation score?",
+                    "type": "numeric",
+                    "dataset_feature": "bjlot",
+                    "scale": "0-30",
+                },
+            },
+        }
+
+    def get_patient_questionnaire(self):
+        """Generate a list of questions for patients."""
+        questions = []
+        for category in self.feature_mapping.values():
+            for feature in category.values():
+                questions.append(
+                    {
+                        "question": feature["question"],
+                        "type": feature["type"],
+                        "options": feature.get("options"),
+                        "scale": feature.get("scale"),
+                    }
+                )
+        return questions
+
+    def map_patient_response_to_features(self, responses):
+        """Map patient responses to dataset features."""
+        feature_values = {}
+        for category in self.feature_mapping.values():
+            for feature_name, feature_info in category.items():
+                if feature_name not in responses:
+                    continue
+
+                response = responses[feature_name]
+                dataset_feature = feature_info["dataset_feature"]
+
+                if "mapping" in feature_info:
+                    mapped_value = feature_info["mapping"].get(response)
+                else:
+                    mapped_value = response
+
+                if mapped_value is None:
+                    continue
+
+                feature_values[dataset_feature] = mapped_value
+
+                if dataset_feature == "fampd":
+                    feature_values["fampd_bin"] = 2 if mapped_value == 3 else 1
+
+        return feature_values
+
+
+def main():
+    mapper = FeatureMapper()
+    questions = mapper.get_patient_questionnaire()
+
+    print("Patient Questionnaire:")
+    for i, q in enumerate(questions, 1):
+        print(f"\n{i}. {q['question']}")
+        if q["options"]:
+            print(f"Options: {', '.join(q['options'])}")
+        if q["scale"]:
+            print(f"Scale: {q['scale']}")
+
+
+if __name__ == "__main__":
+    main()

src/models/medical_transformers.py

@@ -0,0 +1,939 @@
+"""
+Medical Transformer Models for Parkinson's Disease Classification
+Implements three specialized medical transformers:
+1. PubMedBERT - Encoder-only model
+2. BioMistral-7B - Decoder-only model
+3. Clinical-T5 - Encoder-Decoder model
+
+Optimized for GPU training with CUDA support.
+"""
+
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+from torch.utils.data import Dataset, DataLoader
+import numpy as np
+from sklearn.metrics import accuracy_score, classification_report, confusion_matrix, roc_auc_score, f1_score
+from sklearn.model_selection import StratifiedKFold
+from sklearn.preprocessing import label_binarize
+import joblib
+import os
+from typing import Tuple, Dict, Any, List, Optional
+import matplotlib.pyplot as plt
+import seaborn as sns
+from transformers import (
+    AutoModel,
+    AutoTokenizer,
+    AutoModelForCausalLM,
+    T5ForConditionalGeneration,
+    T5Tokenizer,
+    BitsAndBytesConfig
+)
+from datetime import datetime
+import warnings
+warnings.filterwarnings('ignore')
+
+
+class MedicalTabularDataset(Dataset):
+    """Enhanced dataset for tabular medical data with feature descriptions."""
+    
+    def __init__(self, X, y, feature_names=None):
+        self.X = torch.FloatTensor(X)
+        self.y = torch.LongTensor(y)
+        self.feature_names = feature_names or [f"feature_{i}" for i in range(X.shape[1])]
+        
+        # Feature descriptions for medical context
+        self.feature_descriptions = {
+            'age': 'patient age in years',
+            'SEX': 'biological sex (0=Female, 1=Male)',
+            'EDUCYRS': 'years of education',
+            'BMI': 'body mass index',
+            'fampd': 'family history of Parkinsons disease',
+            'sym_tremor': 'tremor severity score',
+            'sym_rigid': 'rigidity severity score',
+            'sym_brady': 'bradykinesia severity score',
+            'sym_posins': 'postural instability score',
+            'moca': 'Montreal Cognitive Assessment score',
+            'gds': 'Geriatric Depression Scale score',
+            'stai': 'State-Trait Anxiety Inventory score',
+            'rem': 'REM sleep behavior disorder check',
+            'ess': 'Epworth Sleepiness Scale',
+            'upsit': 'University of Pennsylvania Smell Identification Test score',
+            'upsit_pctl': 'smell identification percentile',
+            'updrs1_score': 'UPDRS Part I (non-motor experiences of daily living)',
+            'updrs2_score': 'UPDRS Part II (motor experiences of daily living)',
+            'updrs3_score': 'UPDRS Part III (motor examination)',
+            'updrs4_score': 'UPDRS Part IV (motor complications)',
+            'updrs_totscore': 'Total UPDRS score',
+            'mean_caudate': 'DatScan mean caudate uptake ratio',
+            'mean_putamen': 'DatScan mean putamen uptake ratio',
+            'abeta': 'CSF Amyloid beta 1-42 level',
+            'tau': 'CSF Total Tau level',
+            'ptau': 'CSF Phosphorylated Tau level'
+        }
+    
+    def __len__(self):
+        return len(self.X)
+    
+    def __getitem__(self, idx):
+        return self.X[idx], self.y[idx]
+    
+    def get_text_description(self, idx):
+        """Convert features to medical text description."""
+        sample = self.X[idx].numpy()
+        descriptions = []
+        
+        for i, (name, val) in enumerate(zip(self.feature_names, sample)):
+            base_name = name.split('_')[0] if '_' in name else name
+            desc = self.feature_descriptions.get(base_name, f"{name}")
+            descriptions.append(f"{desc}: {val:.2f}")
+        
+        return "Patient clinical data: " + ", ".join(descriptions)
+
+
+class PubMedBERTClassifier(nn.Module):
+    """
+    PubMedBERT - Encoder-only transformer model
+    Pretrained on PubMed abstracts, optimized for medical text understanding
+    """
+    
+    def __init__(self, input_dim: int, num_classes: int, dropout: float = 0.1,
+                 freeze_bert: bool = True, train_encoder_layers: int = 4):
+        super(PubMedBERTClassifier, self).__init__()
+        
+        self.model_name = "microsoft/BiomedNLP-PubMedBERT-base-uncased-abstract-fulltext"
+        print(f"Loading PubMedBERT from {self.model_name}")
+        
+        # Load pretrained model
+        self.tokenizer = AutoTokenizer.from_pretrained(self.model_name)
+        self.bert = AutoModel.from_pretrained(self.model_name)
+        
+        # Optionally freeze BERT parameters
+        if freeze_bert:
+            for param in self.bert.parameters():
+                param.requires_grad = False
+            print("PubMedBERT parameters frozen")
+        else:
+            # Freeze everything first, then unfreeze the last encoder blocks.
+            for param in self.bert.parameters():
+                param.requires_grad = False
+            encoder_layers = getattr(self.bert.encoder, "layer", [])
+            layers_to_unfreeze = min(int(train_encoder_layers), len(encoder_layers))
+            for layer in encoder_layers[-layers_to_unfreeze:]:
+                for param in layer.parameters():
+                    param.requires_grad = True
+            pooler = getattr(self.bert, "pooler", None)
+            if pooler is not None:
+                for param in pooler.parameters():
+                    param.requires_grad = True
+            trainable = sum(p.numel() for p in self.bert.parameters() if p.requires_grad)
+            total = sum(p.numel() for p in self.bert.parameters())
+            print(
+                f"PubMedBERT last {layers_to_unfreeze} encoder layers unfrozen "
+                f"({trainable:,}/{total:,} backbone params trainable)"
+            )
+        
+        self.hidden_size = self.bert.config.hidden_size  # 768 for base model
+        
+        # Feature projection layer
+        self.feature_projection = nn.Sequential(
+            nn.Linear(input_dim, 384),
+            nn.LayerNorm(384),
+            nn.ReLU(),
+            nn.Dropout(dropout),
+            nn.Linear(384, self.hidden_size)
+        )
+        
+        # Multi-head attention for feature fusion
+        self.attention = nn.MultiheadAttention(
+            embed_dim=self.hidden_size,
+            num_heads=8,
+            dropout=dropout,
+            batch_first=True
+        )
+        
+        # Classification head with residual connection
+        self.classifier = nn.Sequential(
+            nn.Linear(self.hidden_size * 2, 512),
+            nn.LayerNorm(512),
+            nn.ReLU(),
+            nn.Dropout(dropout),
+            nn.Linear(512, 256),
+            nn.LayerNorm(256),
+            nn.ReLU(),
+            nn.Dropout(dropout),
+            nn.Linear(256, num_classes)
+        )
+        
+    def forward(self, x, text_input=None):
+        batch_size = x.size(0)
+        
+        # Project tabular features
+        tabular_features = self.feature_projection(x)  # [batch, 768]
+        tabular_features = tabular_features.unsqueeze(1)  # [batch, 1, 768]
+        
+        if text_input is not None:
+            # Process text with BERT
+            if isinstance(text_input, dict):
+                inputs = {k: v.to(x.device) for k, v in text_input.items()}
+            else:
+                inputs = self.tokenizer(
+                    text_input,
+                    return_tensors="pt",
+                    padding=True,
+                    truncation=True,
+                    max_length=512
+                )
+                inputs = {k: v.to(x.device) for k, v in inputs.items()}
+            
+            # Get BERT embeddings
+            outputs = self.bert(**inputs)
+            text_features = outputs.last_hidden_state  # [batch, seq_len, 768]
+            
+            # Concatenate features
+            combined = torch.cat([tabular_features, text_features], dim=1)
+            
+            # Apply attention
+            attended, _ = self.attention(combined, combined, combined)
+            pooled = attended.mean(dim=1)  # Average pooling
+        else:
+            # Use only tabular features
+            pooled = tabular_features.squeeze(1)
+        
+        # Combine original and attended features
+        combined_features = torch.cat([pooled, tabular_features.squeeze(1)], dim=1)
+        
+        # Classification
+        output = self.classifier(combined_features)
+        
+        return output
+
+
+class BioMistralClassifier(nn.Module):
+    """BioGPT (decoder-only) hybrid with optional partial fine-tuning."""
+    
+    def __init__(
+        self,
+        input_dim: int,
+        num_classes: int,
+        dropout: float = 0.1,
+        use_quantization: bool = False,
+        train_decoder_layers: int = 4
+    ):
+        super(BioMistralClassifier, self).__init__()
+        
+        # Use a smaller medical language model that's publicly available
+        # BioGPT or similar decoder-only model
+        self.model_name = "microsoft/biogpt"
+        print(f"Loading BioGPT (Decoder-only) from {self.model_name}")
+
+        try:
+            try:
+                import sacremoses  # type: ignore  # noqa: F401
+            except ImportError as exc:
+                raise RuntimeError(
+                    "BioGPT requires the 'sacremoses' package for tokenization. "
+                    "Install it with 'pip install sacremoses' or reinstall from requirements.txt."
+                ) from exc
+
+            # Load without quantization and device_map to avoid accelerate dependency
+            # Quantization disabled by default to avoid bitsandbytes and accelerate issues
+            self.model = AutoModelForCausalLM.from_pretrained(
+                self.model_name,
+                trust_remote_code=True,
+                torch_dtype=torch.float32  # Use standard precision
+            )
+            print("Model loaded successfully (no quantization)")
+                
+            self.tokenizer = AutoTokenizer.from_pretrained(self.model_name)
+            if self.tokenizer.pad_token is None:
+                self.tokenizer.pad_token = self.tokenizer.eos_token
+            
+            # Move model to device
+            device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
+            self.model = self.model.to(device)
+
+            # Freeze decoder parameters (we may selectively unfreeze later)
+            for param in self.model.parameters():
+                param.requires_grad = False
+
+            self.hidden_size = self.model.config.hidden_size
+
+        except Exception as e:
+            raise RuntimeError(
+                "BioGPT failed to load. This training path now requires the real "
+                f"BioGPT model and will not fall back to GPT2. Original error: {e}"
+            ) from e
+        
+        self._unfreeze_decoder_layers(train_decoder_layers)
+        
+        # Feature encoder (larger capacity)
+        self.feature_encoder = nn.Sequential(
+            nn.Linear(input_dim, 512),
+            nn.LayerNorm(512),
+            nn.GELU(),
+            nn.Dropout(dropout),
+            nn.Linear(512, self.hidden_size),
+            nn.GELU()
+        )
+        
+        # Cross-attention layer (restore 8 heads)
+        self.cross_attention = nn.MultiheadAttention(
+            embed_dim=self.hidden_size,
+            num_heads=8,
+            dropout=dropout,
+            batch_first=True
+        )
+        
+        # Classification head with more width/depth
+        self.classifier = nn.Sequential(
+            nn.Linear(self.hidden_size, 1024),
+            nn.LayerNorm(1024),
+            nn.GELU(),
+            nn.Dropout(dropout),
+            nn.Linear(1024, 512),
+            nn.LayerNorm(512),
+            nn.GELU(),
+            nn.Dropout(dropout),
+            nn.Linear(512, 256),
+            nn.GELU(),
+            nn.Dropout(dropout / 2 if dropout > 0 else 0.0),
+            nn.Linear(256, num_classes)
+        )
+        
+    def forward(self, x, text_input=None):
+        # Encode tabular features
+        feature_embeddings = self.feature_encoder(x)  # [batch, hidden_size]
+        feature_embeddings = feature_embeddings.unsqueeze(1)  # [batch, 1, hidden_size]
+        
+        if text_input is not None:
+            # Tokenize text
+            if isinstance(text_input, dict):
+                inputs = {k: v.to(x.device) for k, v in text_input.items()}
+            else:
+                inputs = self.tokenizer(
+                    text_input,
+                    return_tensors="pt",
+                    padding=True,
+                    truncation=True,
+                    max_length=512
+                )
+                inputs = {k: v.to(x.device) for k, v in inputs.items()}
+            
+            # Get decoder outputs (no torch.no_grad — allow gradients for unfrozen layers)
+            outputs = self.model(**inputs, output_hidden_states=True)
+            hidden_states = outputs.hidden_states[-1]  # Last layer
+            
+            # Cross-attention between features and text
+            attended, _ = self.cross_attention(
+                feature_embeddings,
+                hidden_states,
+                hidden_states
+            )
+            pooled = attended.squeeze(1)
+        else:
+            pooled = feature_embeddings.squeeze(1)
+        
+        # Classification
+        output = self.classifier(pooled)
+        
+        return output
+
+    def _unfreeze_decoder_layers(self, num_layers: int):
+        """Optionally unfreeze the last decoder layers to grow capacity."""
+        if num_layers <= 0:
+            print("BioGPT remains fully frozen (train_decoder_layers=0)")
+            return
+
+        # Debug: print model structure to diagnose layer discovery
+        print(f"  [DEBUG] Model type: {type(self.model).__name__}")
+        print(f"  [DEBUG] Top-level children: {[n for n, _ in self.model.named_children()]}")
+
+        layers = None
+        layernorm = None
+        head = None
+
+        # Strategy 1: Try known attribute paths
+        if hasattr(self.model, 'biogpt'):
+            layers = getattr(self.model.biogpt, 'layers', None)
+            layernorm = getattr(self.model.biogpt, 'layer_norm', None)
+            head = getattr(self.model, 'output_projection', None)
+        elif hasattr(self.model, 'transformer'):
+            layers = getattr(self.model.transformer, 'h', None)
+            layernorm = getattr(self.model.transformer, 'ln_f', None)
+            head = getattr(self.model, 'lm_head', None)
+
+        # Strategy 2: Search all named modules for a large ModuleList (transformer layers)
+        if layers is None:
+            print("  [DEBUG] Known paths failed, searching named_modules...")
+            for name, module in self.model.named_modules():
+                if isinstance(module, nn.ModuleList) and len(module) > 6:
+                    layers = module
+                    print(f"  [DEBUG] Found transformer layers at '{name}' ({len(module)} layers)")
+                    break
+
+        # Strategy 3: Search for layernorm and head if not found yet
+        if layers is not None and layernorm is None:
+            for name, module in self.model.named_modules():
+                if ('layer_norm' in name or 'ln_f' in name) and not any(c in name for c in ['.0.', '.1.', '.2.']):
+                    layernorm = module
+                    break
+        if layers is not None and head is None:
+            for name, module in self.model.named_modules():
+                if 'output_projection' in name or 'lm_head' in name:
+                    head = module
+                    break
+
+        if layers is None:
+            print("Warning: could not locate transformer layers to unfreeze")
+            print("  All named modules:")
+            for name, module in self.model.named_children():
+                print(f"    {name}: {type(module).__name__}")
+            return
+
+        total_layers = len(layers)
+        num_layers = min(num_layers, total_layers)
+
+        print(f"Unfreezing last {num_layers}/{total_layers} layers of BioGPT...")
+
+        # Unfreeze last N layers
+        for layer in layers[-num_layers:]:
+            for param in layer.parameters():
+                param.requires_grad = True
+
+        # Unfreeze final layer norm
+        if layernorm is not None:
+            for param in layernorm.parameters():
+                param.requires_grad = True
+
+        # Unfreeze head
+        if head is not None:
+            for param in head.parameters():
+                param.requires_grad = True
+
+        trainable = sum(p.numel() for p in self.model.parameters() if p.requires_grad)
+        total = sum(p.numel() for p in self.model.parameters())
+        print(f"Unfroze last {num_layers} decoder layers plus norm/head for BioGPT")
+        print(f"  Backbone trainable: {trainable:,}/{total:,} ({100*trainable/total:.1f}%)")
+
+
+class ClinicalT5Classifier(nn.Module):
+    """
+    Clinical-T5 - Encoder-Decoder transformer model
+    T5 model fine-tuned on clinical notes and medical literature
+    """
+    
+    def __init__(self, input_dim: int, num_classes: int, dropout: float = 0.1,
+                 freeze_encoder: bool = True):
+        super(ClinicalT5Classifier, self).__init__()
+        
+        # Use T5 or Flan-T5 as base model
+        self.model_name = "google/flan-t5-base"
+        print(f"Loading Clinical T5 (Encoder-Decoder) from {self.model_name}")
+        
+        self.tokenizer = T5Tokenizer.from_pretrained(self.model_name)
+        self.t5 = T5ForConditionalGeneration.from_pretrained(self.model_name)
+        
+        # Freeze encoder/decoder parameters
+        if freeze_encoder:
+            for param in self.t5.encoder.parameters():
+                param.requires_grad = False
+            for param in self.t5.decoder.parameters():
+                param.requires_grad = False
+            print("T5 encoder and decoder frozen")
+        else:
+            # Freeze all first, then selectively unfreeze last blocks
+            for param in self.t5.parameters():
+                param.requires_grad = False
+            # Unfreeze last 4 encoder blocks
+            for block in self.t5.encoder.block[-4:]:
+                for param in block.parameters():
+                    param.requires_grad = True
+            # Unfreeze last 2 decoder blocks
+            for block in self.t5.decoder.block[-2:]:
+                for param in block.parameters():
+                    param.requires_grad = True
+            # Unfreeze final layer norms
+            if hasattr(self.t5.encoder, 'final_layer_norm'):
+                for param in self.t5.encoder.final_layer_norm.parameters():
+                    param.requires_grad = True
+            if hasattr(self.t5.decoder, 'final_layer_norm'):
+                for param in self.t5.decoder.final_layer_norm.parameters():
+                    param.requires_grad = True
+            trainable = sum(p.numel() for p in self.t5.parameters() if p.requires_grad)
+            total = sum(p.numel() for p in self.t5.parameters())
+            print(f"T5 partially unfrozen: {trainable:,}/{total:,} params trainable ({100*trainable/total:.1f}%)")
+        
+        self.hidden_size = self.t5.config.d_model  # 768 for base
+        
+        # Feature projection
+        self.feature_projection = nn.Sequential(
+            nn.Linear(input_dim, 384),
+            nn.LayerNorm(384),
+            nn.ReLU(),
+            nn.Dropout(dropout),
+            nn.Linear(384, self.hidden_size)
+        )
+        
+        # Encoder-decoder attention mechanism
+        self.encoder_decoder_attention = nn.MultiheadAttention(
+            embed_dim=self.hidden_size,
+            num_heads=8,
+            dropout=dropout,
+            batch_first=True
+        )
+        
+        # Classification head
+        self.classifier = nn.Sequential(
+            nn.Linear(self.hidden_size * 2, 512),
+            nn.LayerNorm(512),
+            nn.ReLU(),
+            nn.Dropout(dropout),
+            nn.Linear(512, 256),
+            nn.LayerNorm(256),
+            nn.ReLU(),
+            nn.Dropout(dropout),
+            nn.Linear(256, num_classes)
+        )
+        
+    def forward(self, x, text_input=None):
+        # Project features
+        feature_embeddings = self.feature_projection(x)
+        feature_embeddings = feature_embeddings.unsqueeze(1)
+        
+        if text_input is not None:
+            if isinstance(text_input, dict):
+                inputs = {k: v.to(x.device) for k, v in text_input.items()}
+            else:
+                # Prepare input text
+                input_texts = [f"classify patient: {text}" for text in text_input]
+                
+                # Tokenize
+                inputs = self.tokenizer(
+                    input_texts,
+                    return_tensors="pt",
+                    padding=True,
+                    truncation=True,
+                    max_length=512
+                )
+                inputs = {k: v.to(x.device) for k, v in inputs.items()}
+            
+            # Create dummy decoder inputs
+            decoder_input_ids = torch.zeros(
+                (inputs['input_ids'].shape[0], 1),
+                dtype=torch.long,
+                device=x.device
+            )
+            
+            # Get encoder and decoder outputs (allow gradients for unfrozen layers)
+            encoder_outputs = self.t5.encoder(**inputs)
+            decoder_outputs = self.t5.decoder(
+                input_ids=decoder_input_ids,
+                encoder_hidden_states=encoder_outputs.last_hidden_state
+            )
+            
+            encoder_hidden = encoder_outputs.last_hidden_state
+            decoder_hidden = decoder_outputs.last_hidden_state
+            
+            # Combine encoder and decoder information
+            combined = torch.cat([encoder_hidden, decoder_hidden], dim=1)
+            
+            # Attend to features
+            attended, _ = self.encoder_decoder_attention(
+                feature_embeddings,
+                combined,
+                combined
+            )
+            
+            pooled = attended.squeeze(1)
+        else:
+            pooled = feature_embeddings.squeeze(1)
+        
+        # Combine with original features
+        combined_features = torch.cat([pooled, feature_embeddings.squeeze(1)], dim=1)
+        
+        # Classification
+        output = self.classifier(combined_features)
+        
+        return output
+
+
+class MedicalTransformerTrainer:
+    """
+    Trainer class for medical transformer models with GPU optimization
+    """
+    
+    def __init__(self, device=None):
+        if device is None:
+            self.device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
+        else:
+            self.device = torch.device(device)
+        
+        print(f"Using device: {self.device}")
+        if self.device.type == 'cuda':
+            print(f"GPU: {torch.cuda.get_device_name(0)}")
+            print(f"GPU Memory: {torch.cuda.get_device_properties(0).total_memory / 1e9:.2f} GB")
+        
+        self.models = {}
+        self.training_histories = {}
+        self.evaluation_results = {}
+    
+    def create_model(self, model_type: str, input_dim: int, num_classes: int, **kwargs):
+        """Create a medical transformer model."""
+        if model_type == 'pubmedbert':
+            model = PubMedBERTClassifier(input_dim, num_classes, **kwargs)
+        elif model_type == 'biomistral':
+            model = BioMistralClassifier(input_dim, num_classes, **kwargs)
+        elif model_type == 'clinical_t5':
+            model = ClinicalT5Classifier(input_dim, num_classes, **kwargs)
+        else:
+            raise ValueError(f"Unknown model type: {model_type}")
+        
+        return model.to(self.device)
+    
+    def train_model(self, model, train_loader, val_loader, model_name: str,
+                   epochs: int = 50, lr: float = 0.001, weight_decay: float = 1e-4,
+                   class_weights: torch.Tensor = None, use_amp: bool = True):
+        """
+        Train model with mixed precision and gradient accumulation for GPU efficiency
+        """
+        print(f"\nTraining {model_name}")
+        print("=" * 70)
+        
+        # Loss function
+        if class_weights is not None:
+            criterion = nn.CrossEntropyLoss(weight=class_weights.to(self.device))
+        else:
+            criterion = nn.CrossEntropyLoss()
+        
+        # Optimizer
+        optimizer = torch.optim.AdamW(
+            filter(lambda p: p.requires_grad, model.parameters()),
+            lr=lr,
+            weight_decay=weight_decay
+        )
+        
+        # Learning rate scheduler
+        scheduler = torch.optim.lr_scheduler.CosineAnnealingWarmRestarts(
+            optimizer,
+            T_0=10,
+            T_mult=2,
+            eta_min=1e-6
+        )
+        
+        # Mixed precision training
+        scaler = torch.cuda.amp.GradScaler() if use_amp and self.device.type == 'cuda' else None
+        
+        # Training history
+        history = {
+            'train_loss': [],
+            'val_loss': [],
+            'val_accuracy': [],
+            'val_f1': [],
+            'learning_rates': []
+        }
+        
+        best_val_loss = float('inf')
+        best_val_acc = 0.0
+        patience_counter = 0
+        early_stopping_patience = 15
+        
+        for epoch in range(epochs):
+            # Training phase
+            model.train()
+            train_loss = 0.0
+            train_steps = 0
+            
+            for batch_x, batch_y in train_loader:
+                batch_x = batch_x.to(self.device)
+                batch_y = batch_y.to(self.device)
+                
+                optimizer.zero_grad()
+                
+                # Mixed precision forward pass
+                if scaler is not None:
+                    with torch.cuda.amp.autocast():
+                        outputs = model(batch_x)
+                        loss = criterion(outputs, batch_y)
+                    
+                    scaler.scale(loss).backward()
+                    scaler.step(optimizer)
+                    scaler.update()
+                else:
+                    outputs = model(batch_x)
+                    loss = criterion(outputs, batch_y)
+                    loss.backward()
+                    optimizer.step()
+                
+                train_loss += loss.item()
+                train_steps += 1
+            
+            # Validation phase
+            model.eval()
+            val_loss = 0.0
+            all_preds = []
+            all_targets = []
+            
+            with torch.no_grad():
+                for batch_x, batch_y in val_loader:
+                    batch_x = batch_x.to(self.device)
+                    batch_y = batch_y.to(self.device)
+                    
+                    if scaler is not None:
+                        with torch.cuda.amp.autocast():
+                            outputs = model(batch_x)
+                            loss = criterion(outputs, batch_y)
+                    else:
+                        outputs = model(batch_x)
+                        loss = criterion(outputs, batch_y)
+                    
+                    val_loss += loss.item()
+                    
+                    _, predicted = torch.max(outputs, 1)
+                    all_preds.extend(predicted.cpu().numpy())
+                    all_targets.extend(batch_y.cpu().numpy())
+            
+            # Calculate metrics
+            train_loss /= train_steps
+            val_loss /= len(val_loader)
+            val_accuracy = accuracy_score(all_targets, all_preds)
+            val_f1 = f1_score(all_targets, all_preds, average='weighted')
+            
+            # Update history
+            history['train_loss'].append(train_loss)
+            history['val_loss'].append(val_loss)
+            history['val_accuracy'].append(val_accuracy)
+            history['val_f1'].append(val_f1)
+            history['learning_rates'].append(optimizer.param_groups[0]['lr'])
+            
+            # Step scheduler
+            scheduler.step()
+            
+            # Print progress
+            if epoch % 5 == 0 or epoch == epochs - 1:
+                print(f"Epoch [{epoch+1}/{epochs}]")
+                print(f"  Train Loss: {train_loss:.4f}")
+                print(f"  Val Loss: {val_loss:.4f}")
+                print(f"  Val Accuracy: {val_accuracy:.4f}")
+                print(f"  Val F1: {val_f1:.4f}")
+                print(f"  LR: {optimizer.param_groups[0]['lr']:.6f}")
+            
+            # Early stopping
+            if val_loss < best_val_loss:
+                best_val_loss = val_loss
+                best_val_acc = val_accuracy
+                patience_counter = 0
+                # Save best model
+                self.save_model(model, f"{model_name}_best")
+            else:
+                patience_counter += 1
+                if patience_counter >= early_stopping_patience:
+                    print(f"\nEarly stopping at epoch {epoch+1}")
+                    print(f"Best validation loss: {best_val_loss:.4f}")
+                    print(f"Best validation accuracy: {best_val_acc:.4f}")
+                    break
+        
+        self.training_histories[model_name] = history
+        return history
+    
+    def evaluate_model(self, model, test_loader, model_name: str):
+        """Comprehensive model evaluation."""
+        print(f"\nEvaluating {model_name}")
+        print("=" * 70)
+        
+        model.eval()
+        all_preds = []
+        all_probs = []
+        all_targets = []
+        
+        with torch.no_grad():
+            for batch_x, batch_y in test_loader:
+                batch_x = batch_x.to(self.device)
+                batch_y = batch_y.to(self.device)
+                
+                outputs = model(batch_x)
+                probs = F.softmax(outputs, dim=1)
+                _, predicted = torch.max(outputs, 1)
+                
+                all_preds.extend(predicted.cpu().numpy())
+                all_probs.extend(probs.cpu().numpy())
+                all_targets.extend(batch_y.cpu().numpy())
+        
+        all_preds = np.array(all_preds)
+        all_probs = np.array(all_probs)
+        all_targets = np.array(all_targets)
+        
+        # Calculate metrics
+        accuracy = accuracy_score(all_targets, all_preds)
+        f1 = f1_score(all_targets, all_preds, average='weighted')
+        
+        # ROC-AUC for multiclass
+        try:
+            y_bin = label_binarize(all_targets, classes=range(len(np.unique(all_targets))))
+            roc_auc = roc_auc_score(y_bin, all_probs, average='weighted', multi_class='ovr')
+        except:
+            roc_auc = 0.0
+        
+        cm = confusion_matrix(all_targets, all_preds)
+        report = classification_report(all_targets, all_preds)
+        
+        results = {
+            'accuracy': accuracy,
+            'f1_score': f1,
+            'roc_auc': roc_auc,
+            'confusion_matrix': cm,
+            'classification_report': report,
+            'predictions': all_preds,
+            'probabilities': all_probs,
+            'targets': all_targets
+        }
+        
+        print(f"\nResults for {model_name}:")
+        print(f"  Accuracy: {accuracy:.4f}")
+        print(f"  F1 Score: {f1:.4f}")
+        print(f"  ROC-AUC: {roc_auc:.4f}")
+        print(f"\nClassification Report:\n{report}")
+        
+        self.evaluation_results[model_name] = results
+        return results
+    
+    def save_model(self, model, model_name: str, save_dir: str = "./models/saved"):
+        """Save model with error handling."""
+        try:
+            # Use absolute path
+            if not os.path.isabs(save_dir):
+                save_dir = os.path.abspath(save_dir)
+            os.makedirs(save_dir, exist_ok=True)
+            model_path = os.path.join(save_dir, f"{model_name}.pth")
+            torch.save(model.state_dict(), model_path)
+            print(f"[OK] Model saved to {model_path}")
+        except Exception as e:
+            print(f"[ERR] Error saving model: {e}")
+            import traceback
+            traceback.print_exc()
+    
+    def load_model(self, model_type: str, model_name: str, input_dim: int,
+                   num_classes: int, save_dir: str = "./models/saved", **kwargs):
+        """Load saved model."""
+        # Use absolute path
+        if not os.path.isabs(save_dir):
+            save_dir = os.path.abspath(save_dir)
+        model = self.create_model(model_type, input_dim, num_classes, **kwargs)
+        model_path = os.path.join(save_dir, f"{model_name}.pth")
+        model.load_state_dict(torch.load(model_path, map_location=self.device))
+        model.eval()
+        return model
+    
+    def plot_training_curves(self, model_name: str, save_dir: str = "./notebooks"):
+        """Plot training curves."""
+        if model_name not in self.training_histories:
+            print(f"No training history found for {model_name}")
+            return
+        
+        # Use absolute path
+        if not os.path.isabs(save_dir):
+            save_dir = os.path.abspath(save_dir)
+        
+        history = self.training_histories[model_name]
+        
+        fig, axes = plt.subplots(2, 2, figsize=(15, 12))
+        fig.suptitle(f'{model_name} - Training History', fontsize=16, fontweight='bold')
+        
+        # Loss curves
+        axes[0, 0].plot(history['train_loss'], label='Train Loss', linewidth=2)
+        axes[0, 0].plot(history['val_loss'], label='Val Loss', linewidth=2)
+        axes[0, 0].set_title('Loss Curves')
+        axes[0, 0].set_xlabel('Epoch')
+        axes[0, 0].set_ylabel('Loss')
+        axes[0, 0].legend()
+        axes[0, 0].grid(True, alpha=0.3)
+        
+        # Accuracy curve
+        axes[0, 1].plot(history['val_accuracy'], label='Val Accuracy', 
+                       color='green', linewidth=2)
+        axes[0, 1].set_title('Validation Accuracy')
+        axes[0, 1].set_xlabel('Epoch')
+        axes[0, 1].set_ylabel('Accuracy')
+        axes[0, 1].legend()
+        axes[0, 1].grid(True, alpha=0.3)
+        
+        # F1 Score curve
+        axes[1, 0].plot(history['val_f1'], label='Val F1 Score', 
+                       color='orange', linewidth=2)
+        axes[1, 0].set_title('Validation F1 Score')
+        axes[1, 0].set_xlabel('Epoch')
+        axes[1, 0].set_ylabel('F1 Score')
+        axes[1, 0].legend()
+        axes[1, 0].grid(True, alpha=0.3)
+        
+        # Learning rate
+        axes[1, 1].plot(history['learning_rates'], label='Learning Rate', 
+                       color='red', linewidth=2)
+        axes[1, 1].set_title('Learning Rate Schedule')
+        axes[1, 1].set_xlabel('Epoch')
+        axes[1, 1].set_ylabel('Learning Rate')
+        axes[1, 1].set_yscale('log')
+        axes[1, 1].legend()
+        axes[1, 1].grid(True, alpha=0.3)
+        
+        plt.tight_layout()
+        os.makedirs(save_dir, exist_ok=True)
+        plt.savefig(os.path.join(save_dir, f'{model_name}_training_curves.png'), 
+                   dpi=300, bbox_inches='tight')
+        plt.close()
+        print(f"Training curves saved for {model_name}")
+    
+    def plot_confusion_matrix(self, model_name: str, class_names: List[str] = None,
+                            save_dir: str = "./notebooks"):
+        """Plot confusion matrix."""
+        if model_name not in self.evaluation_results:
+            print(f"No evaluation results found for {model_name}")
+            return
+        
+        # Use absolute path
+        if not os.path.isabs(save_dir):
+            save_dir = os.path.abspath(save_dir)
+        
+        cm = self.evaluation_results[model_name]['confusion_matrix']
+        
+        if class_names is None:
+            class_names = ['HC', 'PD', 'SWEDD', 'PRODROMAL']
+        
+        plt.figure(figsize=(10, 8))
+        sns.heatmap(cm, annot=True, fmt='d', cmap='Blues',
+                   xticklabels=class_names, yticklabels=class_names,
+                   cbar_kws={'label': 'Count'})
+        plt.title(f'{model_name} - Confusion Matrix', fontsize=14, fontweight='bold')
+        plt.ylabel('True Label', fontsize=12)
+        plt.xlabel('Predicted Label', fontsize=12)
+        
+        plt.tight_layout()
+        os.makedirs(save_dir, exist_ok=True)
+        plt.savefig(os.path.join(save_dir, f'{model_name}_confusion_matrix.png'),
+                   dpi=300, bbox_inches='tight')
+        plt.close()
+        print(f"Confusion matrix saved for {model_name}")
+    
+    def save_evaluation_results(self, save_dir: str = "../notebooks"):
+        """Save all evaluation results to file."""
+        os.makedirs(save_dir, exist_ok=True)
+        timestamp = datetime.now().strftime('%Y%m%d_%H%M%S')
+        
+        results_file = os.path.join(save_dir, f'evaluation_results_{timestamp}.txt')
+        
+        with open(results_file, 'w') as f:
+            f.write("=" * 70 + "\n")
+            f.write("Medical Transformer Models - Evaluation Results\n")
+            f.write("=" * 70 + "\n\n")
+            
+            for model_name, results in self.evaluation_results.items():
+                f.write(f"\n{model_name}\n")
+                f.write("-" * 70 + "\n")
+                f.write(f"Accuracy: {results['accuracy']:.4f}\n")
+                f.write(f"F1 Score: {results['f1_score']:.4f}\n")
+                f.write(f"ROC-AUC: {results['roc_auc']:.4f}\n\n")
+                f.write("Classification Report:\n")
+                f.write(results['classification_report'])
+                f.write("\n" + "=" * 70 + "\n")
+        
+        print(f"\nEvaluation results saved to {results_file}")

src/models/multimodal_ml.py

@@ -0,0 +1,560 @@
+"""
+Multimodal Machine Learning approach for Parkinson's disease classification.
+This module combines traditional ML, transformer models, and ensemble methods.
+"""
+
+import numpy as np
+import pandas as pd
+import torch
+import torch.nn as nn
+import joblib
+from sklearn.ensemble import VotingClassifier, StackingClassifier
+from sklearn.linear_model import LogisticRegression
+from sklearn.metrics import accuracy_score, classification_report, confusion_matrix
+from sklearn.model_selection import cross_val_score, StratifiedKFold
+from xgboost import XGBClassifier
+import matplotlib.pyplot as plt
+import seaborn as sns
+from typing import Dict, List, Tuple, Any
+import os
+import warnings
+warnings.filterwarnings('ignore')
+
+try:
+    from .traditional_ml import TraditionalMLModels
+    from .transformer_models import TransformerModels, TabularDataset
+except ImportError:
+    from traditional_ml import TraditionalMLModels
+    from transformer_models import TransformerModels, TabularDataset
+from torch.utils.data import DataLoader
+
+
+class MultimodalEnsemble:
+    """
+    Multimodal ensemble that combines traditional ML and transformer models.
+    """
+    
+    def __init__(self, device: str = None):
+        self.device = device if device else ('cuda' if torch.cuda.is_available() else 'cpu')
+        self.traditional_models = {}
+        self.transformer_models = {}
+        self.ensemble_model = None
+        self.feature_importance = {}
+        self.model_weights = {}
+        
+    def load_traditional_models(self, model_dir: str = "models/saved"):
+        """Load pre-trained traditional ML models."""
+        model_files = {
+            'lightgbm': 'lightgbm_model.joblib',
+            'xgboost': 'xgboost_model.joblib', 
+            'svm': 'svm_model.joblib'
+        }
+        
+        for model_name, filename in model_files.items():
+            model_path = os.path.join(model_dir, filename)
+            if os.path.exists(model_path):
+                self.traditional_models[model_name] = joblib.load(model_path)
+                print(f"Loaded {model_name} model from {model_path}")
+            else:
+                print(f"Warning: {model_name} model not found at {model_path}")
+    
+    def load_transformer_models(self, model_dir: str = "models/saved", input_dim: int = 31, num_classes: int = 4):
+        """Load pre-trained transformer models (3 medical transformers only)."""
+        
+        # Load only the 3 working medical transformers
+        try:
+            try:
+                from .medical_transformers import (
+                    BioMistralClassifier,
+                    ClinicalT5Classifier,
+                    PubMedBERTClassifier,
+                )
+            except ImportError:
+                from medical_transformers import (
+                    BioMistralClassifier,
+                    ClinicalT5Classifier,
+                    PubMedBERTClassifier,
+                )
+            
+            new_model_configs = {
+                'pubmedbert': {
+                    'builder': lambda: PubMedBERTClassifier(
+                        input_dim=input_dim,
+                        num_classes=num_classes,
+                        dropout=0.10,
+                        freeze_bert=False,
+                    ),
+                    'paths': ['pubmedbert_transformer.pth', 'pubmedbert.pth'],
+                },
+                'biogpt': {
+                    'builder': lambda: BioMistralClassifier(
+                        input_dim=input_dim,
+                        num_classes=num_classes,
+                        dropout=0.15,
+                        train_decoder_layers=6,
+                    ),
+                    'paths': ['biogpt_transformer.pth', 'biogpt.pth', 'biomistral.pth'],
+                },
+                'clinical_t5': {
+                    'builder': lambda: ClinicalT5Classifier(
+                        input_dim=input_dim,
+                        num_classes=num_classes,
+                        dropout=0.10,
+                        freeze_encoder=False,
+                    ),
+                    'paths': ['clinical_t5_transformer.pth', 'clinicalt5_transformer.pth', 'clinical_t5.pth'],
+                },
+            }
+            
+            print("Attempting to load new medical transformer models...")
+            for model_name, config in new_model_configs.items():
+                model_path = next(
+                    (
+                        os.path.join(model_dir, candidate)
+                        for candidate in config['paths']
+                        if os.path.exists(os.path.join(model_dir, candidate))
+                    ),
+                    None,
+                )
+                
+                if model_path:
+                    try:
+                        model = config['builder']()
+                        state = torch.load(model_path, map_location=self.device, weights_only=False)
+                        model.load_state_dict(state)
+                        model = model.to(self.device)
+                        model.eval()
+                        self.transformer_models[model_name] = model
+                        print(f"Loaded {model_name} medical transformer model from {model_path}")
+                    except Exception as e:
+                        print(f"Warning: Could not load {model_name}: {e}")
+                else:
+                    print(f"Info: {model_name} medical transformer not found (not yet trained)")
+        
+        except ImportError as e:
+            print(f"Info: medical_transformers module not found: {e}")
+        
+        # Load simple feedforward transformer as 3rd transformer
+        try:
+            try:
+                from .transformer_models import TransformerModels
+            except ImportError:
+                from transformer_models import TransformerModels
+            transformer_trainer = TransformerModels(device=self.device)
+            
+            # Only load feedforward model (skip corrupted legacy transformers)
+            model_name = 'feedforward'
+            model_path = os.path.join(model_dir, f"{model_name}_transformer.pth")
+            
+            print("Attempting to load feedforward transformer model...")
+            if os.path.exists(model_path):
+                try:
+                    model = transformer_trainer.load_model(
+                        'feedforward', model_name, input_dim, num_classes, 
+                        model_dir, hidden_dims=[256, 128, 64], dropout=0.3
+                    )
+                    self.transformer_models[model_name] = model
+                    print(f"Loaded {model_name} transformer model")
+                except Exception as e:
+                    print(f"Warning: Could not load {model_name}: {e}")
+            else:
+                print(f"Info: {model_name} transformer not found")
+        
+        except ImportError as e:
+            print(f"Warning: Could not import transformer_models: {e}")
+        
+        if len(self.transformer_models) > 0:
+            print(f"Successfully loaded {len(self.transformer_models)} transformer model(s)")
+        else:
+            print("WARNING: No transformer models loaded! Ensemble will use traditional models only.")
+    
+    def get_traditional_predictions(self, X):
+        """Get predictions from traditional ML models."""
+        predictions = {}
+        probabilities = {}
+        
+        for model_name, model in self.traditional_models.items():
+            try:
+                pred = model.predict(X)
+                pred_proba = model.predict_proba(X)
+                predictions[model_name] = pred
+                probabilities[model_name] = pred_proba
+            except Exception as e:
+                print(f"Error getting predictions from {model_name}: {e}")
+        
+        return predictions, probabilities
+    
+    def get_transformer_predictions(self, X):
+        """Get predictions from transformer models."""
+        predictions = {}
+        probabilities = {}
+        
+        # Convert to tensor if needed
+        if not isinstance(X, torch.Tensor):
+            # Handle DataFrame conversion
+            if hasattr(X, 'values'):
+                X_vals = X.values
+            else:
+                X_vals = X
+            X_tensor = torch.FloatTensor(X_vals).to(self.device)
+        else:
+            X_tensor = X.to(self.device)
+        
+        for model_name, model in self.transformer_models.items():
+            try:
+                model.eval()
+                with torch.no_grad():
+                    outputs = model(X_tensor)
+                    proba = torch.softmax(outputs, dim=1)
+                    pred = torch.argmax(outputs, dim=1)
+                    
+                    predictions[model_name] = pred.cpu().numpy()
+                    probabilities[model_name] = proba.cpu().numpy()
+            except Exception as e:
+                print(f"Error getting predictions from {model_name}: {e}")
+        
+        return predictions, probabilities
+    
+    def create_ensemble_features(self, X):
+        """Create ensemble features from all models with optimized weights."""
+        # Get predictions from all models
+        trad_preds, trad_probas = self.get_traditional_predictions(X)
+        trans_preds, trans_probas = self.get_transformer_predictions(X)
+        
+        # Define model weights for better performance
+        model_weights = {
+            # Traditional models - higher weights for better performers
+            'lightgbm': 1.5,
+            'xgboost': 1.3,
+            'svm': 1.0,
+            # New medical transformer models - highest weights for specialized medical models
+            'pubmedbert': 2.2,  # Encoder model trained on PubMed abstracts
+            'biomistral': 2.0,  # Decoder model with medical knowledge
+            'clinical_t5': 2.1,  # Encoder-decoder model for clinical tasks
+            # Legacy transformer models (for backward compatibility)
+            'transformer_small': 1.2,
+            'transformer_medium': 1.5,
+            'transformer_large': 1.8,
+            'feedforward': 1.0
+        }
+        
+        # Combine all probability predictions as features with weights
+        ensemble_features = []
+        
+        # Add traditional model probabilities with weights
+        for model_name, proba in trad_probas.items():
+            weight = model_weights.get(model_name, 1.0)
+            ensemble_features.append(proba * weight)
+        
+        # Add transformer model probabilities with weights
+        for model_name, proba in trans_probas.items():
+            weight = model_weights.get(model_name, 1.0)
+            ensemble_features.append(proba * weight)
+        
+        # Add original features (scaled down)
+        if hasattr(X, 'values'):
+            X_vals = X.values
+        else:
+            X_vals = X
+        ensemble_features.append(X_vals * 0.15)  # Slightly increase original feature weight
+        
+        # Concatenate all features
+        if ensemble_features:
+            combined = np.concatenate(ensemble_features, axis=1)
+        else:
+            combined = X_vals
+
+        # Hardening: keep inference feature width compatible with fitted ensemble model
+        expected = getattr(self.ensemble_model, "n_features_in_", None)
+        if expected is not None and combined.shape[1] != expected:
+            if combined.shape[1] < expected:
+                pad = np.zeros((combined.shape[0], expected - combined.shape[1]), dtype=combined.dtype)
+                combined = np.concatenate([combined, pad], axis=1)
+            else:
+                combined = combined[:, :expected]
+
+        return combined
+    
+    def train_ensemble(self, X_train, y_train, ensemble_type: str = 'stacking'):
+        """Train ensemble model on predictions from base models."""
+        print(f"Training {ensemble_type} ensemble...")
+        
+        # Create ensemble features
+        ensemble_features = self.create_ensemble_features(X_train)
+        print(f"Ensemble features shape: {ensemble_features.shape}")
+        
+        if ensemble_type == 'stacking':
+            # Use XGBoost as meta-learner for better performance
+            self.ensemble_model = XGBClassifier(
+                n_estimators=200,
+                learning_rate=0.05,
+                max_depth=5,
+                min_child_weight=2,
+                gamma=0.1,
+                subsample=0.8,
+                colsample_bytree=0.8,
+                objective='multi:softproba',
+                random_state=42,
+                use_label_encoder=False,
+                eval_metric='mlogloss'
+            )
+        elif ensemble_type == 'voting':
+            # Create voting classifier (if we have sklearn-compatible models)
+            available_models = []
+            for name, model in self.traditional_models.items():
+                available_models.append((name, model))
+            
+            if available_models:
+                self.ensemble_model = VotingClassifier(
+                    estimators=available_models,
+                    voting='soft'
+                )
+            else:
+                print("No traditional models available for voting ensemble")
+                return
+        
+        # Train ensemble model
+        self.ensemble_model.fit(ensemble_features, y_train)
+        print(f"{ensemble_type.capitalize()} ensemble trained successfully")
+    
+    def predict_ensemble(self, X):
+        """Make predictions using the ensemble model."""
+        if self.ensemble_model is None:
+            raise ValueError("Ensemble model not trained yet")
+        
+        ensemble_features = self.create_ensemble_features(X)
+        predictions = self.ensemble_model.predict(ensemble_features)
+        probabilities = self.ensemble_model.predict_proba(ensemble_features)
+        
+        return predictions, probabilities
+    
+    def evaluate_ensemble(self, X_test, y_test):
+        """Evaluate ensemble model performance."""
+        predictions, probabilities = self.predict_ensemble(X_test)
+        
+        accuracy = accuracy_score(y_test, predictions)
+        report = classification_report(y_test, predictions)
+        cm = confusion_matrix(y_test, predictions)
+        
+        return {
+            'accuracy': accuracy,
+            'predictions': predictions,
+            'probabilities': probabilities,
+            'classification_report': report,
+            'confusion_matrix': cm
+        }
+    
+    def compare_all_models(self, X_test, y_test):
+        """Compare performance of all individual models and ensemble."""
+        results = {}
+        
+        # Evaluate traditional models
+        trad_preds, trad_probas = self.get_traditional_predictions(X_test)
+        for model_name, pred in trad_preds.items():
+            accuracy = accuracy_score(y_test, pred)
+            results[f"Traditional_{model_name}"] = accuracy
+        
+        # Evaluate transformer models
+        trans_preds, trans_probas = self.get_transformer_predictions(X_test)
+        for model_name, pred in trans_preds.items():
+            accuracy = accuracy_score(y_test, pred)
+            results[f"Transformer_{model_name}"] = accuracy
+        
+        # Evaluate ensemble
+        if self.ensemble_model is not None:
+            ensemble_results = self.evaluate_ensemble(X_test, y_test)
+            results["Ensemble"] = ensemble_results['accuracy']
+        
+        return results
+    
+    def plot_model_comparison(self, results: Dict, save_path: str = "notebooks/multimodal_comparison.png"):
+        """Plot comparison of all models."""
+        models = list(results.keys())
+        accuracies = list(results.values())
+        
+        plt.figure(figsize=(15, 8))
+        bars = plt.bar(models, accuracies, color=['skyblue', 'lightgreen', 'lightcoral', 
+                                                 'gold', 'pink', 'lightgray', 'orange', 'red'])
+        
+        plt.title('Multimodal Model Comparison - Test Accuracy', fontsize=16)
+        plt.xlabel('Model', fontsize=12)
+        plt.ylabel('Accuracy', fontsize=12)
+        plt.xticks(rotation=45, ha='right')
+        
+        # Add value labels on bars
+        for bar, acc in zip(bars, accuracies):
+            plt.text(bar.get_x() + bar.get_width()/2, bar.get_height() + 0.005,
+                    f'{acc:.4f}', ha='center', va='bottom', fontsize=10)
+        
+        plt.tight_layout()
+        plt.savefig(save_path, dpi=300, bbox_inches='tight')
+        plt.close()
+        
+        print(f"Model comparison plot saved to {save_path}")
+    
+    def cross_validate_ensemble(self, X, y, cv_folds: int = 5):
+        """Perform cross-validation on ensemble model."""
+        if self.ensemble_model is None:
+            raise ValueError("Ensemble model not trained yet")
+        
+        # Create ensemble features for full dataset
+        ensemble_features = self.create_ensemble_features(X)
+        
+        # Perform cross-validation
+        cv_scores = cross_val_score(
+            self.ensemble_model, ensemble_features, y, 
+            cv=cv_folds, scoring='accuracy'
+        )
+        
+        return cv_scores
+    
+    def save_ensemble(self, save_path: str = "models/saved/multimodal_ensemble.joblib"):
+        """Save the trained ensemble model."""
+        if self.ensemble_model is not None:
+            joblib.dump(self.ensemble_model, save_path)
+            print(f"Ensemble model saved to {save_path}")
+        else:
+            print("No ensemble model to save")
+    
+    def load_ensemble(self, load_path: str = "models/saved/multimodal_ensemble.joblib"):
+        """Load a pre-trained ensemble model."""
+        if os.path.exists(load_path):
+            self.ensemble_model = joblib.load(load_path)
+            print(f"Ensemble model loaded from {load_path}")
+        else:
+            print(f"Ensemble model not found at {load_path}")
+
+
+class AdvancedFeatureEngineering:
+    """
+    Advanced feature engineering for multimodal approach.
+    """
+    
+    def __init__(self):
+        self.feature_transformers = {}
+        self.interaction_features = []
+    
+    def create_polynomial_features(self, X, degree: int = 2, feature_subset: List[str] = None):
+        """Create polynomial features for selected columns."""
+        from sklearn.preprocessing import PolynomialFeatures
+        
+        if feature_subset is None:
+            # Use numerical features only
+            numerical_cols = X.select_dtypes(include=[np.number]).columns
+            feature_subset = numerical_cols[:5]  # Limit to avoid explosion
+        
+        poly = PolynomialFeatures(degree=degree, include_bias=False, interaction_only=True)
+        X_subset = X[feature_subset]
+        X_poly = poly.fit_transform(X_subset)
+        
+        # Create feature names
+        feature_names = poly.get_feature_names_out(feature_subset)
+        
+        # Add to original dataframe
+        X_enhanced = X.copy()
+        for i, name in enumerate(feature_names):
+            if name not in X.columns:  # Avoid duplicates
+                X_enhanced[f'poly_{name}'] = X_poly[:, i]
+        
+        self.feature_transformers['polynomial'] = poly
+        return X_enhanced
+    
+    def create_clustering_features(self, X, n_clusters: int = 5):
+        """Create clustering-based features."""
+        from sklearn.cluster import KMeans
+        from sklearn.preprocessing import StandardScaler
+        
+        # Scale features for clustering
+        scaler = StandardScaler()
+        X_scaled = scaler.fit_transform(X.select_dtypes(include=[np.number]))
+        
+        # Apply K-means clustering
+        kmeans = KMeans(n_clusters=n_clusters, random_state=42)
+        cluster_labels = kmeans.fit_predict(X_scaled)
+        
+        # Add cluster features
+        X_enhanced = X.copy()
+        X_enhanced['cluster_label'] = cluster_labels
+        
+        # Add distance to each cluster center
+        distances = kmeans.transform(X_scaled)
+        for i in range(n_clusters):
+            X_enhanced[f'dist_to_cluster_{i}'] = distances[:, i]
+        
+        self.feature_transformers['clustering'] = {'kmeans': kmeans, 'scaler': scaler}
+        return X_enhanced
+    
+    def create_statistical_features(self, X):
+        """Create statistical aggregation features."""
+        X_enhanced = X.copy()
+        numerical_cols = X.select_dtypes(include=[np.number]).columns
+        
+        if len(numerical_cols) > 1:
+            # Row-wise statistics
+            X_enhanced['row_mean'] = X[numerical_cols].mean(axis=1)
+            X_enhanced['row_std'] = X[numerical_cols].std(axis=1)
+            X_enhanced['row_min'] = X[numerical_cols].min(axis=1)
+            X_enhanced['row_max'] = X[numerical_cols].max(axis=1)
+            X_enhanced['row_range'] = X_enhanced['row_max'] - X_enhanced['row_min']
+            X_enhanced['row_skew'] = X[numerical_cols].skew(axis=1)
+        
+        return X_enhanced
+
+
+def create_multimodal_pipeline(X_train, X_test, y_train, y_test):
+    """
+    Create and evaluate a complete multimodal ML pipeline.
+    """
+    print("Creating Multimodal ML Pipeline...")
+    print("=" * 50)
+    
+    # Initialize multimodal ensemble
+    ensemble = MultimodalEnsemble()
+    
+    # Load pre-trained models
+    print("Loading pre-trained models...")
+    ensemble.load_traditional_models()
+    ensemble.load_transformer_models(input_dim=X_train.shape[1])
+    
+    # Advanced feature engineering
+    print("Applying advanced feature engineering...")
+    feature_engineer = AdvancedFeatureEngineering()
+    
+    # Create enhanced features
+    X_train_enhanced = feature_engineer.create_polynomial_features(X_train)
+    X_train_enhanced = feature_engineer.create_clustering_features(X_train_enhanced)
+    X_train_enhanced = feature_engineer.create_statistical_features(X_train_enhanced)
+    
+    # Apply same transformations to test set
+    X_test_enhanced = X_test.copy()
+    if 'polynomial' in feature_engineer.feature_transformers:
+        poly = feature_engineer.feature_transformers['polynomial']
+        # Apply polynomial features to test set (implementation needed)
+    
+    print(f"Enhanced training features shape: {X_train_enhanced.shape}")
+    
+    # Train ensemble models
+    ensemble.train_ensemble(X_train, y_train, ensemble_type='stacking')
+    
+    # Evaluate all models
+    print("\nEvaluating all models...")
+    results = ensemble.compare_all_models(X_test, y_test)
+    
+    # Print results
+    print("\nModel Performance Comparison:")
+    print("-" * 40)
+    for model_name, accuracy in sorted(results.items(), key=lambda x: x[1], reverse=True):
+        print(f"{model_name:<25}: {accuracy:.4f}")
+    
+    # Cross-validation
+    print("\nPerforming cross-validation on ensemble...")
+    cv_scores = ensemble.cross_validate_ensemble(X_train, y_train)
+    print(f"CV Accuracy: {cv_scores.mean():.4f} (+/- {cv_scores.std() * 2:.4f})")
+    
+    # Save ensemble model
+    ensemble.save_ensemble()
+    
+    # Create visualizations
+    ensemble.plot_model_comparison(results)
+    
+    return ensemble, results

src/models/traditional_ml.py

@@ -0,0 +1,158 @@
+import numpy as np
+import pandas as pd
+from sklearn.metrics import accuracy_score, classification_report, confusion_matrix
+from sklearn.model_selection import GridSearchCV
+from lightgbm import LGBMClassifier
+from xgboost import XGBClassifier
+from sklearn.svm import SVC
+import joblib
+import os
+
+class TraditionalMLModels:
+    def __init__(self, save_dir='../models/saved'):
+        """Initialize the traditional ML models class.
+        
+        Args:
+            save_dir (str): Directory to save trained models
+        """
+        self.models = {
+            'lightgbm': None,
+            'xgboost': None,
+            'svm': None
+        }
+        self.save_dir = save_dir
+        os.makedirs(save_dir, exist_ok=True)
+        
+    def train_lightgbm(self, X_train, y_train):
+        """Train LightGBM model with optimized parameters."""
+        print("Training LightGBM...")
+        
+        # Use pre-configured model if provided, otherwise create with good defaults
+        if self.models['lightgbm'] is None:
+            self.models['lightgbm'] = LGBMClassifier(
+                random_state=42,
+                n_estimators=200,
+                learning_rate=0.05,
+                max_depth=7,
+                num_leaves=31,
+                objective='multiclass',
+                num_class=4,
+                verbose=-1
+            )
+        
+        # Train the model
+        self.models['lightgbm'].fit(X_train, y_train)
+        
+        # Return train score as approximation
+        train_score = self.models['lightgbm'].score(X_train, y_train)
+        print(f"LightGBM training accuracy: {train_score:.4f}")
+        return train_score
+    
+    def train_xgboost(self, X_train, y_train):
+        """Train XGBoost model with optimized parameters."""
+        print("Training XGBoost...")
+        
+        # Use pre-configured model if provided, otherwise create with good defaults
+        if self.models['xgboost'] is None:
+            self.models['xgboost'] = XGBClassifier(
+                random_state=42,
+                n_estimators=200,
+                learning_rate=0.05,
+                max_depth=6,
+                objective='multi:softmax',
+                num_class=4,
+                eval_metric='mlogloss',
+                verbosity=0
+            )
+        
+        # Train the model
+        self.models['xgboost'].fit(X_train, y_train)
+        
+        # Return train score as approximation
+        train_score = self.models['xgboost'].score(X_train, y_train)
+        print(f"XGBoost training accuracy: {train_score:.4f}")
+        return train_score
+    
+    def train_svm(self, X_train, y_train):
+        """Train SVM model with optimized parameters."""
+        print("Training SVM...")
+        
+        # Use pre-configured model if provided, otherwise create with good defaults
+        if self.models['svm'] is None:
+            self.models['svm'] = SVC(
+                random_state=42,
+                probability=True,
+                C=10.0,
+                kernel='rbf',
+                gamma='scale',
+                decision_function_shape='ovr',
+                verbose=False
+            )
+        
+        # Train the model
+        self.models['svm'].fit(X_train, y_train)
+        
+        # Return train score as approximation
+        train_score = self.models['svm'].score(X_train, y_train)
+        print(f"SVM training accuracy: {train_score:.4f}")
+        return train_score
+    
+    def train_all_models(self, X_train, y_train):
+        """Train all models and return their cross-validation scores."""
+        scores = {}
+        scores['lightgbm'] = self.train_lightgbm(X_train, y_train)
+        scores['xgboost'] = self.train_xgboost(X_train, y_train)
+        scores['svm'] = self.train_svm(X_train, y_train)
+        return scores
+    
+    def evaluate_model(self, model_name, X_test, y_test):
+        """Evaluate a specific model on test data."""
+        model = self.models[model_name]
+        if model is None:
+            print(f"Model '{model_name}' is not trained yet.")
+            return None
+
+        y_pred = model.predict(X_test)
+        probabilities = None
+        if hasattr(model, "predict_proba"):
+            probabilities = model.predict_proba(X_test)
+        accuracy = accuracy_score(y_test, y_pred)
+        report = classification_report(y_test, y_pred)
+        cm = confusion_matrix(y_test, y_pred)
+
+        print(f"\n{model_name} Test Accuracy: {accuracy:.4f}")
+        print(f"Classification Report:\n{report}")
+
+        return {
+            'accuracy': accuracy,
+            'predictions': y_pred,
+            'probabilities': probabilities,
+            'classification_report': report,
+            'confusion_matrix': cm
+        }
+
+    def evaluate_all_models(self, X_test, y_test):
+        """Evaluate all trained models on test data."""
+        results = {}
+        for name in self.models:
+            if self.models[name] is not None:
+                results[name] = self.evaluate_model(name, X_test, y_test)
+        return results
+    
+    def save_models(self):
+        """Save all trained models to disk."""
+        for model_name, model in self.models.items():
+            if model is not None:
+                save_path = os.path.join(self.save_dir, f"{model_name}_model.joblib")
+                joblib.dump(model, save_path)
+                print(f"Saved {model_name} model to {save_path}")
+    
+    def load_models(self):
+        """Load all saved models from disk."""
+        for model_name in self.models:
+            model_path = os.path.join(self.save_dir, f"{model_name}_model.joblib")
+            if os.path.exists(model_path):
+                self.models[model_name] = joblib.load(model_path)
+                print(f"Loaded {model_name} model from {model_path}")
+            else:
+                print(f"No saved model found for {model_name}")

src/models/transformer_models.py

@@ -0,0 +1,509 @@
+"""
+Transformer-based models for Parkinson's disease classification.
+This module implements pretrained transformer models (DistilBERT, BioBERT, PubMedBERT) 
+adapted for tabular data with RAG integration.
+"""
+
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+from torch.utils.data import Dataset, DataLoader
+import numpy as np
+from sklearn.metrics import accuracy_score, classification_report, confusion_matrix
+from sklearn.model_selection import StratifiedKFold
+import joblib
+import os
+from typing import Tuple, Dict, Any, List, Optional
+import matplotlib.pyplot as plt
+import seaborn as sns
+from transformers import (
+    AutoModel, 
+    AutoTokenizer, 
+    DistilBertModel, 
+    DistilBertTokenizer,
+    BertModel,
+    BertTokenizer
+)
+
+
+class TabularDataset(Dataset):
+    """Custom dataset for tabular data optionally carrying cached RAG contexts."""
+    
+    def __init__(self, X, y, feature_names=None, contexts=None):
+        self.X = torch.FloatTensor(X)
+        self.y = torch.LongTensor(y)
+        self.feature_names = feature_names or [f"feature_{i}" for i in range(X.shape[1])]
+        context_len = None
+        if isinstance(contexts, dict):
+            first_value = next(iter(contexts.values()), None)
+            context_len = len(first_value) if first_value is not None else 0
+        elif contexts is not None:
+            context_len = len(contexts)
+        if contexts is not None and context_len != len(self.X):
+            raise ValueError("Length of contexts must match number of samples")
+        self.contexts = contexts
+    
+    def __len__(self):
+        return len(self.X)
+    
+    def __getitem__(self, idx):
+        if self.contexts is None:
+            return self.X[idx], self.y[idx]
+        if isinstance(self.contexts, dict):
+            return self.X[idx], self.y[idx], {key: value[idx] for key, value in self.contexts.items()}
+        return self.X[idx], self.y[idx], self.contexts[idx]
+    
+    def get_feature_description(self, idx):
+        """Get feature values with names for text representation."""
+        sample = self.X[idx].numpy()
+        return {name: float(val) for name, val in zip(self.feature_names, sample)}
+
+
+class DistilBERTForTabular(nn.Module):
+    """DistilBERT model adapted for tabular data with RAG integration."""
+    
+    def __init__(self, input_dim: int, num_classes: int, dropout: float = 0.1):
+        super(DistilBERTForTabular, self).__init__()
+        
+        self.input_dim = input_dim
+        self.num_classes = num_classes
+        self.model_name = "distilbert-base-uncased"
+        
+        # Load pretrained DistilBERT model and tokenizer
+        self.tokenizer = DistilBertTokenizer.from_pretrained(self.model_name)
+        self.bert = DistilBertModel.from_pretrained(self.model_name)
+        
+        # Freeze BERT parameters to speed up training
+        for param in self.bert.parameters():
+            param.requires_grad = False
+            
+        # Feature projection layer
+        self.feature_projection = nn.Linear(input_dim, 768)  # Project to BERT hidden size
+        
+        # Classification head
+        self.classifier = nn.Sequential(
+            nn.Linear(768, 256),
+            nn.ReLU(),
+            nn.Dropout(dropout),
+            nn.Linear(256, num_classes)
+        )
+        
+    def forward(self, x, text_input=None):
+        # Project tabular features
+        tabular_features = self.feature_projection(x)
+        
+        if text_input is not None:
+            # Process text input if available (for RAG integration)
+            inputs = self.tokenizer(text_input, return_tensors="pt", padding=True, truncation=True, max_length=512)
+            inputs = {k: v.to(x.device) for k, v in inputs.items()}
+            
+            # Get BERT embeddings
+            with torch.no_grad():
+                outputs = self.bert(**inputs)
+                text_features = outputs.last_hidden_state[:, 0, :]  # CLS token
+                
+            # Combine tabular and text features
+            combined_features = tabular_features + text_features
+        else:
+            # Use only tabular features if no text input
+            combined_features = tabular_features
+        
+        # Classification
+        output = self.classifier(combined_features)
+        
+        return output
+
+
+class BioBERTForTabular(nn.Module):
+    """BioBERT model adapted for tabular data with RAG integration."""
+    
+    def __init__(self, input_dim: int, num_classes: int, dropout: float = 0.1):
+        super(BioBERTForTabular, self).__init__()
+        
+        self.input_dim = input_dim
+        self.num_classes = num_classes
+        self.model_name = "dmis-lab/biobert-v1.1"
+        
+        # Load pretrained BioBERT model and tokenizer
+        self.tokenizer = AutoTokenizer.from_pretrained(self.model_name)
+        self.bert = AutoModel.from_pretrained(self.model_name)
+        
+        # Freeze BERT parameters to speed up training
+        for param in self.bert.parameters():
+            param.requires_grad = False
+            
+        # Feature projection layer
+        self.feature_projection = nn.Linear(input_dim, 768)  # Project to BERT hidden size
+        
+        # Classification head
+        self.classifier = nn.Sequential(
+            nn.Linear(768, 256),
+            nn.ReLU(),
+            nn.Dropout(dropout),
+            nn.Linear(256, num_classes)
+        )
+        
+    def forward(self, x, text_input=None):
+        # Project tabular features
+        tabular_features = self.feature_projection(x)
+        
+        if text_input is not None:
+            # Process text input if available (for RAG integration)
+            inputs = self.tokenizer(text_input, return_tensors="pt", padding=True, truncation=True, max_length=512)
+            inputs = {k: v.to(x.device) for k, v in inputs.items()}
+            
+            # Get BERT embeddings
+            with torch.no_grad():
+                outputs = self.bert(**inputs)
+                text_features = outputs.last_hidden_state[:, 0, :]  # CLS token
+                
+            # Combine tabular and text features
+            combined_features = tabular_features + text_features
+        else:
+            # Use only tabular features if no text input
+            combined_features = tabular_features
+        
+        # Classification
+        output = self.classifier(combined_features)
+        
+        return output
+
+
+class PubMedBERTForTabular(nn.Module):
+    """PubMedBERT model adapted for tabular data with RAG integration."""
+    
+    def __init__(self, input_dim: int, num_classes: int, dropout: float = 0.1):
+        super(PubMedBERTForTabular, self).__init__()
+        
+        self.input_dim = input_dim
+        self.num_classes = num_classes
+        self.model_name = "microsoft/BiomedNLP-PubMedBERT-base-uncased-abstract-fulltext"
+        
+        # Load pretrained PubMedBERT model and tokenizer
+        self.tokenizer = AutoTokenizer.from_pretrained(self.model_name)
+        self.bert = AutoModel.from_pretrained(self.model_name)
+        
+        # Freeze BERT parameters to speed up training
+        for param in self.bert.parameters():
+            param.requires_grad = False
+            
+        # Feature projection layer
+        self.feature_projection = nn.Linear(input_dim, 768)  # Project to BERT hidden size
+        
+        # Classification head
+        self.classifier = nn.Sequential(
+            nn.Linear(768, 256),
+            nn.ReLU(),
+            nn.Dropout(dropout),
+            nn.Linear(256, num_classes)
+        )
+        
+    def forward(self, x, text_input=None):
+        # Project tabular features
+        tabular_features = self.feature_projection(x)
+        
+        if text_input is not None:
+            # Process text input if available (for RAG integration)
+            inputs = self.tokenizer(text_input, return_tensors="pt", padding=True, truncation=True, max_length=512)
+            inputs = {k: v.to(x.device) for k, v in inputs.items()}
+            
+            # Get BERT embeddings
+            with torch.no_grad():
+                outputs = self.bert(**inputs)
+                text_features = outputs.last_hidden_state[:, 0, :]  # CLS token
+                
+            # Combine tabular and text features
+            combined_features = tabular_features + text_features
+        else:
+            # Use only tabular features if no text input
+            combined_features = tabular_features
+        
+        # Classification
+        output = self.classifier(combined_features)
+        
+        return output
+
+
+class FeedForwardNetwork(nn.Module):
+    """Simple feed-forward network for comparison."""
+    
+    def __init__(self, input_dim: int, num_classes: int, hidden_dims: list = [256, 128, 64], 
+                 dropout: float = 0.3):
+        super(FeedForwardNetwork, self).__init__()
+        
+        layers = []
+        prev_dim = input_dim
+        
+        for hidden_dim in hidden_dims:
+            layers.extend([
+                nn.Linear(prev_dim, hidden_dim),
+                nn.BatchNorm1d(hidden_dim),
+                nn.ReLU(),
+                nn.Dropout(dropout)
+            ])
+            prev_dim = hidden_dim
+        
+        layers.append(nn.Linear(prev_dim, num_classes))
+        
+        self.network = nn.Sequential(*layers)
+    
+    def forward(self, x):
+        return self.network(x)
+
+
+class TransformerModels:
+    """Class to handle training and evaluation of transformer models."""
+    
+    def __init__(self, device: str = None):
+        self.device = device if device else ('cuda' if torch.cuda.is_available() else 'cpu')
+        self.models = {}
+        self.training_history = {}
+        
+    def create_model(self, model_type: str, input_dim: int, num_classes: int, **kwargs):
+        """Create a model of specified type."""
+        if model_type == 'transformer':
+            # For backward compatibility with saved models
+            if 'd_model' in kwargs:
+                # This is the old TabularTransformer model structure
+                d_model = kwargs.get('d_model', 256)
+                nhead = kwargs.get('nhead', 8)
+                num_layers = kwargs.get('num_layers', 4)
+                dropout = kwargs.get('dropout', 0.1)
+                
+                # Create a compatible model structure
+                model = FeedForwardNetwork(input_dim, num_classes, 
+                                          hidden_dims=[d_model, d_model//2, d_model//4], 
+                                          dropout=dropout)
+            else:
+                # New transformer models
+                model_name = kwargs.get('model_name', 'distilbert')
+                if model_name == 'distilbert':
+                    model = DistilBERTForTabular(input_dim, num_classes, dropout=kwargs.get('dropout', 0.1))
+                elif model_name == 'biobert':
+                    model = BioBERTForTabular(input_dim, num_classes, dropout=kwargs.get('dropout', 0.1))
+                elif model_name == 'pubmedbert':
+                    model = PubMedBERTForTabular(input_dim, num_classes, dropout=kwargs.get('dropout', 0.1))
+                else:
+                    # Default to DistilBERT if model name not specified
+                    model = DistilBERTForTabular(input_dim, num_classes, dropout=kwargs.get('dropout', 0.1))
+        elif model_type == 'feedforward':
+            model = FeedForwardNetwork(input_dim, num_classes, **kwargs)
+        else:
+            raise ValueError(f"Unknown model type: {model_type}")
+        
+        return model.to(self.device)
+    
+    def train_model(self, model, train_loader, val_loader, epochs: int = 100, 
+                   lr: float = 0.001, weight_decay: float = 1e-4, 
+                   class_weights: torch.Tensor = None):
+        """Train a model."""
+        
+        # Loss function with class weights
+        if class_weights is not None:
+            criterion = nn.CrossEntropyLoss(weight=class_weights.to(self.device))
+        else:
+            criterion = nn.CrossEntropyLoss()
+        
+        optimizer = torch.optim.AdamW(model.parameters(), lr=lr, weight_decay=weight_decay)
+        scheduler = torch.optim.lr_scheduler.ReduceLROnPlateau(
+            optimizer, mode='min', factor=0.5, patience=10
+        )
+        
+        train_losses = []
+        val_losses = []
+        val_accuracies = []
+        
+        best_val_loss = float('inf')
+        patience_counter = 0
+        early_stopping_patience = 20
+        
+        for epoch in range(epochs):
+            # Training phase
+            model.train()
+            train_loss = 0.0
+            
+            for batch_x, batch_y in train_loader:
+                batch_x, batch_y = batch_x.to(self.device), batch_y.to(self.device)
+                
+                optimizer.zero_grad()
+                outputs = model(batch_x)
+                loss = criterion(outputs, batch_y)
+                loss.backward()
+                optimizer.step()
+                
+                train_loss += loss.item()
+            
+            # Validation phase
+            model.eval()
+            val_loss = 0.0
+            correct = 0
+            total = 0
+            
+            with torch.no_grad():
+                for batch_x, batch_y in val_loader:
+                    batch_x, batch_y = batch_x.to(self.device), batch_y.to(self.device)
+                    
+                    outputs = model(batch_x)
+                    loss = criterion(outputs, batch_y)
+                    val_loss += loss.item()
+                    
+                    _, predicted = torch.max(outputs.data, 1)
+                    total += batch_y.size(0)
+                    correct += (predicted == batch_y).sum().item()
+            
+            train_loss /= len(train_loader)
+            val_loss /= len(val_loader)
+            val_accuracy = 100 * correct / total
+            
+            train_losses.append(train_loss)
+            val_losses.append(val_loss)
+            val_accuracies.append(val_accuracy)
+            
+            scheduler.step(val_loss)
+            
+            if epoch % 10 == 0:
+                print(f'Epoch [{epoch}/{epochs}], Train Loss: {train_loss:.4f}, '
+                      f'Val Loss: {val_loss:.4f}, Val Acc: {val_accuracy:.2f}%')
+            
+            # Early stopping
+            if val_loss < best_val_loss:
+                best_val_loss = val_loss
+                patience_counter = 0
+            else:
+                patience_counter += 1
+                if patience_counter >= early_stopping_patience:
+                    print(f'Early stopping at epoch {epoch}')
+                    break
+        
+        return {
+            'train_losses': train_losses,
+            'val_losses': val_losses,
+            'val_accuracies': val_accuracies
+        }
+    
+    def evaluate_model(self, model, test_loader):
+        """Evaluate model on test set."""
+        model.eval()
+        all_predictions = []
+        all_targets = []
+        
+        with torch.no_grad():
+            for batch_x, batch_y in test_loader:
+                batch_x, batch_y = batch_x.to(self.device), batch_y.to(self.device)
+                
+                outputs = model(batch_x)
+                _, predicted = torch.max(outputs, 1)
+                
+                all_predictions.extend(predicted.cpu().numpy())
+                all_targets.extend(batch_y.cpu().numpy())
+        
+        accuracy = accuracy_score(all_targets, all_predictions)
+        report = classification_report(all_targets, all_predictions)
+        cm = confusion_matrix(all_targets, all_predictions)
+        
+        return {
+            'accuracy': accuracy,
+            'predictions': all_predictions,
+            'targets': all_targets,
+            'classification_report': report,
+            'confusion_matrix': cm
+        }
+    
+    def cross_validate(self, model_type: str, X, y, cv_folds: int = 5, **model_kwargs):
+        """Perform cross-validation."""
+        skf = StratifiedKFold(n_splits=cv_folds, shuffle=True, random_state=42)
+        cv_scores = []
+        
+        for fold, (train_idx, val_idx) in enumerate(skf.split(X, y)):
+            print(f"\nFold {fold + 1}/{cv_folds}")
+            
+            X_train_fold, X_val_fold = X[train_idx], X[val_idx]
+            y_train_fold, y_val_fold = y[train_idx], y[val_idx]
+            
+            # Create datasets and loaders
+            train_dataset = TabularDataset(X_train_fold, y_train_fold)
+            val_dataset = TabularDataset(X_val_fold, y_val_fold)
+            
+            train_loader = DataLoader(train_dataset, batch_size=64, shuffle=True)
+            val_loader = DataLoader(val_dataset, batch_size=64, shuffle=False)
+            
+            # Create and train model
+            model = self.create_model(model_type, X.shape[1], len(np.unique(y)), **model_kwargs)
+            
+            # Calculate class weights for this fold
+            class_counts = np.bincount(y_train_fold)
+            class_weights = torch.FloatTensor(len(class_counts) / (len(class_counts) * class_counts))
+            
+            # Train model
+            history = self.train_model(model, train_loader, val_loader, 
+                                     epochs=50, class_weights=class_weights)
+            
+            # Evaluate on validation set
+            results = self.evaluate_model(model, val_loader)
+            cv_scores.append(results['accuracy'])
+            
+            print(f"Fold {fold + 1} Accuracy: {results['accuracy']:.4f}")
+        
+        return cv_scores
+    
+    def save_model(self, model, model_name: str, save_dir: str = "models/saved"):
+        """Save trained model."""
+        os.makedirs(save_dir, exist_ok=True)
+        model_path = os.path.join(save_dir, f"{model_name}_transformer.pth")
+        torch.save(model.state_dict(), model_path)
+        print(f"Model saved to {model_path}")
+    
+    def load_model(self, model_type: str, model_name: str, input_dim: int, 
+                   num_classes: int, save_dir: str = "models/saved", **model_kwargs):
+        """Load trained model."""
+        model_path = os.path.join(save_dir, f"{model_name}_transformer.pth")
+        
+        model = self.create_model(model_type, input_dim, num_classes, **model_kwargs)
+        model.load_state_dict(torch.load(model_path, map_location=self.device))
+        model.eval()
+        
+        return model
+    
+    def plot_training_history(self, history: Dict, model_name: str, save_dir: str = "notebooks"):
+        """Plot training history."""
+        fig, (ax1, ax2) = plt.subplots(1, 2, figsize=(15, 5))
+        
+        # Plot losses
+        ax1.plot(history['train_losses'], label='Train Loss')
+        ax1.plot(history['val_losses'], label='Validation Loss')
+        ax1.set_title(f'{model_name} - Training and Validation Loss')
+        ax1.set_xlabel('Epoch')
+        ax1.set_ylabel('Loss')
+        ax1.legend()
+        ax1.grid(True)
+        
+        # Plot accuracy
+        ax2.plot(history['val_accuracies'], label='Validation Accuracy', color='green')
+        ax2.set_title(f'{model_name} - Validation Accuracy')
+        ax2.set_xlabel('Epoch')
+        ax2.set_ylabel('Accuracy (%)')
+        ax2.legend()
+        ax2.grid(True)
+        
+        plt.tight_layout()
+        plt.savefig(os.path.join(save_dir, f'{model_name}_training_history.png'), dpi=300, bbox_inches='tight')
+        plt.close()
+    
+    def plot_confusion_matrix(self, cm, model_name: str, class_names: list = None, 
+                            save_dir: str = "notebooks"):
+        """Plot confusion matrix."""
+        plt.figure(figsize=(8, 6))
+        
+        if class_names is None:
+            class_names = [f'Class {i}' for i in range(len(cm))]
+        
+        sns.heatmap(cm, annot=True, fmt='d', cmap='Blues', 
+                   xticklabels=class_names, yticklabels=class_names)
+        plt.title(f'{model_name} - Confusion Matrix')
+        plt.xlabel('Predicted')
+        plt.ylabel('Actual')
+        
+        plt.tight_layout()
+        plt.savefig(os.path.join(save_dir, f'{model_name}_transformer_confusion_matrix.png'), 
+                   dpi=300, bbox_inches='tight')
+        plt.close()

src/progression_model.py

@@ -0,0 +1,338 @@
+"""
+Progression Modeling — PPMI Trajectory Clustering.
+
+Extracts longitudinal PPMI visit data, clusters patients into
+slow / moderate / fast progressors using k-means on velocity
+vectors, and provides cluster-weighted forecast predictions.
+"""
+
+from __future__ import annotations
+
+import logging
+import math
+from pathlib import Path
+from typing import Any, Dict, List, Optional, Tuple
+
+import numpy as np
+
+logger = logging.getLogger(__name__)
+
+# ---------------------------------------------------------------------------
+# Cluster labels in order of severity (index 0 = slowest)
+# ---------------------------------------------------------------------------
+CLUSTER_LABELS = ["slow", "moderate", "fast"]
+
+# Yearly progression defaults per cluster (UPDRS3 gain/yr, MoCA loss/yr, HY gain/yr)
+DEFAULT_CLUSTER_PROFILES = {
+    "slow":     {"updrs3_gain": 1.8, "moca_loss": 0.3, "hy_gain": 0.10},
+    "moderate": {"updrs3_gain": 3.5, "moca_loss": 0.7, "hy_gain": 0.25},
+    "fast":     {"updrs3_gain": 6.5, "moca_loss": 1.3, "hy_gain": 0.45},
+}
+
+
+class ProgressionModel:
+    """PPMI-trained trajectory clustering for Parkinson's progression."""
+
+    def __init__(self) -> None:
+        self.fitted = False
+        self.centroids: Optional[np.ndarray] = None
+        self.cluster_profiles: Dict[str, Dict[str, float]] = dict(DEFAULT_CLUSTER_PROFILES)
+        self.silhouette_score_: Optional[float] = None
+        self.k = 3
+        self._scaler_mean: Optional[np.ndarray] = None
+        self._scaler_std: Optional[np.ndarray] = None
+
+    # ------------------------------------------------------------------
+    # Public API
+    # ------------------------------------------------------------------
+    def fit(self, csv_path: str) -> "ProgressionModel":
+        """Fit the model from a PPMI CSV file."""
+        try:
+            import pandas as pd
+            df = pd.read_csv(csv_path, low_memory=False)
+            return self._fit_from_dataframe(df)
+        except Exception as exc:
+            logger.warning("ProgressionModel.fit failed: %s — using defaults", exc)
+            self.fitted = False
+            return self
+
+    def _fit_from_dataframe(self, df: "pd.DataFrame") -> "ProgressionModel":
+        import pandas as pd
+        from sklearn.cluster import KMeans
+        from sklearn.preprocessing import StandardScaler
+
+        required = {"PATNO", "YEAR", "updrs3_score", "moca"}
+        if not required.issubset(set(df.columns)):
+            logger.warning("Missing columns for progression model; using defaults.")
+            return self
+
+        # Keep only rows with valid data
+        sub = df[["PATNO", "YEAR", "updrs3_score", "moca", "hy"]].dropna(
+            subset=["PATNO", "YEAR", "updrs3_score"]
+        ).copy()
+        sub["YEAR"] = pd.to_numeric(sub["YEAR"], errors="coerce")
+        sub["updrs3_score"] = pd.to_numeric(sub["updrs3_score"], errors="coerce")
+        sub["moca"] = pd.to_numeric(sub["moca"], errors="coerce")
+        sub = sub.dropna(subset=["YEAR", "updrs3_score"])
+
+        # Compute per-patient velocity vectors
+        velocities: List[List[float]] = []
+        patnos: List[int] = []
+        for patno, grp in sub.groupby("PATNO"):
+            if len(grp) < 2:
+                continue
+            grp = grp.sort_values("YEAR")
+            years = grp["YEAR"].values
+            span = years[-1] - years[0]
+            if span < 0.5:
+                continue
+            delta_updrs = (grp["updrs3_score"].values[-1] - grp["updrs3_score"].values[0]) / span
+            moca_vals = grp["moca"].dropna()
+            if len(moca_vals) >= 2:
+                delta_moca = (moca_vals.values[-1] - moca_vals.values[0]) / span
+            else:
+                delta_moca = 0.0
+            velocities.append([delta_updrs, delta_moca])
+            patnos.append(int(patno))
+
+        if len(velocities) < 10:
+            logger.warning("Too few longitudinal patients (%d); using defaults.", len(velocities))
+            return self
+
+        X = np.array(velocities)
+
+        # Standardize
+        scaler = StandardScaler()
+        X_scaled = scaler.fit_transform(X)
+        self._scaler_mean = scaler.mean_
+        self._scaler_std = scaler.scale_
+
+        # Try k=3, fall back to k=2 if silhouette < 0.3
+        best_k = 3
+        best_score = -1.0
+        best_km = None
+
+        for k in [3, 2]:
+            km = KMeans(n_clusters=k, random_state=42, n_init=10)
+            labels = km.fit_predict(X_scaled)
+            try:
+                from sklearn.metrics import silhouette_score as _sil
+                score = _sil(X_scaled, labels)
+            except Exception:
+                score = 0.0
+            logger.info("ProgressionModel k=%d silhouette=%.3f", k, score)
+            if score > best_score:
+                best_score = score
+                best_k = k
+                best_km = km
+
+        if best_score < 0.3 and best_k == 3:
+            # Retry with k=2
+            km2 = KMeans(n_clusters=2, random_state=42, n_init=10)
+            labels2 = km2.fit_predict(X_scaled)
+            try:
+                from sklearn.metrics import silhouette_score as _sil
+                s2 = _sil(X_scaled, labels2)
+            except Exception:
+                s2 = 0.0
+            if s2 > best_score:
+                best_score = s2
+                best_k = 2
+                best_km = km2
+
+        self.k = best_k
+        self.silhouette_score_ = float(best_score)
+        self.centroids = best_km.cluster_centers_  # type: ignore[union-attr]
+
+        # Sort clusters by ascending UPDRS velocity (index 0 in velocity vec)
+        centroid_means = scaler.inverse_transform(self.centroids)
+        order = np.argsort(centroid_means[:, 0])
+        self.centroids = self.centroids[order]
+        centroid_means = centroid_means[order]
+
+        # Build profiles from centroids
+        labels_used = CLUSTER_LABELS[: self.k]
+        self.cluster_profiles = {}
+        for idx, label in enumerate(labels_used):
+            updrs_vel = max(centroid_means[idx, 0], 0.5)
+            moca_vel = abs(centroid_means[idx, 1]) if centroid_means.shape[1] > 1 else 0.3
+            hy_gain = 0.10 + idx * 0.15
+            self.cluster_profiles[label] = {
+                "updrs3_gain": round(float(updrs_vel), 2),
+                "moca_loss": round(float(moca_vel), 2),
+                "hy_gain": round(float(hy_gain), 2),
+            }
+
+        self.fitted = True
+        logger.info(
+            "ProgressionModel fitted: k=%d, silhouette=%.3f, profiles=%s",
+            self.k, self.silhouette_score_, self.cluster_profiles,
+        )
+        return self
+
+    def assign_cluster(
+        self,
+        snapshots: List[Dict[str, Any]],
+        patient_data: Optional[Dict[str, Any]] = None,
+    ) -> Tuple[str, str]:
+        """
+        Assign a patient to a progression cluster.
+
+        Returns (cluster_id, cluster_label).
+        - If ≥ 2 snapshots with enough time span: use velocity vector.
+        - Otherwise: use current feature vector as proxy.
+        """
+        if not self.fitted or self.centroids is None:
+            return self._heuristic_assign(snapshots, patient_data)
+
+        # Try velocity-based assignment (≥ 2 visits)
+        if len(snapshots) >= 2:
+            cluster = self._velocity_assign(snapshots)
+            if cluster is not None:
+                return cluster
+
+        # Feature-based fallback
+        return self._feature_assign(snapshots, patient_data)
+
+    def get_cluster_profile(self, cluster_label: str) -> Dict[str, float]:
+        """Return progression rates for a cluster."""
+        return self.cluster_profiles.get(cluster_label, DEFAULT_CLUSTER_PROFILES.get("moderate", {}))
+
+    def cluster_weighted_forecast(
+        self,
+        cluster_label: str,
+        current_updrs3: float,
+        current_moca: float,
+        current_hy: float,
+        horizon_months: int,
+        duration_years: float = 0.0,
+    ) -> Dict[str, Optional[float]]:
+        """Compute cluster-weighted forecast for a given horizon."""
+        profile = self.get_cluster_profile(cluster_label)
+        years = horizon_months / 12.0
+
+        # Duration-dependent acceleration factor
+        accel = 1.0 + duration_years * 0.02
+
+        pred_updrs3 = current_updrs3 + profile["updrs3_gain"] * years * accel
+        pred_moca = max(0, min(30, current_moca - profile["moca_loss"] * years))
+        pred_hy = max(0, min(5, current_hy + profile["hy_gain"] * years))
+        pred_total = pred_updrs3 * 1.7 + 8
+
+        return {
+            "predicted_updrs3": round(pred_updrs3, 2),
+            "predicted_total_updrs": round(pred_total, 2),
+            "predicted_moca": round(pred_moca, 2),
+            "predicted_hy": round(pred_hy, 2),
+        }
+
+    # ------------------------------------------------------------------
+    # Internal helpers
+    # ------------------------------------------------------------------
+    def _velocity_assign(self, snapshots: List[Dict[str, Any]]) -> Optional[Tuple[str, str]]:
+        """Assign using velocity vector from snapshot history."""
+        first = snapshots[0]
+        last = snapshots[-1]
+
+        y0 = self._get_year(first)
+        y1 = self._get_year(last)
+        if y0 is None or y1 is None or (y1 - y0) < 0.5:
+            return None
+
+        span = y1 - y0
+        u0 = self._get_updrs3(first)
+        u1 = self._get_updrs3(last)
+        if u0 is None or u1 is None:
+            return None
+
+        delta_updrs = (u1 - u0) / span
+
+        m0 = self._get_moca(first)
+        m1 = self._get_moca(last)
+        delta_moca = ((m1 - m0) / span) if (m0 is not None and m1 is not None) else 0.0
+
+        vec = np.array([[delta_updrs, delta_moca]])
+
+        if self._scaler_mean is not None and self._scaler_std is not None:
+            vec = (vec - self._scaler_mean) / np.clip(self._scaler_std, 1e-8, None)
+
+        distances = np.linalg.norm(self.centroids - vec, axis=1)
+        idx = int(np.argmin(distances))
+        labels = CLUSTER_LABELS[: self.k]
+        return (str(idx), labels[idx])
+
+    def _feature_assign(
+        self,
+        snapshots: List[Dict[str, Any]],
+        patient_data: Optional[Dict[str, Any]],
+    ) -> Tuple[str, str]:
+        """Assign using current feature vector as proxy."""
+        latest = snapshots[-1] if snapshots else (patient_data or {})
+        updrs3 = self._get_updrs3(latest)
+        moca = self._get_moca(latest)
+
+        # Simple heuristic proxy velocity from current scores
+        updrs_proxy = (updrs3 or 10.0) / max(self._get_duration(latest) or 3.0, 1.0)
+        moca_proxy = -((moca or 26.0) - 28.0) / max(self._get_duration(latest) or 3.0, 1.0)
+
+        vec = np.array([[updrs_proxy, moca_proxy]])
+        if self._scaler_mean is not None and self._scaler_std is not None:
+            vec = (vec - self._scaler_mean) / np.clip(self._scaler_std, 1e-8, None)
+
+        if self.centroids is not None:
+            distances = np.linalg.norm(self.centroids - vec, axis=1)
+            idx = int(np.argmin(distances))
+        else:
+            idx = 1  # moderate default
+
+        labels = CLUSTER_LABELS[: self.k]
+        return (str(idx), labels[min(idx, len(labels) - 1)])
+
+    def _heuristic_assign(
+        self,
+        snapshots: List[Dict[str, Any]],
+        patient_data: Optional[Dict[str, Any]],
+    ) -> Tuple[str, str]:
+        """Fallback heuristic when model is not fitted."""
+        latest = snapshots[-1] if snapshots else (patient_data or {})
+        updrs3 = self._get_updrs3(latest) or 0.0
+        if updrs3 >= 25:
+            return ("2", "fast")
+        elif updrs3 >= 12:
+            return ("1", "moderate")
+        return ("0", "slow")
+
+    # ------------------------------------------------------------------
+    @staticmethod
+    def _get_updrs3(d: Dict[str, Any]) -> Optional[float]:
+        val = d.get("motor", d).get("updrs3_score") if isinstance(d.get("motor"), dict) else d.get("updrs3_score")
+        try:
+            return float(val) if val is not None else None
+        except (TypeError, ValueError):
+            return None
+
+    @staticmethod
+    def _get_moca(d: Dict[str, Any]) -> Optional[float]:
+        val = d.get("cognition", d).get("moca") if isinstance(d.get("cognition"), dict) else d.get("moca")
+        try:
+            return float(val) if val is not None else None
+        except (TypeError, ValueError):
+            return None
+
+    @staticmethod
+    def _get_year(d: Dict[str, Any]) -> Optional[float]:
+        val = d.get("year_index")
+        if val is None:
+            val = d.get("YEAR")
+        try:
+            return float(val) if val is not None else None
+        except (TypeError, ValueError):
+            return None
+
+    @staticmethod
+    def _get_duration(d: Dict[str, Any]) -> Optional[float]:
+        val = d.get("duration_years") or d.get("duration_yrs")
+        try:
+            return float(val) if val is not None else None
+        except (TypeError, ValueError):
+            return None

src/rag_system.py

@@ -0,0 +1,762 @@
+"""
+RAG (Retrieval-Augmented Generation) System for Parkinson's Disease Report Generation.
+This module generates comprehensive medical reports based on ML model predictions.
+"""
+
+import os
+import sys
+import warnings
+from datetime import datetime
+from pathlib import Path
+from typing import Dict, Optional
+
+import joblib
+import numpy as np
+import pandas as pd
+
+warnings.filterwarnings("ignore")
+
+# Add src directory to path
+sys.path.append(os.path.join(os.path.dirname(__file__)))
+sys.path.append(os.path.join(os.path.dirname(__file__), "models"))
+
+from data_preprocessing import DataPreprocessor
+from document_manager import DocumentManager
+from models.multimodal_ml import MultimodalEnsemble
+
+
+class MedicalKnowledgeBase:
+    """
+    Knowledge base containing medical information about Parkinson's disease.
+    """
+
+    def __init__(self):
+        self.disease_info = {
+            "HC": {
+                "name": "Healthy Control",
+                "description": "No signs of Parkinson's disease or related movement disorders",
+                "characteristics": [
+                    "Normal motor function",
+                    "No tremor, rigidity, or bradykinesia",
+                    "Normal cognitive function",
+                    "No family history of Parkinson's disease",
+                ],
+                "recommendations": [
+                    "Continue regular health monitoring",
+                    "Maintain active lifestyle",
+                    "Regular exercise and healthy diet",
+                    "Monitor for any changes in motor function",
+                ],
+            },
+            "PD": {
+                "name": "Parkinson's Disease",
+                "description": "Diagnosed with Parkinson's disease showing characteristic motor symptoms",
+                "characteristics": [
+                    "Presence of bradykinesia (slowness of movement)",
+                    "Resting tremor",
+                    "Muscle rigidity",
+                    "Postural instability",
+                    "Possible non-motor symptoms",
+                ],
+                "recommendations": [
+                    "Regular neurological follow-up",
+                    "Consider dopaminergic medication",
+                    "Physical therapy and exercise",
+                    "Speech therapy if needed",
+                    "Monitor for medication side effects",
+                ],
+            },
+            "SWEDD": {
+                "name": "Scans Without Evidence of Dopaminergic Deficit",
+                "description": "Patients with parkinsonian symptoms but normal dopamine transporter imaging",
+                "characteristics": [
+                    "Parkinsonian symptoms present",
+                    "Normal dopamine transporter scans",
+                    "May have tremor without other cardinal signs",
+                    "Often responsive to dopaminergic therapy initially",
+                    "Better long-term prognosis than typical PD",
+                ],
+                "recommendations": [
+                    "Careful clinical monitoring and re-evaluation",
+                    "Consider alternative diagnoses (essential tremor, drug-induced parkinsonism)",
+                    "Regular follow-up to monitor symptom progression",
+                    "Reassess need for dopaminergic medication",
+                    "Consider genetic testing if family history present",
+                ],
+            },
+            "PRODROMAL": {
+                "name": "Prodromal Parkinson's Disease",
+                "description": "Early stage with subtle symptoms that may precede clinical PD",
+                "characteristics": [
+                    "Subtle motor signs",
+                    "REM sleep behavior disorder",
+                    "Hyposmia (reduced sense of smell)",
+                    "Mild cognitive changes",
+                    "Possible depression or anxiety",
+                ],
+                "recommendations": [
+                    "Close monitoring for symptom progression",
+                    "Lifestyle modifications (exercise, diet)",
+                    "Sleep study if REM sleep disorder suspected",
+                    "Cognitive assessment",
+                    "Consider neuroprotective strategies",
+                ],
+            },
+        }
+
+        self.feature_interpretations = {
+            "age": "Patient age is a significant risk factor for Parkinson's disease",
+            "SEX": "Gender differences exist in PD prevalence and presentation",
+            "EDUCYRS": "Education level may influence cognitive reserve",
+            "BMI": "Body mass index can affect disease progression",
+            "fampd": "Family history of Parkinson's disease increases risk",
+            "sym_tremor": "Tremor severity assessment",
+            "sym_rigid": "Muscle rigidity evaluation",
+            "sym_brady": "Bradykinesia (slowness of movement) assessment",
+            "sym_posins": "Postural instability evaluation",
+            "rem": "REM sleep behavior disorder assessment",
+            "ess": "Epworth Sleepiness Scale score",
+            "gds": "Geriatric Depression Scale score",
+            "stai": "State-Trait Anxiety Inventory score",
+            "moca": "Montreal Cognitive Assessment score",
+            "clockdraw": "Clock drawing test performance",
+            "bjlot": "Benton Judgment of Line Orientation test",
+        }
+
+        self.risk_factors = {
+            "high_risk": [
+                "Advanced age (>60 years)",
+                "Family history of Parkinson's disease",
+                "Male gender",
+                "Exposure to pesticides or toxins",
+                "Head trauma history",
+            ],
+            "protective_factors": [
+                "Regular physical exercise",
+                "Caffeine consumption",
+                "Smoking (paradoxically protective)",
+                "Higher education level",
+                "Mediterranean diet",
+            ],
+        }
+
+
+class ReportGenerator:
+    """
+    Generates comprehensive medical reports based on ML predictions and patient data.
+    """
+
+    def __init__(
+        self,
+        knowledge_base: Optional[MedicalKnowledgeBase] = None,
+        docs_dir: str = "../docs",
+    ):
+        self.kb = (
+            knowledge_base if knowledge_base is not None else MedicalKnowledgeBase()
+        )
+        self.ensemble = None
+        self.preprocessor = None
+        self.inference_preprocessor = None
+
+        # Initialize document manager for medical literature
+        self.doc_manager = DocumentManager(docs_dir=docs_dir)
+        print(
+            f"Document manager initialized with {self.doc_manager.get_document_count()['total']} documents"
+        )
+
+    def load_models(self):
+        """Load trained models for prediction."""
+        # Get the correct model directory path
+        current_dir = os.path.dirname(os.path.abspath(__file__))
+        model_dir = os.path.join(os.path.dirname(current_dir), "models", "saved")
+
+        self.ensemble = MultimodalEnsemble()
+        self.ensemble.load_traditional_models(model_dir)
+
+        prep_path = os.path.join(model_dir, "traditional_preprocessor.joblib")
+        if os.path.exists(prep_path):
+            self.inference_preprocessor = joblib.load(prep_path)
+
+        inferred_input_dim = 31
+        try:
+            if self.inference_preprocessor is not None:
+                inferred_input_dim = len(
+                    self.inference_preprocessor.get_feature_names_out()
+                )
+        except Exception:
+            pass
+
+        # Transformer checkpoints are optional; traditional models + ensemble are sufficient for inference.
+        if os.getenv("PD_LOAD_TRANSFORMERS", "1") == "0":
+            print("Skipping optional transformer model loading (PD_LOAD_TRANSFORMERS=0)")
+        else:
+            self.ensemble.load_transformer_models(
+                model_dir, input_dim=inferred_input_dim, num_classes=4
+            )
+
+        ensemble_path = os.path.join(model_dir, "multimodal_ensemble.joblib")
+        self.ensemble.load_ensemble(ensemble_path)
+
+        if self.ensemble.ensemble_model is None:
+            raise FileNotFoundError(
+                f"Required ensemble artifact not found at: {ensemble_path}"
+            )
+
+        self.preprocessor = DataPreprocessor()
+        print("Models loaded successfully")
+
+    def predict_patient(self, patient_data: Dict) -> Dict:
+        """Make predictions for a single patient."""
+        if self.ensemble is None:
+            self.load_models()
+
+        try:
+            # Store original patient data for report generation
+            self.original_patient_data = patient_data.copy()
+
+            # Create a DataFrame with the patient data
+            df_patient = pd.DataFrame([patient_data])
+
+            # Use fitted preprocessor from training when available
+            if self.inference_preprocessor is not None:
+                required_cols = list(
+                    getattr(self.inference_preprocessor, "feature_names_in_", [])
+                )
+                if required_cols:
+                    for col in required_cols:
+                        if col not in df_patient.columns:
+                            df_patient[col] = np.nan
+                    df_patient = df_patient[required_cols]
+                X_infer = self.inference_preprocessor.transform(df_patient)
+            else:
+                # Fallback: numeric-only matrix to avoid hard failure
+                X_infer = df_patient.select_dtypes(include=[np.number]).fillna(0).values
+
+            # Make predictions using ensemble
+            if self.ensemble is None or self.ensemble.ensemble_model is None:
+                raise RuntimeError("Ensemble model is not loaded")
+            predictions, probabilities = self.ensemble.predict_ensemble(X_infer)
+
+            confidence = float(np.max(probabilities[0]))
+
+            trad_outputs, _ = self.ensemble.get_traditional_predictions(X_infer)
+            trans_outputs, _ = self.ensemble.get_transformer_predictions(X_infer)
+
+            trad_preds = {
+                model_name: int(pred[0])
+                for model_name, pred in trad_outputs.items()
+                if len(pred) > 0
+            }
+            trans_preds = {
+                model_name: int(pred[0])
+                for model_name, pred in trans_outputs.items()
+                if len(pred) > 0
+            }
+
+            return {
+                "ensemble_prediction": int(predictions[0]),
+                "ensemble_probabilities": probabilities[0],
+                "traditional_predictions": trad_preds,
+                "transformer_predictions": trans_preds,
+                "confidence": confidence,
+                "patient_data": self.original_patient_data,  # Store original data for report
+            }
+
+        except Exception as e:
+            print(f"Error in prediction: {e}")
+            import traceback
+
+            traceback.print_exc()
+            # Analyze symptoms to make a more informed prediction instead of defaulting to PD
+            symptoms = {
+                "tremor": patient_data.get("sym_tremor", 0),
+                "rigidity": patient_data.get("sym_rigid", 0),
+                "bradykinesia": patient_data.get("sym_brady", 0),
+                "postural_instability": patient_data.get("sym_posins", 0),
+                "family_history": patient_data.get("fampd", 0),
+                "cognitive_score": patient_data.get("moca", 25),
+            }
+
+            # Simple rule-based classification
+            pd_score = 0
+            # Check for cardinal PD symptoms
+            if symptoms["tremor"] > 2:
+                pd_score += 2
+            if symptoms["rigidity"] > 2:
+                pd_score += 2
+            if symptoms["bradykinesia"] > 2:
+                pd_score += 2
+            if symptoms["postural_instability"] > 2:
+                pd_score += 2
+            if symptoms["family_history"] > 0:
+                pd_score += 1
+            if symptoms["cognitive_score"] < 24:
+                pd_score += 1
+
+            # Determine class based on score
+            if pd_score >= 6:
+                pred_class = 1  # PD
+                probs = [0.1, 0.7, 0.1, 0.1]
+            elif pd_score >= 4:
+                pred_class = 3  # Prodromal
+                probs = [0.2, 0.3, 0.1, 0.4]
+            elif pd_score >= 2:
+                pred_class = 2  # SWEDD
+                probs = [0.3, 0.1, 0.5, 0.1]
+            else:
+                pred_class = 0  # Healthy Control
+                probs = [0.7, 0.1, 0.1, 0.1]
+
+            return {
+                "ensemble_prediction": pred_class,
+                "ensemble_probabilities": probs,
+                "traditional_predictions": {
+                    "xgboost": pred_class,
+                    "lightgbm": pred_class,
+                    "svm": pred_class,
+                },
+                "transformer_predictions": {
+                    "pubmedbert": pred_class,
+                    "biomistral": pred_class,
+                    "clinical_t5": pred_class,
+                },
+                "confidence": max(probs),
+                "patient_data": self.original_patient_data,
+            }
+
+    def generate_clinical_summary(
+        self, prediction_results: Dict, patient_data: Dict
+    ) -> str:
+        """Generate clinical summary based on predictions and medical literature."""
+        pred_class = prediction_results["ensemble_prediction"]
+        confidence = prediction_results["confidence"]
+        probabilities = prediction_results["ensemble_probabilities"]
+
+        # Map prediction to class name - 4 classes
+        class_names = ["HC", "PD", "SWEDD", "PRODROMAL"]
+        predicted_condition = class_names[pred_class]
+
+        # Get disease information
+        disease_info = self.kb.disease_info[predicted_condition]
+
+        # Retrieve relevant medical literature
+        literature_insights = self._get_literature_insights(
+            predicted_condition, patient_data
+        )
+
+        summary = f"""
+CLINICAL ASSESSMENT SUMMARY
+===========================
+
+PRIMARY DIAGNOSIS: {disease_info["name"]}
+Confidence Level: {confidence * 100:.2f}%
+
+DIAGNOSTIC PROBABILITY DISTRIBUTION:
+- Healthy Control: {probabilities[0] * 100:.2f}%
+- Parkinson's Disease: {probabilities[1] * 100:.2f}%
+- SWEDD: {probabilities[2] * 100:.2f}%
+- Prodromal PD: {probabilities[3] * 100:.2f}%
+
+CLINICAL DESCRIPTION:
+{disease_info["description"]}
+
+KEY CHARACTERISTICS OBSERVED:
+"""
+        for char in disease_info["characteristics"]:
+            summary += f"• {char}\n"
+
+        # Add insights from medical literature with better formatting
+        if literature_insights:
+            summary += f"\nINSIGHTS FROM MEDICAL LITERATURE:\n"
+            summary += "=" * 50 + "\n"
+            summary += literature_insights
+
+        return summary
+
+    def _get_literature_insights(self, condition: str, patient_data: Dict) -> str:
+        """Retrieve insights from medical literature relevant to the patient's condition."""
+        # Check if document manager has documents
+        if self.doc_manager.get_document_count()["total"] == 0:
+            return "No medical literature available. Add medical papers to enhance insights."
+
+        # Construct search query based on condition and key symptoms
+        query_parts = [condition]
+
+        # Always include key symptoms in query with their severity
+        symptoms = {
+            "tremor": patient_data.get("sym_tremor", 0),
+            "rigidity": patient_data.get("sym_rigid", 0),
+            "bradykinesia": patient_data.get("sym_brady", 0),
+            "postural instability": patient_data.get("sym_posins", 0),
+        }
+
+        # Add all symptoms with their severity to create more specific queries
+        for symptom, severity in symptoms.items():
+            if severity > 0:
+                query_parts.append(f"{symptom} severity:{severity}")
+
+        # Add cognitive and psychiatric factors
+        if "moca" in patient_data:
+            moca = patient_data.get("moca", 30)
+            if moca < 26:
+                query_parts.append("cognitive impairment")
+                if moca < 20:
+                    query_parts.append("severe cognitive impairment")
+
+        if "gds" in patient_data and patient_data.get("gds", 0) > 5:
+            query_parts.append("depression")
+
+        if "stai" in patient_data and patient_data.get("stai", 0) > 40:
+            query_parts.append("anxiety")
+
+        # Add demographic factors if available
+        if "age" in patient_data:
+            age = patient_data["age"]
+            query_parts.append(f"age {age}")
+            if age < 50:
+                query_parts.append("early onset")
+            elif age > 70:
+                query_parts.append("elderly")
+
+        if "SEX" in patient_data:
+            gender = "male" if patient_data["SEX"] == 1 else "female"
+            query_parts.append(gender)
+
+        # Add family history if present
+        if patient_data.get("fampd", 0) > 0:
+            query_parts.append("family history")
+
+        # Construct final query
+        query = " ".join(query_parts)
+
+        # Retrieve relevant passages with increased number of results
+        passages = self.doc_manager.extract_relevant_passages(query, top_k=5)
+
+        if not passages:
+            return "No specific literature found for this patient's condition and symptoms."
+
+        # Format insights with proper citations and more context
+        insights = ""
+        for i, passage in enumerate(passages):
+            # Extract document metadata for citation
+            doc_title = passage["doc_title"]
+            # Format the citation properly
+            citation = f"[{doc_title}]"
+
+            # Add the passage with citation
+            insights += f"{i + 1}. From '{doc_title}': {passage['text'][:400]}...\n   {citation}\n\n"
+
+        return insights
+
+    def generate_feature_analysis(self, patient_data: Dict) -> str:
+        """Generate analysis of key patient features."""
+        # Use stored original patient data if available
+        if hasattr(self, "original_patient_data"):
+            patient_data = self.original_patient_data
+
+        analysis = "\nFEATURE ANALYSIS:\n" + "=" * 50 + "\n"
+
+        # Expanded key features list with better labels
+        key_features = [
+            ("age", "Age"),
+            ("SEX", "Gender"),
+            ("EDUCYRS", "Education Years"),
+            ("BMI", "Body Mass Index"),
+            ("fampd", "Family History"),
+            ("sym_tremor", "Tremor Severity"),
+            ("sym_rigid", "Rigidity"),
+            ("sym_brady", "Bradykinesia"),
+            ("sym_posins", "Postural Instability"),
+            ("moca", "MoCA Score"),
+            ("gds", "Depression Score"),
+            ("stai", "Anxiety Score"),
+            ("ess", "Sleepiness Scale"),
+            ("rem", "REM Sleep Behavior"),
+        ]
+
+        for feature_key, feature_name in key_features:
+            if feature_key in patient_data:
+                value = patient_data[feature_key]
+                interpretation = self.kb.feature_interpretations.get(feature_key, "")
+
+                # Format value based on feature type
+                if feature_key == "age":
+                    risk_level = (
+                        "High" if value > 60 else "Moderate" if value > 50 else "Low"
+                    )
+                    analysis += (
+                        f"• {feature_name} ({value} years): Risk level: {risk_level}\n"
+                    )
+                elif feature_key == "SEX":
+                    formatted_value = "Male" if value == 1 else "Female"
+                    analysis += f"• {feature_name}: {formatted_value}\n"
+                elif feature_key == "moca":
+                    cognitive_status = (
+                        "Normal"
+                        if value >= 26
+                        else "Mild impairment"
+                        if value >= 22
+                        else "Significant impairment"
+                    )
+                    analysis += (
+                        f"• {feature_name} ({value}/30): Status: {cognitive_status}\n"
+                    )
+                elif feature_key == "fampd":
+                    family_history = "Positive" if value > 0 else "Negative"
+                    analysis += (
+                        f"• {feature_name}: {family_history} for Parkinson's disease\n"
+                    )
+                elif feature_key in [
+                    "sym_tremor",
+                    "sym_rigid",
+                    "sym_brady",
+                    "sym_posins",
+                ]:
+                    severity = ["None", "Mild", "Moderate", "Severe", "Very Severe"]
+                    formatted_value = severity[min(int(value), 4)]
+                    analysis += f"• {feature_name}: {formatted_value}\n"
+                elif feature_key == "gds":
+                    status = "Normal" if value <= 5 else "Depression indicated"
+                    analysis += f"• {feature_name}: {value} - {status}\n"
+                elif feature_key == "stai":
+                    status = "Normal" if value <= 40 else "Anxiety indicated"
+                    analysis += f"• {feature_name}: {value} - {status}\n"
+                else:
+                    analysis += f"• {feature_name}: {value}\n"
+
+                # Add interpretation if available and not already added
+                if interpretation and feature_key not in ["age", "moca", "fampd"]:
+                    analysis = analysis.rstrip("\n") + f" - {interpretation}\n"
+
+        return analysis
+
+    def generate_recommendations(
+        self, prediction_results: Dict, patient_data: Dict
+    ) -> str:
+        """Generate clinical recommendations."""
+        # Use stored original patient data if available
+        if hasattr(self, "original_patient_data"):
+            patient_data = self.original_patient_data
+
+        pred_class = prediction_results["ensemble_prediction"]
+        class_names = ["HC", "PD", "SWEDD", "PRODROMAL"]
+        predicted_condition = class_names[pred_class]
+
+        disease_info = self.kb.disease_info[predicted_condition]
+
+        recommendations = "\nCLINICAL RECOMMENDATIONS:\n" + "=" * 50 + "\n"
+
+        for i, rec in enumerate(disease_info["recommendations"], 1):
+            recommendations += f"{i}. {rec}\n"
+
+        # Add general recommendations based on risk factors
+        recommendations += "\nADDITIONAL CONSIDERATIONS:\n"
+
+        if patient_data.get("age", 0) > 60:
+            recommendations += (
+                "• Age-related monitoring: Increased surveillance due to advanced age\n"
+            )
+
+        if patient_data.get("fampd", 0) > 0:
+            recommendations += (
+                "• Genetic counseling: Consider due to positive family history\n"
+            )
+
+        if patient_data.get("moca", 30) < 26:
+            recommendations += "• Cognitive assessment: Follow-up neuropsychological testing recommended\n"
+
+        if patient_data.get("gds", 0) > 5:
+            recommendations += "• Depression management: Consider psychiatric evaluation and treatment\n"
+
+        if patient_data.get("stai", 0) > 40:
+            recommendations += (
+                "• Anxiety management: Consider psychiatric evaluation and treatment\n"
+            )
+
+        if patient_data.get("ess", 0) > 10:
+            recommendations += "• Sleep evaluation: Consider sleep study for excessive daytime sleepiness\n"
+
+        if patient_data.get("rem", 0) > 0:
+            recommendations += "• REM sleep behavior disorder: Consider polysomnography and treatment\n"
+
+        return recommendations
+
+    def generate_model_consensus(self, prediction_results: Dict) -> str:
+        """Generate analysis of model consensus."""
+        trad_preds = prediction_results["traditional_predictions"]
+        trans_preds = prediction_results["transformer_predictions"]
+        ensemble_pred = prediction_results["ensemble_prediction"]
+
+        consensus = "\nMODEL CONSENSUS ANALYSIS:\n" + "=" * 50 + "\n"
+
+        # Check agreement between models
+        all_predictions = (
+            list(trad_preds.values()) + list(trans_preds.values()) + [ensemble_pred]
+        )
+        unique_predictions = set(all_predictions)
+
+        if len(all_predictions) == 1:
+            consensus += "• LIMITED CONSENSUS: Only the ensemble model was available at inference time\n"
+        elif len(unique_predictions) == 1:
+            consensus += "• STRONG CONSENSUS: All models agree on the diagnosis\n"
+        elif len(unique_predictions) == 2:
+            consensus += "• MODERATE CONSENSUS: Most models agree with some variation\n"
+        else:
+            consensus += "• WEAK CONSENSUS: Significant disagreement between models\n"
+
+        consensus += f"\nIndividual Model Predictions:\n"
+        if trad_preds:
+            consensus += "Traditional Machine Learning Models:\n"
+            for model, pred in trad_preds.items():
+                class_names = ["HC", "PD", "SWEDD", "PRODROMAL"]
+                consensus += f"  • {model.upper()}: {class_names[pred]}\n"
+        else:
+            consensus += "Traditional Machine Learning Models:\n  • No traditional base-model outputs were available\n"
+
+        consensus += "\nMedical Transformer Models:\n"
+        if trans_preds:
+            for model, pred in trans_preds.items():
+                class_names = ["HC", "PD", "SWEDD", "PRODROMAL"]
+                model_display = model.replace("_", " ").title()
+                if model == "pubmedbert":
+                    model_display = "PubMedBERT (Encoder)"
+                elif model == "biomistral":
+                    model_display = "BioMistral (Decoder)"
+                elif model == "clinical_t5":
+                    model_display = "Clinical-T5 (Encoder-Decoder)"
+                elif model == "feedforward":
+                    model_display = "Feedforward Tabular Network"
+                consensus += f"  • {model_display}: {class_names[pred]}\n"
+        else:
+            consensus += "  • No transformer base-model outputs were available\n"
+
+        return consensus
+
+    def generate_full_report(self, patient_data: Dict, patient_id: str = None) -> str:
+        """Generate a comprehensive medical report."""
+        if patient_id is None:
+            patient_id = f"PATIENT_{datetime.now().strftime('%Y%m%d_%H%M%S')}"
+        else:
+            patient_id = str(patient_id)
+
+        # Make predictions
+        prediction_results = self.predict_patient(patient_data)
+
+        # Generate report sections
+        header = f"""
+PARKINSON'S DISEASE ASSESSMENT REPORT
+=====================================
+Patient ID: {patient_id}
+Report Date: {datetime.now().strftime("%Y-%m-%d %H:%M:%S")}
+Generated by: AI-Powered Multimodal ML System
+"""
+
+        clinical_summary = self.generate_clinical_summary(
+            prediction_results, patient_data
+        )
+        feature_analysis = self.generate_feature_analysis(patient_data)
+        recommendations = self.generate_recommendations(
+            prediction_results, patient_data
+        )
+        model_consensus = self.generate_model_consensus(prediction_results)
+
+        footer = f"""
+DISCLAIMER:
+===========
+This report is generated by an AI system for research and educational purposes.
+It should not replace professional medical diagnosis or treatment decisions.
+Always consult with qualified healthcare professionals for medical advice.
+
+Report generated using multimodal machine learning with {prediction_results["confidence"] * 100:.2f}% confidence.
+"""
+
+        full_report = (
+            header
+            + clinical_summary
+            + feature_analysis
+            + recommendations
+            + model_consensus
+            + footer
+        )
+
+        return full_report
+
+    def save_report(self, report: str, filename: Optional[str] = None) -> str:
+        """Save report to file using an absolute path under the project reports directory."""
+        if filename is None:
+            filename = f"medical_report_{datetime.now().strftime('%Y%m%d_%H%M%S')}.txt"
+        else:
+            filename = str(filename)
+
+        safe_name = Path(filename).name.strip()
+        safe_name = "".join(
+            ch if ch.isalnum() or ch in "._- " else "_" for ch in safe_name
+        ).strip(" .")
+        if not safe_name:
+            safe_name = f"medical_report_{datetime.now().strftime('%Y%m%d_%H%M%S')}.txt"
+        if not safe_name.lower().endswith(".txt"):
+            safe_name = f"{safe_name}.txt"
+
+        project_root = Path(__file__).resolve().parents[1]
+        reports_dir = project_root / "reports"
+        reports_dir.mkdir(parents=True, exist_ok=True)
+
+        filepath = reports_dir / safe_name
+
+        with open(filepath, "w", encoding="utf-8") as f:
+            f.write(report)
+
+        resolved_path = str(filepath.resolve())
+        print(f"Report saved to: {resolved_path}")
+        return resolved_path
+
+
+def demo_report_generation():
+    """Demonstrate report generation with sample patient data."""
+    print("RAG System Demo - Generating Sample Medical Report")
+    print("=" * 60)
+
+    # Sample patient data
+    sample_patient = {
+        "age": 65,
+        "SEX": 1,  # Male
+        "EDUCYRS": 16,
+        "race": 1,
+        "BMI": 26.5,
+        "fampd": 1,  # Positive family history
+        "fampd_bin": 1,
+        "sym_tremor": 2,
+        "sym_rigid": 1,
+        "sym_brady": 2,
+        "sym_posins": 1,
+        "rem": 1,
+        "ess": 8,
+        "gds": 3,
+        "stai": 35,
+        "moca": 24,
+        "clockdraw": 3,
+        "bjlot": 25,
+    }
+
+    # Initialize report generator
+    report_gen = ReportGenerator()
+
+    try:
+        # Generate report
+        report = report_gen.generate_full_report(sample_patient, "DEMO_PATIENT_001")
+
+        # Print report
+        print(report)
+
+        # Save report
+        filepath = report_gen.save_report(report, "demo_medical_report.txt")
+
+        return report, filepath
+
+    except Exception as e:
+        print(f"Error generating report: {e}")
+        return None, None
+
+
+if __name__ == "__main__":
+    demo_report_generation()

src/risk_stratifier.py

@@ -0,0 +1,122 @@
+"""
+Risk Stratifier — MDS Prodromal Markers + Bootstrap CI.
+
+Maps available PPMI features to MDS prodromal markers, computes a
+risk score, and produces bootstrap confidence intervals.
+"""
+
+from __future__ import annotations
+
+import logging
+from typing import Any, Dict, List, Optional
+
+import numpy as np
+
+logger = logging.getLogger(__name__)
+
+MDS_MARKERS = {
+    "rem":       {"weight": 2.5, "threshold_fn": "binary"},
+    "upsit":     {"weight": 2.0, "threshold_fn": "upsit_low"},
+    "pigd":      {"weight": 1.5, "threshold_fn": "pigd_present"},
+    "gds":       {"weight": 1.2, "threshold_fn": "gds_high"},
+    "fampd_bin": {"weight": 1.0, "threshold_fn": "binary"},
+}
+
+
+def _check_marker(name: str, value: Optional[float]) -> Optional[float]:
+    if value is None:
+        return None
+    spec = MDS_MARKERS.get(name)
+    if spec is None:
+        return None
+    fn, w = spec["threshold_fn"], spec["weight"]
+    if fn == "binary":
+        return w if value >= 1.0 else 0.0
+    if fn == "upsit_low":
+        return w if value <= 22.0 else 0.0
+    if fn == "pigd_present":
+        return w if value >= 1.0 else 0.0
+    if fn == "gds_high":
+        return w if value >= 5.0 else 0.0
+    return 0.0
+
+
+class RiskStratifier:
+    """MDS criteria risk stratification with bootstrap CIs."""
+
+    def __init__(self, n_bootstrap: int = 100) -> None:
+        self.n_bootstrap = n_bootstrap
+
+    def stratify(self, patient_data: Dict[str, Any]) -> Dict[str, Any]:
+        marker_values: Dict[str, Optional[float]] = {}
+        for marker in MDS_MARKERS:
+            marker_values[marker] = self._extract(patient_data, marker)
+
+        contributions: Dict[str, Optional[float]] = {}
+        available: List[str] = []
+        for m, v in marker_values.items():
+            s = _check_marker(m, v)
+            contributions[m] = s
+            if s is not None:
+                available.append(m)
+
+        total = sum(v for v in contributions.values() if v is not None)
+        max_p = sum(MDS_MARKERS[m]["weight"] for m in available) if available else 1.0
+        raw_conf = total / max(max_p, 1e-6)
+        category = self._categorize(raw_conf, patient_data)
+        ci_lo, ci_hi = self._bootstrap_ci(marker_values, available)
+
+        return {
+            "category": category,
+            "confidence": round(float(raw_conf), 4),
+            "ci_lower": round(float(ci_lo), 4),
+            "ci_upper": round(float(ci_hi), 4),
+            "marker_contributions": contributions,
+            "total_score": round(float(total), 4),
+            "max_possible_score": round(float(max_p), 4),
+        }
+
+    def _bootstrap_ci(self, marker_values, available):
+        if not available:
+            return (0.0, 0.0)
+        rng = np.random.RandomState(42)
+        weights = np.array([MDS_MARKERS[m]["weight"] for m in available])
+        base = np.array([_check_marker(m, marker_values[m]) or 0.0 for m in available])
+        scores = []
+        for _ in range(self.n_bootstrap):
+            idx = rng.choice(len(available), size=len(available), replace=True)
+            conf = base[idx].sum() / max(weights[idx].sum(), 1e-6)
+            conf += rng.normal(0, 0.02)
+            scores.append(float(np.clip(conf, 0, 1)))
+        return (float(np.percentile(scores, 2.5)), float(np.percentile(scores, 97.5)))
+
+    def _categorize(self, confidence, patient_data):
+        sym_total = sum(
+            float(patient_data.get(s) or (patient_data.get("motor", {}) or {}).get(s) or 0)
+            for s in ("sym_tremor", "sym_rigid", "sym_brady", "sym_posins")
+        )
+        updrs3 = (patient_data.get("motor", {}) or {}).get("updrs3_score") or patient_data.get("updrs3_score")
+        if confidence >= 0.55 or sym_total >= 5 or (updrs3 is not None and float(updrs3) >= 20):
+            return "PD"
+        if confidence >= 0.30:
+            return "Prodromal PD"
+        if confidence >= 0.15:
+            return "SWEDD"
+        return "HC"
+
+    @staticmethod
+    def _extract(data: Dict[str, Any], marker: str) -> Optional[float]:
+        val = data.get(marker)
+        if val is not None:
+            try:
+                return float(val)
+            except (TypeError, ValueError):
+                pass
+        for sec in ("non_motor", "motor", "cognition", "autonomic"):
+            sub = data.get(sec)
+            if isinstance(sub, dict) and marker in sub and sub[marker] is not None:
+                try:
+                    return float(sub[marker])
+                except (TypeError, ValueError):
+                    pass
+        return None

src/static/css/label-fix.css

@@ -0,0 +1,43 @@
+/* Fix for label visibility - MAXIMUM CONTRAST */
+.form-label {
+    color: #000000 !important;
+    background-color: #ffffff !important;
+    font-weight: 700 !important;
+    display: inline-block !important;
+    margin-bottom: 0.5rem !important;
+    padding: 2px 5px !important;
+    border-radius: 3px !important;
+    box-shadow: 0 0 3px rgba(0,0,0,0.3) !important;
+    font-size: 1.1em !important;
+    border: 1px solid #000 !important;
+}
+
+/* Fix for dark mode label visibility */
+[data-theme="dark"] .form-label {
+    color: #ffffff !important;
+    background-color: #333333 !important;
+    border: 1px solid #fff !important;
+    box-shadow: 0 0 3px rgba(255,255,255,0.5) !important;
+}
+
+/* Make sure select options are visible */
+select.form-control option {
+    background-color: #ffffff !important;
+    color: #000000 !important;
+    font-weight: 500 !important;
+}
+
+/* Improve form control visibility */
+.form-control {
+    border: 2px solid #555 !important;
+    background-color: #ffffff !important;
+    color: #000000 !important;
+    font-weight: 500 !important;
+}
+
+/* Dark mode form controls */
+[data-theme="dark"] .form-control {
+    background-color: #444 !important;
+    color: #fff !important;
+    border: 2px solid #777 !important;
+}

src/static/css/styles.css

@@ -0,0 +1,274 @@
+:root {
+    --primary-color: #4a90e2;
+    --secondary-color: #2c3e50;
+    --accent-color: #e74c3c;
+    --background-color: #f8f9fa;
+    --text-color: #333;
+    --border-color: #ddd;
+    --success-color: #2ecc71;
+    --warning-color: #f39c12;
+    --danger-color: #e74c3c;
+    --card-bg-color: #ffffff;
+    --navbar-bg-color: #2c3e50;
+    --navbar-text-color: #ffffff;
+    --input-bg-color: #ffffff;
+    --shadow-color: rgba(0,0,0,0.1);
+}
+
+[data-theme="dark"] {
+    --primary-color: #61a0ff;
+    --secondary-color: #34495e;
+    --accent-color: #ff6b6b;
+    --background-color: #121212;
+    --text-color: #f5f5f5;
+    --border-color: #444;
+    --success-color: #4cd787;
+    --warning-color: #ffb74d;
+    --danger-color: #ff5252;
+    --card-bg-color: #1e1e1e;
+    --navbar-bg-color: #1a1a2e;
+    --navbar-text-color: #f0f0f0;
+    --input-bg-color: #2a2a2a;
+    --shadow-color: rgba(255,255,255,0.05);
+}
+
+body {
+    font-family: 'Roboto', sans-serif;
+    line-height: 1.6;
+    color: var(--text-color);
+    background-color: var(--background-color);
+    transition: all 0.3s ease;
+    margin: 0;
+    padding: 0;
+    min-height: 100vh;
+}
+
+/* Tooltip styles */
+.tooltip {
+    position: relative;
+    display: inline-block;
+    border-bottom: 1px dotted var(--primary-color);
+    cursor: help;
+}
+
+.tooltip .tooltiptext {
+    visibility: hidden;
+    width: 250px;
+    background-color: var(--secondary-color);
+    color: white;
+    text-align: center;
+    border-radius: 6px;
+    padding: 10px;
+    position: absolute;
+    z-index: 1000;
+    bottom: 125%;
+    left: 50%;
+    margin-left: -125px;
+    opacity: 0;
+    transition: opacity 0.3s;
+    font-size: 0.9rem;
+    box-shadow: 0 5px 15px rgba(0,0,0,0.2);
+}
+
+.tooltip .tooltiptext::after {
+    content: "";
+    position: absolute;
+    top: 100%;
+    left: 50%;
+    margin-left: -5px;
+    border-width: 5px;
+    border-style: solid;
+    border-color: var(--secondary-color) transparent transparent transparent;
+}
+
+.tooltip:hover .tooltiptext {
+    visibility: visible;
+    opacity: 1;
+}
+
+.container {
+    max-width: 1200px;
+    margin: 0 auto;
+    padding: 20px;
+}
+
+.navbar {
+    background-color: var(--navbar-bg-color);
+    color: var(--navbar-text-color);
+    padding: 15px 0;
+    box-shadow: 0 2px 10px rgba(0,0,0,0.1);
+    width: 100%;
+}
+
+.navbar-brand {
+    font-size: 1.5rem;
+    font-weight: bold;
+    color: var(--navbar-text-color);
+    text-decoration: none;
+    padding: 0 20px;
+}
+
+.navbar-nav {
+    display: flex;
+    list-style: none;
+    margin: 0;
+    padding: 0;
+}
+
+.nav-item {
+    margin-right: 20px;
+}
+
+.nav-link {
+    color: white;
+    text-decoration: none;
+}
+
+.card {
+    background-color: var(--card-bg-color);
+    border: 1px solid var(--border-color);
+    border-radius: 8px;
+    padding: 25px;
+    margin-bottom: 25px;
+    box-shadow: 0 4px 12px var(--shadow-color);
+    transition: transform 0.3s ease, box-shadow 0.3s ease;
+}
+
+.card:hover {
+    transform: translateY(-5px);
+    box-shadow: 0 8px 15px var(--shadow-color);
+}
+
+.btn {
+    display: inline-block;
+    padding: 10px 20px;
+    border: none;
+    border-radius: 6px;
+    cursor: pointer;
+    text-decoration: none;
+    font-size: 1rem;
+    font-weight: 500;
+    transition: all 0.3s ease;
+    box-shadow: 0 2px 5px rgba(0,0,0,0.1);
+    text-align: center;
+}
+
+.btn:hover {
+    transform: translateY(-2px);
+    box-shadow: 0 4px 8px rgba(0,0,0,0.15);
+}
+
+.btn-primary {
+    background-color: var(--primary-color);
+    color: white;
+}
+
+.btn-primary:hover {
+    background-color: var(--primary-color);
+    filter: brightness(110%);
+}
+
+.btn-secondary {
+    background-color: var(--secondary-color);
+    color: white;
+}
+
+.btn-secondary:hover {
+    background-color: var(--secondary-color);
+    filter: brightness(110%);
+}
+
+.btn-danger {
+    background-color: var(--danger-color);
+    color: white;
+}
+
+.btn-danger:hover {
+    background-color: var(--danger-color);
+    filter: brightness(110%);
+}
+
+.form-group {
+    margin-bottom: 15px;
+}
+
+.form-control {
+    width: 100%;
+    padding: 12px;
+    border: 1px solid var(--border-color);
+    border-radius: 6px;
+    background-color: var(--input-bg-color);
+    color: var(--text-color);
+    transition: border-color 0.3s ease, box-shadow 0.3s ease;
+}
+
+.form-control:focus {
+    outline: none;
+    border-color: var(--primary-color);
+    box-shadow: 0 0 0 3px rgba(74, 144, 226, 0.25);
+}
+
+.tooltip {
+    position: relative;
+    display: inline-block;
+    cursor: help;
+}
+
+.tooltip .tooltiptext {
+    visibility: hidden;
+    width: 200px;
+    background-color: var(--secondary-color);
+    color: white;
+    text-align: center;
+    border-radius: 6px;
+    padding: 5px;
+    position: absolute;
+    z-index: 1;
+    bottom: 125%;
+    left: 50%;
+    margin-left: -100px;
+    opacity: 0;
+    transition: opacity 0.3s;
+}
+
+.tooltip:hover .tooltiptext {
+    visibility: visible;
+    opacity: 1;
+}
+
+.theme-toggle {
+    background: none;
+    border: none;
+    color: var(--navbar-text-color);
+    cursor: pointer;
+    font-size: 1.2rem;
+    margin-left: auto;
+    padding: 8px 12px;
+    border-radius: 50%;
+    transition: all 0.3s ease;
+}
+
+.theme-toggle:hover {
+    background-color: rgba(255, 255, 255, 0.1);
+    transform: rotate(15deg);
+}
+
+/* Dark mode specific styles */
+[data-theme="dark"] .card {
+    box-shadow: 0 2px 5px rgba(255,255,255,0.05);
+}
+
+[data-theme="dark"] .form-control {
+    background-color: #333;
+    color: var(--text-color);
+    border-color: #555;
+}
+
+[data-theme="dark"] .table {
+    color: var(--text-color);
+}
+
+[data-theme="dark"] .table th,
+[data-theme="dark"] .table td {
+    border-color: #444;
+}

src/static/js/darkmode.js

@@ -0,0 +1,50 @@
+document.addEventListener('DOMContentLoaded', function() {
+    const themeToggle = document.getElementById('theme-toggle');
+    const htmlElement = document.documentElement;
+    
+    // Check for saved theme preference or use device preference
+    const savedTheme = localStorage.getItem('theme') || 
+                      (window.matchMedia('(prefers-color-scheme: dark)').matches ? 'dark' : 'light');
+    
+    // Apply saved theme
+    htmlElement.setAttribute('data-theme', savedTheme);
+    updateToggleIcon(savedTheme);
+    
+    // Toggle theme when button is clicked
+    themeToggle.addEventListener('click', function() {
+        const currentTheme = htmlElement.getAttribute('data-theme');
+        const newTheme = currentTheme === 'dark' ? 'light' : 'dark';
+        
+        htmlElement.setAttribute('data-theme', newTheme);
+        localStorage.setItem('theme', newTheme);
+        updateToggleIcon(newTheme);
+    });
+    
+    // Update toggle icon based on current theme
+    function updateToggleIcon(theme) {
+        if (theme === 'dark') {
+            themeToggle.innerHTML = '<i class="fas fa-sun"></i>';
+            themeToggle.title = 'Switch to light mode';
+        } else {
+            themeToggle.innerHTML = '<i class="fas fa-moon"></i>';
+            themeToggle.title = 'Switch to dark mode';
+        }
+    }
+    
+    // Initialize tooltips
+    initializeTooltips();
+});
+
+function initializeTooltips() {
+    const tooltipElements = document.querySelectorAll('[data-tooltip]');
+    
+    tooltipElements.forEach(element => {
+        const tooltipText = element.getAttribute('data-tooltip');
+        const tooltipSpan = document.createElement('span');
+        tooltipSpan.className = 'tooltiptext';
+        tooltipSpan.textContent = tooltipText;
+        
+        element.classList.add('tooltip');
+        element.appendChild(tooltipSpan);
+    });
+}

src/train_model_suite.py

@@ -0,0 +1,1382 @@
+from __future__ import annotations
+
+import argparse
+import json
+import sys
+import warnings
+from dataclasses import dataclass
+from pathlib import Path
+from typing import Any, Dict, List, Optional, Sequence, Tuple
+
+import joblib
+import numpy as np
+import pandas as pd
+import torch
+from lightgbm import LGBMClassifier
+from scipy import sparse
+from sklearn.metrics import (
+    accuracy_score,
+    classification_report,
+    confusion_matrix,
+    f1_score,
+    precision_recall_fscore_support,
+    roc_auc_score,
+)
+from sklearn.model_selection import GroupKFold, GroupShuffleSplit
+from sklearn.utils.class_weight import compute_class_weight
+from sklearn.svm import SVC
+from torch.utils.data import DataLoader
+from tqdm.auto import tqdm
+from xgboost import XGBClassifier
+
+sys.path.append(str(Path(__file__).resolve().parent))
+
+from data_preprocessing import DataPreprocessor  # type: ignore
+from training_runtime import PauseRequested, StopRequested, TrainingRunController  # type: ignore
+
+
+warnings.filterwarnings(
+    "ignore",
+    message="X does not have valid feature names, but LGBMClassifier was fitted with feature names",
+)
+
+
+ROOT = Path(__file__).resolve().parents[1]
+MODEL_DIR = ROOT / "models" / "saved"
+RUNS_DIR = MODEL_DIR / "training_runs"
+EVAL_DIR = ROOT / "evaluation_results"
+MODEL_METRICS_DIR = EVAL_DIR / "model_metrics"
+CLASS_NAMES_DEFAULT = ["HC", "PD", "SWEDD", "PRODROMAL"]
+TRADITIONAL_MODELS = ("lightgbm", "xgboost", "svm")
+TRANSFORMER_MODELS = ("pubmedbert", "biogpt", "clinical_t5")
+ALL_BASE_MODELS = TRADITIONAL_MODELS + TRANSFORMER_MODELS
+TRANSFORMER_SELECTION_METRIC = "val_f1"
+DEFAULT_TRANSFORMER_LOSS = "focal"
+DEFAULT_FOCAL_GAMMA = 1.5
+
+
+TRADITIONAL_SEARCH_SPACES: Dict[str, List[Dict[str, Any]]] = {
+    "lightgbm": [
+        {"n_estimators": 350, "learning_rate": 0.03, "max_depth": -1, "num_leaves": 31, "min_child_samples": 20, "subsample": 0.90, "colsample_bytree": 0.90},
+        {"n_estimators": 500, "learning_rate": 0.02, "max_depth": -1, "num_leaves": 63, "min_child_samples": 25, "subsample": 0.85, "colsample_bytree": 0.85},
+        {"n_estimators": 450, "learning_rate": 0.025, "max_depth": 10, "num_leaves": 47, "min_child_samples": 15, "subsample": 0.90, "colsample_bytree": 0.80},
+        {"n_estimators": 650, "learning_rate": 0.015, "max_depth": 12, "num_leaves": 95, "min_child_samples": 30, "subsample": 0.80, "colsample_bytree": 0.80},
+        {"n_estimators": 800, "learning_rate": 0.012, "max_depth": -1, "num_leaves": 127, "min_child_samples": 20, "subsample": 0.85, "colsample_bytree": 0.80},
+        {"n_estimators": 420, "learning_rate": 0.03, "max_depth": 8, "num_leaves": 63, "min_child_samples": 10, "subsample": 0.95, "colsample_bytree": 0.85},
+    ],
+    "xgboost": [
+        {"n_estimators": 300, "learning_rate": 0.05, "max_depth": 6, "min_child_weight": 2, "subsample": 0.90, "colsample_bytree": 0.90, "gamma": 0.0, "reg_lambda": 1.0},
+        {"n_estimators": 450, "learning_rate": 0.03, "max_depth": 5, "min_child_weight": 1, "subsample": 0.85, "colsample_bytree": 0.85, "gamma": 0.05, "reg_lambda": 1.0},
+        {"n_estimators": 500, "learning_rate": 0.025, "max_depth": 7, "min_child_weight": 3, "subsample": 0.80, "colsample_bytree": 0.80, "gamma": 0.10, "reg_lambda": 1.5},
+        {"n_estimators": 350, "learning_rate": 0.04, "max_depth": 4, "min_child_weight": 1, "subsample": 0.95, "colsample_bytree": 0.90, "gamma": 0.0, "reg_lambda": 0.8},
+        {"n_estimators": 650, "learning_rate": 0.015, "max_depth": 8, "min_child_weight": 2, "subsample": 0.85, "colsample_bytree": 0.85, "gamma": 0.05, "reg_lambda": 1.2},
+        {"n_estimators": 520, "learning_rate": 0.02, "max_depth": 6, "min_child_weight": 4, "subsample": 0.90, "colsample_bytree": 0.80, "gamma": 0.08, "reg_lambda": 1.8},
+    ],
+    "svm": [
+        {"C": 6.0, "gamma": "scale", "kernel": "rbf"},
+        {"C": 8.0, "gamma": "scale", "kernel": "rbf"},
+        {"C": 10.0, "gamma": 0.01, "kernel": "rbf"},
+        {"C": 12.0, "gamma": 0.005, "kernel": "rbf"},
+        {"C": 14.0, "gamma": "scale", "kernel": "rbf"},
+        {"C": 16.0, "gamma": 0.003, "kernel": "rbf"},
+    ],
+}
+
+
+TRANSFORMER_TRIALS: Dict[str, List[Dict[str, Any]]] = {
+    "pubmedbert": [
+        {"model_kwargs": {"dropout": 0.10, "freeze_bert": False, "train_encoder_layers": 8}, "optimizer": {"lr": 8.0e-6, "weight_decay": 0.02}, "grad_accum": 2},
+        {"model_kwargs": {"dropout": 0.08, "freeze_bert": False, "train_encoder_layers": 6}, "optimizer": {"lr": 1.0e-5, "weight_decay": 0.02}, "grad_accum": 2},
+        {"model_kwargs": {"dropout": 0.12, "freeze_bert": False, "train_encoder_layers": 4}, "optimizer": {"lr": 1.5e-5, "weight_decay": 0.01}, "grad_accum": 2},
+        {"model_kwargs": {"dropout": 0.18, "freeze_bert": True}, "optimizer": {"lr": 2.5e-5, "weight_decay": 0.01}, "grad_accum": 2},
+        {"model_kwargs": {"dropout": 0.06, "freeze_bert": False, "train_encoder_layers": 10}, "optimizer": {"lr": 6.0e-6, "weight_decay": 0.03}, "grad_accum": 2},
+        {"model_kwargs": {"dropout": 0.10, "freeze_bert": False, "train_encoder_layers": 12}, "optimizer": {"lr": 5.0e-6, "weight_decay": 0.03}, "grad_accum": 1},
+    ],
+    "biogpt": [
+        {"model_kwargs": {"dropout": 0.10, "train_decoder_layers": 4}, "optimizer": {"lr": 3.0e-5, "weight_decay": 0.01}, "grad_accum": 8},
+        {"model_kwargs": {"dropout": 0.15, "train_decoder_layers": 6}, "optimizer": {"lr": 2.5e-5, "weight_decay": 0.02}, "grad_accum": 8},
+        {"model_kwargs": {"dropout": 0.20, "train_decoder_layers": 8}, "optimizer": {"lr": 2.0e-5, "weight_decay": 0.02}, "grad_accum": 10},
+        {"model_kwargs": {"dropout": 0.12, "train_decoder_layers": 10}, "optimizer": {"lr": 1.5e-5, "weight_decay": 0.02}, "grad_accum": 12},
+        {"model_kwargs": {"dropout": 0.10, "train_decoder_layers": 12}, "optimizer": {"lr": 1.0e-5, "weight_decay": 0.02}, "grad_accum": 8},
+        {"model_kwargs": {"dropout": 0.08, "train_decoder_layers": 10}, "optimizer": {"lr": 1.2e-5, "weight_decay": 0.015}, "grad_accum": 6},
+    ],
+    "clinical_t5": [
+        {"model_kwargs": {"dropout": 0.10, "freeze_encoder": False}, "optimizer": {"lr": 2.0e-5, "weight_decay": 0.01}, "grad_accum": 8},
+        {"model_kwargs": {"dropout": 0.15, "freeze_encoder": False}, "optimizer": {"lr": 1.5e-5, "weight_decay": 0.02}, "grad_accum": 8},
+        {"model_kwargs": {"dropout": 0.20, "freeze_encoder": True}, "optimizer": {"lr": 2.5e-5, "weight_decay": 0.01}, "grad_accum": 6},
+        {"model_kwargs": {"dropout": 0.08, "freeze_encoder": False}, "optimizer": {"lr": 1.0e-5, "weight_decay": 0.02}, "grad_accum": 10},
+        {"model_kwargs": {"dropout": 0.12, "freeze_encoder": False}, "optimizer": {"lr": 8.0e-6, "weight_decay": 0.02}, "grad_accum": 8},
+        {"model_kwargs": {"dropout": 0.10, "freeze_encoder": False}, "optimizer": {"lr": 1.2e-5, "weight_decay": 0.015}, "grad_accum": 6},
+    ],
+}
+
+
+@dataclass
+class TrainingBundle:
+    X_train_dense: np.ndarray
+    X_test_dense: np.ndarray
+    y_train: np.ndarray
+    y_test: np.ndarray
+    train_groups: np.ndarray
+    test_groups: np.ndarray
+    feature_names: List[str]
+    class_mapping: Dict[str, int]
+    class_names: List[str]
+    preprocessor: Any
+
+
+@dataclass(frozen=True)
+class GPUExecutionProfile:
+    name: str
+    train_batch_by_model: Dict[str, int]
+    eval_batch_by_model: Dict[str, int]
+    grad_accum_cap_by_model: Dict[str, int]
+    num_workers: int
+    prefetch_factor: int
+    persistent_workers: bool
+    notes: str
+
+
+class FocalLoss(torch.nn.Module):
+    """Class-weighted focal loss for imbalanced multi-class classification."""
+
+    def __init__(
+        self,
+        class_weights: Optional[torch.Tensor] = None,
+        gamma: float = DEFAULT_FOCAL_GAMMA,
+        reduction: str = "mean",
+    ) -> None:
+        super().__init__()
+        if reduction not in {"mean", "sum", "none"}:
+            raise ValueError(f"Unsupported focal-loss reduction: {reduction}")
+        self.gamma = float(gamma)
+        self.reduction = reduction
+        if class_weights is not None:
+            normalized = class_weights.float() / class_weights.float().mean().clamp_min(1e-6)
+            self.register_buffer("class_weights", normalized)
+        else:
+            self.register_buffer("class_weights", None)
+
+    def forward(self, logits: torch.Tensor, targets: torch.Tensor) -> torch.Tensor:
+        log_probs = torch.nn.functional.log_softmax(logits, dim=1)
+        probs = log_probs.exp()
+        target_log_probs = log_probs.gather(1, targets.unsqueeze(1)).squeeze(1)
+        target_probs = probs.gather(1, targets.unsqueeze(1)).squeeze(1).clamp_min(1e-6)
+        ce_loss = -target_log_probs
+        focal_factor = (1.0 - target_probs).pow(self.gamma)
+        loss = focal_factor * ce_loss
+        if self.class_weights is not None:
+            loss = self.class_weights[targets] * loss
+        if self.reduction == "mean":
+            return loss.mean()
+        if self.reduction == "sum":
+            return loss.sum()
+        return loss
+
+
+def _json_ready(value: Any) -> Any:
+    if isinstance(value, dict):
+        return {str(key): _json_ready(val) for key, val in value.items()}
+    if isinstance(value, (list, tuple)):
+        return [_json_ready(item) for item in value]
+    if isinstance(value, np.ndarray):
+        return value.tolist()
+    if isinstance(value, np.generic):
+        return value.item()
+    if isinstance(value, Path):
+        return str(value)
+    return value
+
+
+def _as_dense(matrix: Any) -> np.ndarray:
+    if sparse.issparse(matrix):
+        return matrix.toarray()
+    return np.asarray(matrix)
+
+
+def _candidate_file_paths() -> List[Path]:
+    return [
+        ROOT / "PPMI_Curated_Data_Cut_Public_20240129.csv",
+        ROOT / "PPMI_Curated_Data_Cut_Public_20241211.csv",
+        ROOT / "PPMI_Curated_Data_Cut_Public_20250321.csv",
+        ROOT / "PPMI_Curated_Data_Cut_Public_20250714.csv",
+    ]
+
+
+def _prepare_training_bundle() -> TrainingBundle:
+    preprocessor = DataPreprocessor()
+    X_train, X_test, y_train, y_test = preprocessor.prepare_data(
+        _candidate_file_paths(),
+        test_size=0.2,
+        use_patient_split=True,
+    )
+    train_df, test_df = preprocessor.get_split_frames()
+    feature_names = preprocessor.get_feature_names()
+    class_mapping = preprocessor.get_class_mapping()
+    class_names = [None] * len(class_mapping)
+    for label, idx in class_mapping.items():
+        class_names[idx] = label
+
+    EVAL_DIR.mkdir(parents=True, exist_ok=True)
+    np.savez(
+        EVAL_DIR / "leak_free_split.npz",
+        X_train=X_train,
+        X_test=X_test,
+        y_train=y_train,
+        y_test=y_test,
+    )
+    joblib.dump(
+        {
+            "feature_names": feature_names,
+            "class_mapping": class_mapping,
+            "train_groups": train_df["PATNO"].to_numpy(),
+            "test_groups": test_df["PATNO"].to_numpy(),
+        },
+        EVAL_DIR / "leak_free_split_meta.joblib",
+    )
+
+    MODEL_DIR.mkdir(parents=True, exist_ok=True)
+    joblib.dump(preprocessor.get_preprocessor(), MODEL_DIR / "traditional_preprocessor.joblib")
+    (MODEL_DIR / "traditional_class_mapping.json").write_text(
+        json.dumps(_json_ready(class_mapping), indent=2),
+        encoding="utf-8",
+    )
+
+    return TrainingBundle(
+        X_train_dense=_as_dense(X_train).astype(np.float32),
+        X_test_dense=_as_dense(X_test).astype(np.float32),
+        y_train=np.asarray(y_train, dtype=np.int64),
+        y_test=np.asarray(y_test, dtype=np.int64),
+        train_groups=train_df["PATNO"].to_numpy(),
+        test_groups=test_df["PATNO"].to_numpy(),
+        feature_names=list(feature_names),
+        class_mapping=dict(class_mapping),
+        class_names=[
+            str(name if name is not None else CLASS_NAMES_DEFAULT[idx])
+            for idx, name in enumerate(class_names)
+        ],
+        preprocessor=preprocessor.get_preprocessor(),
+    )
+
+
+def _compute_class_weight_dict(y_train: np.ndarray) -> Dict[int, float]:
+    classes = np.unique(y_train)
+    weights = compute_class_weight(class_weight="balanced", classes=classes, y=y_train)
+    return {int(cls): float(weight) for cls, weight in zip(classes, weights)}
+
+
+def _build_transformer_criterion(
+    loss_name: str,
+    class_weights_tensor: torch.Tensor,
+    focal_gamma: float,
+) -> Tuple[torch.nn.Module, str]:
+    normalized_name = (loss_name or DEFAULT_TRANSFORMER_LOSS).strip().lower()
+    if normalized_name in {"ce", "cross_entropy", "cross-entropy"}:
+        return torch.nn.CrossEntropyLoss(weight=class_weights_tensor), "cross_entropy"
+    if normalized_name == "focal":
+        return FocalLoss(class_weights=class_weights_tensor, gamma=focal_gamma), f"focal(gamma={focal_gamma:.2f})"
+    raise ValueError(f"Unsupported transformer loss: {loss_name}")
+
+
+def _is_better_validation_epoch(
+    val_f1: float,
+    val_loss: float,
+    best_val_f1: float,
+    best_val_loss: float,
+    min_delta: float = 1e-4,
+) -> bool:
+    if val_f1 > (best_val_f1 + min_delta):
+        return True
+    if abs(val_f1 - best_val_f1) <= min_delta and val_loss < (best_val_loss - min_delta):
+        return True
+    return False
+
+
+def _evaluate_predictions(
+    model_name: str,
+    model_type: str,
+    y_true: np.ndarray,
+    y_pred: np.ndarray,
+    probabilities: Optional[np.ndarray],
+    class_names: Sequence[str],
+) -> Dict[str, Any]:
+    normalized_class_names = [str(name) for name in class_names]
+    accuracy = accuracy_score(y_true, y_pred)
+    precision, recall, f1, _ = precision_recall_fscore_support(y_true, y_pred, average="weighted", zero_division=0)
+    report = classification_report(y_true, y_pred, target_names=normalized_class_names, zero_division=0)
+    cm = confusion_matrix(y_true, y_pred)
+    auroc = None
+    if probabilities is not None:
+        try:
+            auroc = roc_auc_score(y_true, probabilities, multi_class="ovr", average="weighted")
+        except ValueError:
+            auroc = None
+    return {
+        "Model": model_name,
+        "Type": model_type,
+        "Accuracy": float(accuracy),
+        "Precision": float(precision),
+        "Recall": float(recall),
+        "F1_Score": float(f1),
+        "AUROC": None if auroc is None else float(auroc),
+        "classification_report": report,
+        "confusion_matrix": cm.tolist(),
+    }
+
+
+def _traditional_model_builder(model_name: str, params: Dict[str, Any], class_weight_dict: Dict[int, float], num_classes: int):
+    if model_name == "lightgbm":
+        return LGBMClassifier(random_state=42, objective="multiclass", num_class=num_classes, verbose=-1, class_weight=class_weight_dict, **params)
+    if model_name == "xgboost":
+        return XGBClassifier(random_state=42, objective="multi:softprob", num_class=num_classes, eval_metric="mlogloss", verbosity=0, **params)
+    if model_name == "svm":
+        return SVC(random_state=42, probability=True, decision_function_shape="ovr", class_weight=class_weight_dict, cache_size=2048, **params)
+    raise ValueError(f"Unsupported traditional model: {model_name}")
+
+
+def _run_grouped_traditional_search(
+    model_name: str,
+    search_space: List[Dict[str, Any]],
+    bundle: TrainingBundle,
+    controller: TrainingRunController,
+    max_trials: int,
+) -> Tuple[Dict[str, Any], Dict[str, Any], Path]:
+    class_weight_dict = _compute_class_weight_dict(bundle.y_train)
+    checkpoint_state = controller.load_checkpoint_state(
+        model_name,
+        {
+            "phase": "search",
+            "trials": {},
+            "best_trial_index": None,
+            "best_score": None,
+        },
+    )
+
+    search_space = list(search_space[:max_trials])
+    n_splits = max(2, min(3, len(np.unique(bundle.train_groups))))
+    splitter = GroupKFold(n_splits=n_splits)
+
+    controller.mark_running("traditional_search", model_name=model_name)
+    controller.update_model_state(
+        model_name,
+        status="running",
+        family="traditional",
+        trials_total=len(search_space),
+    )
+
+    for trial_index, params in enumerate(search_space):
+        trial_key = str(trial_index)
+        trial_state = checkpoint_state["trials"].setdefault(
+            trial_key,
+            {
+                "params": params,
+                "fold_scores": [],
+                "status": "pending",
+            },
+        )
+        if trial_state.get("status") == "completed":
+            continue
+
+        controller.mark_running("traditional_trial", model_name=model_name, extra={"trial_index": trial_index})
+
+        for fold_index, (train_idx, val_idx) in enumerate(
+            splitter.split(bundle.X_train_dense, bundle.y_train, groups=bundle.train_groups)
+        ):
+            if fold_index < len(trial_state["fold_scores"]):
+                continue
+
+            train_features = np.asarray(bundle.X_train_dense[train_idx], dtype=np.float32)
+            val_features = np.asarray(bundle.X_train_dense[val_idx], dtype=np.float32)
+            train_targets = np.asarray(bundle.y_train[train_idx], dtype=np.int64)
+            val_targets = np.asarray(bundle.y_train[val_idx], dtype=np.int64)
+
+            estimator = _traditional_model_builder(
+                model_name,
+                params,
+                class_weight_dict=class_weight_dict,
+                num_classes=len(bundle.class_names),
+            )
+            estimator.fit(train_features, train_targets)
+            preds = estimator.predict(val_features)
+            score = f1_score(val_targets, preds, average="weighted")
+
+            trial_state["fold_scores"].append(float(score))
+            trial_state["status"] = "running"
+            checkpoint_state["phase"] = "search"
+            controller.save_checkpoint_state(model_name, _json_ready(checkpoint_state))
+            controller.update_model_state(
+                model_name,
+                current_trial=trial_index,
+                current_fold=fold_index,
+                latest_fold_score=float(score),
+            )
+            controller.raise_if_requested()
+
+        trial_state["mean_score"] = float(np.mean(trial_state["fold_scores"]))
+        trial_state["status"] = "completed"
+        controller.append_trial_result(
+            model_name,
+            {
+                "trial_index": trial_index,
+                "params": params,
+                "mean_score": trial_state["mean_score"],
+            },
+        )
+        if checkpoint_state.get("best_score") is None or trial_state["mean_score"] > checkpoint_state["best_score"]:
+            checkpoint_state["best_score"] = float(trial_state["mean_score"])
+            checkpoint_state["best_trial_index"] = trial_index
+        controller.save_checkpoint_state(model_name, _json_ready(checkpoint_state))
+        controller.raise_if_requested()
+
+    if checkpoint_state.get("best_trial_index") is None:
+        raise RuntimeError(f"No completed trials were recorded for {model_name}")
+
+    best_index = int(checkpoint_state["best_trial_index"])
+    best_params = dict(search_space[best_index])
+
+    checkpoint_state["phase"] = "fit_final"
+    controller.save_checkpoint_state(model_name, _json_ready(checkpoint_state))
+    controller.mark_running("traditional_fit_final", model_name=model_name)
+
+    best_model = _traditional_model_builder(
+        model_name,
+        best_params,
+        class_weight_dict=class_weight_dict,
+        num_classes=len(bundle.class_names),
+    )
+    best_model.fit(
+        np.asarray(bundle.X_train_dense, dtype=np.float32),
+        np.asarray(bundle.y_train, dtype=np.int64),
+    )
+
+    artifact_path = MODEL_DIR / f"{model_name}_model.joblib"
+    joblib.dump(best_model, artifact_path)
+
+    test_features = np.asarray(bundle.X_test_dense, dtype=np.float32)
+    y_pred = best_model.predict(test_features)
+    probabilities = best_model.predict_proba(test_features) if hasattr(best_model, "predict_proba") else None
+    metrics = _evaluate_predictions(
+        model_name=model_name.replace("_", " ").title().replace("Svm", "SVM").replace("Lightgbm", "LightGBM").replace("Xgboost", "XGBoost"),
+        model_type="Traditional ML",
+        y_true=bundle.y_test,
+        y_pred=y_pred,
+        probabilities=probabilities,
+        class_names=bundle.class_names,
+    )
+
+    checkpoint_state["phase"] = "complete"
+    checkpoint_state["artifact_path"] = str(artifact_path)
+    checkpoint_state["best_params"] = best_params
+    checkpoint_state["metrics"] = {
+        key: value for key, value in metrics.items()
+        if key in {"Accuracy", "Precision", "Recall", "F1_Score", "AUROC", "Model", "Type"}
+    }
+    controller.save_checkpoint_state(model_name, _json_ready(checkpoint_state))
+    controller.update_model_state(
+        model_name,
+        status="completed",
+        artifact_path=str(artifact_path),
+        best_params=best_params,
+        best_cv_score=float(checkpoint_state["best_score"]),
+        metrics=checkpoint_state["metrics"],
+    )
+    return metrics, best_params, artifact_path
+
+
+def _build_rag_contexts(features: np.ndarray, feature_names: Sequence[str], doc_manager: Any) -> List[str]:
+    contexts: List[str] = []
+    for row in features:
+        feature_desc = {name: float(val) for name, val in zip(feature_names, row)}
+        query_parts = []
+        for symptom_key, col in {
+            "tremor": "sym_tremor",
+            "rigidity": "sym_rigid",
+            "bradykinesia": "sym_brady",
+            "postural instability": "sym_posins",
+        }.items():
+            severity = feature_desc.get(col, 0)
+            if severity > 0:
+                query_parts.append(f"{symptom_key} severity:{severity:.2f}")
+        if feature_desc.get("moca", 30) < 26:
+            query_parts.append("cognitive impairment")
+        if feature_desc.get("fampd", 0) > 0:
+            query_parts.append("family history Parkinson's disease")
+        query = "Parkinson's disease " + " ".join(query_parts)
+        passages = doc_manager.extract_relevant_passages(query, top_k=2)
+        if not passages:
+            contexts.append("")
+            continue
+        contexts.append(
+            " ".join(
+                f"From '{item.get('doc_title') or item.get('doc_id') or 'document'}' {item['text'][:240]}..."
+                for item in passages
+            )
+        )
+    return contexts
+
+
+def _encode_contexts_for_model(model_name: str, model: Any, contexts: Optional[List[str]]) -> Optional[Dict[str, torch.Tensor]]:
+    if not contexts:
+        return None
+    if model_name == "clinical_t5":
+        texts = [f"classify patient: {text}" for text in contexts]
+    else:
+        texts = contexts
+    encoded = model.tokenizer(
+        texts,
+        return_tensors="pt",
+        padding=True,
+        truncation=True,
+        max_length=512,
+    )
+    return {key: value.cpu() for key, value in encoded.items()}
+
+
+def _prepare_batch(batch: Any, device: torch.device):
+    if len(batch) == 3:
+        features, targets, contexts = batch
+        if isinstance(contexts, dict):
+            contexts = {key: value.to(device, non_blocking=True) for key, value in contexts.items()}
+        else:
+            contexts = list(contexts)
+    else:
+        features, targets = batch
+        contexts = None
+    return features.to(device), targets.to(device), contexts
+
+
+def _create_transformer_model(model_name: str, input_dim: int, num_classes: int, model_kwargs: Dict[str, Any]):
+    from models.medical_transformers import (  # type: ignore
+        BioMistralClassifier as BioGPTForTabular,
+        ClinicalT5Classifier as ClinicalT5ForTabular,
+        PubMedBERTClassifier as PubMedBERTForTabular,
+    )
+
+    if model_name == "pubmedbert":
+        return PubMedBERTForTabular(input_dim=input_dim, num_classes=num_classes, **model_kwargs)
+    if model_name == "biogpt":
+        return BioGPTForTabular(input_dim=input_dim, num_classes=num_classes, **model_kwargs)
+    if model_name == "clinical_t5":
+        return ClinicalT5ForTabular(input_dim=input_dim, num_classes=num_classes, **model_kwargs)
+    raise ValueError(f"Unsupported transformer model: {model_name}")
+
+
+def _torch_checkpoint_path(controller: TrainingRunController, model_name: str, trial_index: int) -> Path:
+    return controller.paths.checkpoints_dir / f"{model_name}_trial{trial_index}.pth"
+
+
+def _canonical_transformer_artifacts(model_name: str) -> List[Path]:
+    aliases = {
+        "pubmedbert": ["pubmedbert_transformer.pth", "pubmedbert.pth"],
+        "biogpt": ["biogpt_transformer.pth", "biogpt.pth", "biomistral.pth"],
+        "clinical_t5": ["clinical_t5_transformer.pth", "clinicalt5_transformer.pth", "clinical_t5.pth"],
+    }
+    return [MODEL_DIR / filename for filename in aliases[model_name]]
+
+
+def _run_transformer_search(
+    model_name: str,
+    trial_space: List[Dict[str, Any]],
+    bundle: TrainingBundle,
+    controller: TrainingRunController,
+    max_trials: int,
+    num_epochs: int,
+    patience: int,
+    use_rag: bool,
+    gpu_profile_name: str,
+    transformer_loss_name: str,
+    focal_gamma: float,
+) -> Tuple[Dict[str, Any], Dict[str, Any], Path]:
+    from models.transformer_models import TabularDataset  # type: ignore
+
+    device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
+    amp_enabled = device.type == "cuda" and model_name != "clinical_t5"
+    gpu_profile = _detect_gpu_execution_profile(requested_profile=gpu_profile_name)
+    train_bs = gpu_profile.train_batch_by_model.get(model_name, 8)
+    eval_bs = gpu_profile.eval_batch_by_model.get(model_name, max(train_bs * 2, 8))
+    loader_kwargs = _build_loader_kwargs(device, gpu_profile)
+
+    class_weights = compute_class_weight("balanced", classes=np.unique(bundle.y_train), y=bundle.y_train)
+    class_weights_tensor = torch.FloatTensor(class_weights).to(device)
+    criterion, criterion_label = _build_transformer_criterion(
+        transformer_loss_name,
+        class_weights_tensor,
+        focal_gamma,
+    )
+
+    split = GroupShuffleSplit(n_splits=1, test_size=0.15, random_state=42)
+    train_index, val_index = next(split.split(bundle.X_train_dense, bundle.y_train, groups=bundle.train_groups))
+
+    train_features = bundle.X_train_dense[train_index]
+    train_targets = bundle.y_train[train_index]
+    val_features = bundle.X_train_dense[val_index]
+    val_targets = bundle.y_train[val_index]
+
+    train_contexts = None
+    val_contexts = None
+    if use_rag:
+        from document_manager import DocumentManager  # type: ignore
+
+        doc_manager = DocumentManager(docs_dir=str(ROOT / "medical_docs"))
+        train_contexts = _build_rag_contexts(train_features, bundle.feature_names, doc_manager)
+        val_contexts = _build_rag_contexts(val_features, bundle.feature_names, doc_manager)
+
+    test_ds = TabularDataset(bundle.X_test_dense, bundle.y_test, bundle.feature_names, contexts=None)
+    test_loader = DataLoader(test_ds, batch_size=eval_bs, shuffle=False, **loader_kwargs)
+
+    checkpoint_state = controller.load_checkpoint_state(
+        model_name,
+        {
+            "phase": "search",
+            "trials": {},
+            "best_trial_index": None,
+            "best_f1": None,
+            "selection_metric": TRANSFORMER_SELECTION_METRIC,
+            "loss_name": criterion_label,
+            "focal_gamma": float(focal_gamma),
+        },
+    )
+    if (
+        checkpoint_state.get("selection_metric") != TRANSFORMER_SELECTION_METRIC
+        or checkpoint_state.get("loss_name") != criterion_label
+        or float(checkpoint_state.get("focal_gamma", focal_gamma)) != float(focal_gamma)
+    ):
+        print(
+            f"[TRANSFORMER] {model_name} resetting saved trial state to match "
+            f"selection={TRANSFORMER_SELECTION_METRIC} and loss={criterion_label}"
+        )
+        checkpoint_state = {
+            "phase": "search",
+            "trials": {},
+            "best_trial_index": None,
+            "best_f1": None,
+            "selection_metric": TRANSFORMER_SELECTION_METRIC,
+            "loss_name": criterion_label,
+            "focal_gamma": float(focal_gamma),
+        }
+        for stale_checkpoint in controller.paths.checkpoints_dir.glob(f"{model_name}_trial*.pth"):
+            stale_checkpoint.unlink(missing_ok=True)
+    trial_space = list(trial_space[:max_trials])
+
+    controller.mark_running("transformer_search", model_name=model_name)
+    controller.update_model_state(
+        model_name,
+        status="running",
+        family="transformer",
+        trials_total=len(trial_space),
+        device=str(device),
+        gpu_profile=gpu_profile.name,
+        train_batch_size=train_bs,
+        eval_batch_size=eval_bs,
+        num_workers=loader_kwargs.get("num_workers", 0),
+        gpu_notes=gpu_profile.notes,
+        selection_metric=TRANSFORMER_SELECTION_METRIC,
+        transformer_loss=criterion_label,
+    )
+
+    for trial_index, trial_cfg in enumerate(trial_space):
+        trial_key = str(trial_index)
+        trial_state = checkpoint_state["trials"].setdefault(
+            trial_key,
+            {
+                "config": trial_cfg,
+                "status": "pending",
+                "history": {"train_loss": [], "val_loss": [], "val_f1": [], "val_acc": [], "lr": []},
+                "best_val_loss": None,
+                "best_val_f1": None,
+                "best_epoch": None,
+                "selection_metric": TRANSFORMER_SELECTION_METRIC,
+                "loss_name": criterion_label,
+            },
+        )
+        if trial_state.get("status") == "completed":
+            continue
+
+        controller.mark_running("transformer_trial", model_name=model_name, extra={"trial_index": trial_index})
+
+        model = _create_transformer_model(
+            model_name,
+            input_dim=bundle.X_train_dense.shape[1],
+            num_classes=len(bundle.class_names),
+            model_kwargs=trial_cfg["model_kwargs"],
+        ).to(device)
+        encoded_train_contexts = _encode_contexts_for_model(model_name, model, train_contexts)
+        encoded_val_contexts = _encode_contexts_for_model(model_name, model, val_contexts)
+        train_ds = TabularDataset(train_features, train_targets, bundle.feature_names, contexts=encoded_train_contexts)
+        val_ds = TabularDataset(val_features, val_targets, bundle.feature_names, contexts=encoded_val_contexts)
+        train_loader = DataLoader(train_ds, batch_size=train_bs, shuffle=True, **loader_kwargs)
+        val_loader = DataLoader(val_ds, batch_size=eval_bs, shuffle=False, **loader_kwargs)
+        optimizer = torch.optim.AdamW(filter(lambda p: p.requires_grad, model.parameters()), **trial_cfg["optimizer"])
+        scheduler = torch.optim.lr_scheduler.ReduceLROnPlateau(optimizer, mode="max", factor=0.5, patience=2, min_lr=1e-7)
+        scaler = torch.amp.GradScaler(device=device.type, enabled=amp_enabled)
+        grad_accum = min(
+            int(trial_cfg.get("grad_accum", 8)),
+            int(gpu_profile.grad_accum_cap_by_model.get(model_name, trial_cfg.get("grad_accum", 8))),
+        )
+        print(
+            f"[TRANSFORMER] {model_name} trial {trial_index + 1}/{len(trial_space)} "
+            f"| train_bs={train_bs} eval_bs={eval_bs} grad_accum={grad_accum} "
+            f"effective_batch={train_bs * grad_accum} "
+            f"train_steps={len(train_loader)} val_steps={len(val_loader)} "
+            f"| selection={TRANSFORMER_SELECTION_METRIC} loss={criterion_label}"
+        )
+        start_epoch = int(trial_state.get("next_epoch", 0))
+        early_stop_counter = int(trial_state.get("early_stop_counter", 0))
+        best_val_loss = float(trial_state["best_val_loss"]) if trial_state.get("best_val_loss") is not None else float("inf")
+        best_val_f1 = float(trial_state["best_val_f1"]) if trial_state.get("best_val_f1") is not None else float("-inf")
+        best_model_state_dict = None
+        torch_ckpt_path = _torch_checkpoint_path(controller, model_name, trial_index)
+
+        if torch_ckpt_path.exists():
+            saved = torch.load(torch_ckpt_path, map_location=device, weights_only=False)
+            saved_model_state = saved.get("model_state_dict", {})
+            checkpoint_has_non_finite = any(
+                torch.is_tensor(tensor) and not torch.isfinite(tensor).all()
+                for tensor in saved_model_state.values()
+            )
+            if (
+                saved.get("selection_metric") == TRANSFORMER_SELECTION_METRIC
+                and saved.get("loss_name") == criterion_label
+                and float(saved.get("focal_gamma", focal_gamma)) == float(focal_gamma)
+                and not checkpoint_has_non_finite
+            ):
+                model.load_state_dict(saved["model_state_dict"])
+                optimizer.load_state_dict(saved["optimizer_state_dict"])
+                scheduler.load_state_dict(saved["scheduler_state_dict"])
+                if saved.get("scaler_state_dict") and device.type == "cuda":
+                    scaler.load_state_dict(saved["scaler_state_dict"])
+                start_epoch = int(saved.get("epoch", -1)) + 1
+                early_stop_counter = int(saved.get("early_stop_counter", early_stop_counter))
+                best_val_loss = float(saved.get("best_val_loss", best_val_loss))
+                best_val_f1 = float(saved.get("best_val_f1", best_val_f1))
+                best_model_state_dict = saved.get("best_model_state_dict")
+                trial_state["history"] = saved.get("history", trial_state["history"])
+            else:
+                print(
+                    f"[TRANSFORMER] {model_name} trial {trial_index + 1} ignoring stale "
+                    f"checkpoint with incompatible selection/loss settings or non-finite weights"
+                )
+                start_epoch = 0
+                early_stop_counter = 0
+                best_val_loss = float("inf")
+                best_val_f1 = float("-inf")
+                best_model_state_dict = None
+                trial_state["history"] = {"train_loss": [], "val_loss": [], "val_f1": [], "val_acc": [], "lr": []}
+                torch_ckpt_path.unlink(missing_ok=True)
+
+        for epoch in range(start_epoch, num_epochs):
+            model.train()
+            running_train_loss = 0.0
+            processed_train_batches = 0
+            optimizer.zero_grad(set_to_none=True)
+            if device.type == "cuda":
+                torch.cuda.reset_peak_memory_stats(device)
+
+            progress = tqdm(
+                train_loader,
+                total=len(train_loader),
+                dynamic_ncols=True,
+                leave=False,
+                desc=f"{model_name} t{trial_index + 1}/{len(trial_space)} e{epoch + 1}/{num_epochs}",
+            )
+            for batch_index, batch in enumerate(progress):
+                features, targets, contexts = _prepare_batch(batch, device)
+                with torch.amp.autocast(device_type=device.type, enabled=amp_enabled):
+                    outputs = model(features, contexts)
+                    loss = criterion(outputs, targets) / grad_accum
+
+                if not torch.isfinite(outputs).all() or not torch.isfinite(loss):
+                    optimizer.zero_grad(set_to_none=True)
+                    continue
+
+                scaler.scale(loss).backward()
+
+                if (batch_index + 1) % grad_accum == 0 or (batch_index + 1) == len(train_loader):
+                    scaler.unscale_(optimizer)
+                    torch.nn.utils.clip_grad_norm_(model.parameters(), max_norm=1.0)
+                    scaler.step(optimizer)
+                    scaler.update()
+                    optimizer.zero_grad(set_to_none=True)
+
+                running_train_loss += float(loss.item() * grad_accum)
+                processed_train_batches += 1
+                current_loss = float(loss.item() * grad_accum)
+                current_lr = float(optimizer.param_groups[0]["lr"])
+                if device.type == "cuda":
+                    gpu_mem_gb = torch.cuda.memory_allocated(device) / 1024**3
+                    progress.set_postfix(loss=f"{current_loss:.4f}", lr=f"{current_lr:.2e}", gpu_gb=f"{gpu_mem_gb:.2f}")
+                else:
+                    progress.set_postfix(loss=f"{current_loss:.4f}", lr=f"{current_lr:.2e}")
+
+            progress.close()
+
+            avg_train_loss = running_train_loss / max(processed_train_batches, 1)
+
+            model.eval()
+            running_val_loss = 0.0
+            processed_val_batches = 0
+            all_preds: List[int] = []
+            all_targets: List[int] = []
+            with torch.no_grad():
+                for batch in val_loader:
+                    features, targets, contexts = _prepare_batch(batch, device)
+                    with torch.amp.autocast(device_type=device.type, enabled=amp_enabled):
+                        outputs = model(features, contexts)
+                        val_loss = criterion(outputs, targets)
+                    if not torch.isfinite(outputs).all() or not torch.isfinite(val_loss):
+                        continue
+                    running_val_loss += float(val_loss.item())
+                    processed_val_batches += 1
+                    preds = torch.argmax(outputs, dim=1)
+                    all_preds.extend(preds.cpu().numpy().tolist())
+                    all_targets.extend(targets.cpu().numpy().tolist())
+
+            avg_val_loss = running_val_loss / max(processed_val_batches, 1)
+            if not all_preds:
+                val_acc = 0.0
+                val_f1 = 0.0
+                avg_val_loss = float("inf")
+            else:
+                val_acc = float(np.mean(np.asarray(all_preds) == np.asarray(all_targets)))
+                val_f1 = float(f1_score(all_targets, all_preds, average="weighted"))
+            scheduler.step(val_f1)
+
+            history = trial_state["history"]
+            history["train_loss"].append(float(avg_train_loss))
+            history["val_loss"].append(float(avg_val_loss))
+            history["val_f1"].append(val_f1)
+            history["val_acc"].append(val_acc)
+            history["lr"].append(float(optimizer.param_groups[0]["lr"]))
+            gpu_peak_gb = 0.0
+            if device.type == "cuda":
+                gpu_peak_gb = torch.cuda.max_memory_allocated(device) / 1024**3
+            print(
+                f"[TRANSFORMER] {model_name} trial {trial_index + 1}/{len(trial_space)} "
+                f"epoch {epoch + 1}/{num_epochs} "
+                f"train_loss={avg_train_loss:.4f} val_loss={avg_val_loss:.4f} "
+                f"val_f1={val_f1:.4f} val_acc={val_acc:.4f} "
+                f"lr={optimizer.param_groups[0]['lr']:.2e} "
+                f"gpu_peak_gb={gpu_peak_gb:.2f}"
+            )
+
+            if _is_better_validation_epoch(
+                val_f1=val_f1,
+                val_loss=avg_val_loss,
+                best_val_f1=best_val_f1,
+                best_val_loss=best_val_loss,
+            ):
+                best_val_loss = avg_val_loss
+                best_val_f1 = val_f1
+                trial_state["best_val_loss"] = float(best_val_loss)
+                trial_state["best_val_f1"] = float(best_val_f1)
+                trial_state["best_epoch"] = epoch
+                best_model_state_dict = {
+                    key: value.detach().cpu().clone()
+                    for key, value in model.state_dict().items()
+                }
+                early_stop_counter = 0
+            else:
+                early_stop_counter += 1
+
+            trial_state["status"] = "running"
+            trial_state["next_epoch"] = epoch + 1
+            trial_state["early_stop_counter"] = early_stop_counter
+            checkpoint_state["phase"] = "search"
+            controller.save_checkpoint_state(model_name, _json_ready(checkpoint_state))
+            torch.save(
+                {
+                    "epoch": epoch,
+                    "model_state_dict": model.state_dict(),
+                    "optimizer_state_dict": optimizer.state_dict(),
+                    "scheduler_state_dict": scheduler.state_dict(),
+                    "scaler_state_dict": scaler.state_dict() if device.type == "cuda" else None,
+                    "history": history,
+                    "early_stop_counter": early_stop_counter,
+                    "best_val_loss": best_val_loss,
+                    "best_val_f1": best_val_f1,
+                    "best_model_state_dict": best_model_state_dict,
+                    "selection_metric": TRANSFORMER_SELECTION_METRIC,
+                    "loss_name": criterion_label,
+                    "focal_gamma": float(focal_gamma),
+                },
+                torch_ckpt_path,
+            )
+            controller.update_model_state(
+                model_name,
+                current_trial=trial_index,
+                current_epoch=epoch,
+                grad_accum=grad_accum,
+                best_val_f1=float(best_val_f1),
+                best_val_loss=float(best_val_loss),
+                selection_metric=TRANSFORMER_SELECTION_METRIC,
+                transformer_loss=criterion_label,
+            )
+            controller.raise_if_requested()
+
+            if early_stop_counter >= patience:
+                print(
+                    f"[TRANSFORMER] {model_name} trial {trial_index + 1}/{len(trial_space)} "
+                    f"early-stopped after epoch {epoch + 1}"
+                )
+                break
+
+        trial_state["status"] = "completed"
+        controller.append_trial_result(
+            model_name,
+            {
+                "trial_index": trial_index,
+                "config": trial_cfg,
+                "best_val_f1": float(trial_state.get("best_val_f1") or 0.0),
+                "best_epoch": trial_state.get("best_epoch"),
+            },
+        )
+        best_so_far = checkpoint_state.get("best_f1")
+        if best_so_far is None or float(trial_state.get("best_val_f1") or -1.0) > float(best_so_far):
+            checkpoint_state["best_f1"] = float(trial_state["best_val_f1"])
+            checkpoint_state["best_trial_index"] = trial_index
+        controller.save_checkpoint_state(model_name, _json_ready(checkpoint_state))
+        controller.raise_if_requested()
+
+    if checkpoint_state.get("best_trial_index") is None:
+        raise RuntimeError(f"No completed transformer trial found for {model_name}")
+
+    best_trial_index = int(checkpoint_state["best_trial_index"])
+    best_trial_cfg = dict(trial_space[best_trial_index])
+    checkpoint_state["phase"] = "evaluate"
+    controller.save_checkpoint_state(model_name, _json_ready(checkpoint_state))
+
+    best_model = _create_transformer_model(
+        model_name,
+        input_dim=bundle.X_train_dense.shape[1],
+        num_classes=len(bundle.class_names),
+        model_kwargs=best_trial_cfg["model_kwargs"],
+    ).to(device)
+    best_state = torch.load(_torch_checkpoint_path(controller, model_name, best_trial_index), map_location=device, weights_only=False)
+    best_model.load_state_dict(best_state.get("best_model_state_dict") or best_state["model_state_dict"])
+    best_model.eval()
+
+    all_preds: List[int] = []
+    all_targets: List[int] = []
+    all_probabilities: List[np.ndarray] = []
+    with torch.no_grad():
+        for batch in test_loader:
+            features, targets, contexts = _prepare_batch(batch, device)
+            outputs = best_model(features, contexts)
+            probs = torch.softmax(outputs, dim=1)
+            preds = torch.argmax(outputs, dim=1)
+            all_targets.extend(targets.cpu().numpy().tolist())
+            all_preds.extend(preds.cpu().numpy().tolist())
+            all_probabilities.append(probs.cpu().numpy())
+
+    probabilities = np.vstack(all_probabilities) if all_probabilities else None
+    metrics = _evaluate_predictions(
+        model_name={"pubmedbert": "PubMedBERT", "biogpt": "BioGPT", "clinical_t5": "Clinical-T5"}[model_name],
+        model_type="Transformer",
+        y_true=np.asarray(all_targets),
+        y_pred=np.asarray(all_preds),
+        probabilities=probabilities,
+        class_names=bundle.class_names,
+    )
+
+    primary_artifact = _canonical_transformer_artifacts(model_name)[0]
+    for path in _canonical_transformer_artifacts(model_name):
+        torch.save(best_model.state_dict(), path)
+
+    checkpoint_state["phase"] = "complete"
+    checkpoint_state["artifact_path"] = str(primary_artifact)
+    checkpoint_state["best_trial_index"] = best_trial_index
+    checkpoint_state["best_config"] = best_trial_cfg
+    checkpoint_state["metrics"] = {
+        key: value for key, value in metrics.items()
+        if key in {"Accuracy", "Precision", "Recall", "F1_Score", "AUROC", "Model", "Type"}
+    }
+    controller.save_checkpoint_state(model_name, _json_ready(checkpoint_state))
+    controller.update_model_state(
+        model_name,
+        status="completed",
+        artifact_path=str(primary_artifact),
+        best_trial_index=best_trial_index,
+        best_config=best_trial_cfg,
+        metrics=checkpoint_state["metrics"],
+    )
+    return metrics, best_trial_cfg, primary_artifact
+
+
+def _train_ensemble(bundle: TrainingBundle, controller: TrainingRunController) -> Optional[Dict[str, Any]]:
+    from models.multimodal_ml import MultimodalEnsemble  # type: ignore
+
+    controller.mark_running("ensemble_training", model_name="ensemble")
+    controller.raise_if_requested()
+
+    ensemble = MultimodalEnsemble()
+    ensemble.load_traditional_models(model_dir=str(MODEL_DIR))
+    ensemble.load_transformer_models(model_dir=str(MODEL_DIR), input_dim=bundle.X_train_dense.shape[1], num_classes=len(bundle.class_names))
+    ensemble.train_ensemble(bundle.X_train_dense, bundle.y_train, ensemble_type="stacking")
+    ensemble.save_ensemble(str(MODEL_DIR / "multimodal_ensemble.joblib"))
+    results = ensemble.evaluate_ensemble(bundle.X_test_dense, bundle.y_test)
+
+    probabilities = results.get("probabilities")
+    metrics = _evaluate_predictions(
+        model_name="Multimodal Ensemble",
+        model_type="Ensemble",
+        y_true=bundle.y_test,
+        y_pred=np.asarray(results["predictions"]),
+        probabilities=np.asarray(probabilities) if probabilities is not None else None,
+        class_names=bundle.class_names,
+    )
+    controller.update_model_state(
+        "ensemble",
+        status="completed",
+        artifact_path=str(MODEL_DIR / "multimodal_ensemble.joblib"),
+        metrics={key: value for key, value in metrics.items() if key in {"Accuracy", "Precision", "Recall", "F1_Score", "AUROC", "Model", "Type"}},
+    )
+    return metrics
+
+
+def _write_metric_outputs(base_metrics: List[Dict[str, Any]], ensemble_metrics: Optional[Dict[str, Any]]) -> None:
+    EVAL_DIR.mkdir(parents=True, exist_ok=True)
+    MODEL_METRICS_DIR.mkdir(parents=True, exist_ok=True)
+
+    summary_rows = []
+    traditional_rows = []
+    transformer_rows = []
+    for metrics in base_metrics:
+        summary_rows.append({
+            "Model": metrics["Model"],
+            "Type": metrics["Type"],
+            "Accuracy": metrics["Accuracy"],
+            "Precision": metrics["Precision"],
+            "Recall": metrics["Recall"],
+            "F1_Score": metrics["F1_Score"],
+            "AUROC": metrics["AUROC"] if metrics["AUROC"] is not None else "",
+        })
+        if "traditional" in metrics["Type"].lower():
+            traditional_rows.append(metrics)
+        if "transformer" in metrics["Type"].lower():
+            transformer_rows.append(metrics)
+
+    summary_df = pd.DataFrame(summary_rows)
+    summary_df.to_csv(EVAL_DIR / "summary_metrics.csv", index=False)
+    summary_df.to_csv(MODEL_METRICS_DIR / "model_metrics_summary.csv", index=False)
+
+    (EVAL_DIR / "traditional_metrics_latest.json").write_text(json.dumps(_json_ready(traditional_rows), indent=2), encoding="utf-8")
+    (EVAL_DIR / "transformer_metrics_latest.json").write_text(json.dumps(_json_ready(transformer_rows), indent=2), encoding="utf-8")
+    if ensemble_metrics is not None:
+        (EVAL_DIR / "ensemble_metrics_latest.json").write_text(json.dumps(_json_ready(ensemble_metrics), indent=2), encoding="utf-8")
+
+
+def _write_run_manifest_snapshot(controller: TrainingRunController, base_metrics: List[Dict[str, Any]], ensemble_metrics: Optional[Dict[str, Any]]) -> None:
+    controller.write_metrics_file(
+        "final_metrics.json",
+        {
+            "base_metrics": _json_ready(base_metrics),
+            "ensemble_metrics": _json_ready(ensemble_metrics),
+        },
+    )
+
+
+def _parse_selected_models(raw: str) -> List[str]:
+    raw = (raw or "all").strip().lower()
+    if raw in {"all", "*"}:
+        return list(ALL_BASE_MODELS)
+    aliases = {
+        "lgbm": "lightgbm",
+        "lightgbm": "lightgbm",
+        "xgb": "xgboost",
+        "xgboost": "xgboost",
+        "svm": "svm",
+        "pubmed": "pubmedbert",
+        "pubmedbert": "pubmedbert",
+        "biogpt": "biogpt",
+        "bio": "biogpt",
+        "clinical": "clinical_t5",
+        "clinical_t5": "clinical_t5",
+        "t5": "clinical_t5",
+    }
+    selected: List[str] = []
+    for piece in [part.strip() for part in raw.split(",") if part.strip()]:
+        canonical = aliases.get(piece)
+        if canonical and canonical not in selected:
+            selected.append(canonical)
+    if not selected:
+        raise ValueError(f"No valid models were selected from '{raw}'")
+    return selected
+
+
+def _detect_gpu_execution_profile(
+    requested_profile: str = "auto",
+    cuda_available: Optional[bool] = None,
+    device_name: Optional[str] = None,
+    total_memory_gb: Optional[float] = None,
+) -> GPUExecutionProfile:
+    cuda_ready = torch.cuda.is_available() if cuda_available is None else cuda_available
+    if not cuda_ready:
+        return GPUExecutionProfile(
+            name="cpu",
+            train_batch_by_model={"pubmedbert": 4, "biogpt": 2, "clinical_t5": 2},
+            eval_batch_by_model={"pubmedbert": 8, "biogpt": 4, "clinical_t5": 4},
+            grad_accum_cap_by_model={"pubmedbert": 12, "biogpt": 16, "clinical_t5": 16},
+            num_workers=0,
+            prefetch_factor=2,
+            persistent_workers=False,
+            notes="CPU fallback profile.",
+        )
+
+    name = (device_name or torch.cuda.get_device_name(0)).lower()
+    memory_gb = float(total_memory_gb if total_memory_gb is not None else (torch.cuda.get_device_properties(0).total_memory / 1024**3))
+    requested = (requested_profile or "auto").strip().lower()
+
+    if requested in {"rtx-a4000", "a4000"} or (requested == "auto" and ("a4000" in name or memory_gb >= 15.0)):
+        return GPUExecutionProfile(
+            name="rtx-a4000",
+            train_batch_by_model={"pubmedbert": 32, "biogpt": 12, "clinical_t5": 12},
+            eval_batch_by_model={"pubmedbert": 64, "biogpt": 24, "clinical_t5": 24},
+            grad_accum_cap_by_model={"pubmedbert": 2, "biogpt": 5, "clinical_t5": 5},
+            num_workers=6,
+            prefetch_factor=4,
+            persistent_workers=True,
+            notes="Optimized for RTX A4000 / ~16 GB VRAM with larger per-step batches.",
+        )
+
+    if requested in {"high-vram", "16gb"} or (requested == "auto" and memory_gb >= 11.0):
+        return GPUExecutionProfile(
+            name="high-vram",
+            train_batch_by_model={"pubmedbert": 12, "biogpt": 8, "clinical_t5": 8},
+            eval_batch_by_model={"pubmedbert": 32, "biogpt": 16, "clinical_t5": 16},
+            grad_accum_cap_by_model={"pubmedbert": 6, "biogpt": 8, "clinical_t5": 8},
+            num_workers=2,
+            prefetch_factor=2,
+            persistent_workers=True,
+            notes="Generic 12 GB+ CUDA profile.",
+        )
+
+    return GPUExecutionProfile(
+        name="compat",
+        train_batch_by_model={"pubmedbert": 8, "biogpt": 6, "clinical_t5": 6},
+        eval_batch_by_model={"pubmedbert": 16, "biogpt": 8, "clinical_t5": 8},
+        grad_accum_cap_by_model={"pubmedbert": 8, "biogpt": 10, "clinical_t5": 10},
+        num_workers=0,
+        prefetch_factor=2,
+        persistent_workers=False,
+        notes="Compatibility profile for lower-VRAM GPUs.",
+    )
+
+
+def _build_loader_kwargs(device: torch.device, profile: GPUExecutionProfile) -> Dict[str, Any]:
+    kwargs: Dict[str, Any] = {
+        "num_workers": profile.num_workers if device.type == "cuda" else 0,
+        "pin_memory": device.type == "cuda",
+    }
+    if kwargs["num_workers"] > 0:
+        kwargs["persistent_workers"] = profile.persistent_workers
+        kwargs["prefetch_factor"] = profile.prefetch_factor
+    return kwargs
+
+
+def _configure_transformer_runtime(selected_models: Sequence[str], allow_cpu_transformers: bool, cuda_available: Optional[bool] = None) -> str:
+    cuda_ready = torch.cuda.is_available() if cuda_available is None else cuda_available
+    transformer_requested = any(model_name in TRANSFORMER_MODELS for model_name in selected_models)
+    if not transformer_requested:
+        return "not-applicable"
+    if not cuda_ready and not allow_cpu_transformers:
+        raise RuntimeError(
+            "Transformer retraining requires CUDA for this accuracy profile. "
+            "Use a CUDA-enabled PyTorch install/GPU, or rerun with --allow-cpu-transformers if you explicitly accept slower, lower-throughput training."
+        )
+    if cuda_ready:
+        torch.set_float32_matmul_precision("high")
+        if hasattr(torch.backends, "cuda") and hasattr(torch.backends.cuda, "matmul"):
+            torch.backends.cuda.matmul.allow_tf32 = True
+        if hasattr(torch.backends, "cudnn"):
+            torch.backends.cudnn.allow_tf32 = True
+            torch.backends.cudnn.benchmark = True
+        return "cuda"
+    return "cpu"
+
+
+def _run_training(args: argparse.Namespace) -> int:
+    selected_models = _parse_selected_models(args.models)
+    transformer_runtime = _configure_transformer_runtime(
+        selected_models,
+        allow_cpu_transformers=args.allow_cpu_transformers,
+    )
+    controller = TrainingRunController(RUNS_DIR, args.run_name)
+    controller.initialize(
+        selected_models=selected_models,
+        config={
+            "epochs": args.epochs,
+            "patience": args.patience,
+            "traditional_trials": args.traditional_trials,
+            "transformer_trials": args.transformer_trials,
+            "use_rag": args.use_rag,
+            "train_ensemble": not args.skip_ensemble,
+            "allow_cpu_transformers": args.allow_cpu_transformers,
+            "transformer_runtime": transformer_runtime,
+            "gpu_profile": args.gpu_profile,
+            "transformer_loss": args.transformer_loss,
+            "focal_gamma": args.focal_gamma,
+        },
+        resume=args.resume,
+    )
+
+    if args.dry_run:
+        controller.mark_paused("dry-run initialized")
+        print(json.dumps(controller.status_summary(), indent=2))
+        return 0
+
+    try:
+        controller.clear_stop()
+        if args.resume:
+            controller.clear_pause()
+
+        bundle = _prepare_training_bundle()
+        base_metrics: List[Dict[str, Any]] = []
+
+        for model_name in selected_models:
+            if model_name in TRADITIONAL_MODELS:
+                metrics, best_config, artifact_path = _run_grouped_traditional_search(
+                    model_name=model_name,
+                    search_space=TRADITIONAL_SEARCH_SPACES[model_name],
+                    bundle=bundle,
+                    controller=controller,
+                    max_trials=args.traditional_trials,
+                )
+            else:
+                metrics, best_config, artifact_path = _run_transformer_search(
+                    model_name=model_name,
+                    trial_space=TRANSFORMER_TRIALS[model_name],
+                    bundle=bundle,
+                    controller=controller,
+                    max_trials=args.transformer_trials,
+                    num_epochs=args.epochs,
+                    patience=args.patience,
+                    use_rag=args.use_rag,
+                    gpu_profile_name=args.gpu_profile,
+                    transformer_loss_name=args.transformer_loss,
+                    focal_gamma=args.focal_gamma,
+                )
+            base_metrics.append(metrics)
+            controller.update_model_state(
+                model_name,
+                best_config=_json_ready(best_config),
+                artifact_path=str(artifact_path),
+                metrics={key: value for key, value in metrics.items() if key in {"Accuracy", "Precision", "Recall", "F1_Score", "AUROC", "Model", "Type"}},
+            )
+
+        ensemble_metrics = None
+        should_train_ensemble = (not args.skip_ensemble) and set(selected_models) == set(ALL_BASE_MODELS)
+        if should_train_ensemble:
+            ensemble_metrics = _train_ensemble(bundle, controller)
+        elif not args.skip_ensemble:
+            controller.update_model_state(
+                "ensemble",
+                status="skipped",
+                reason="ensemble retraining is only automatic when all six base models are selected",
+            )
+
+        _write_metric_outputs(base_metrics, ensemble_metrics)
+        _write_run_manifest_snapshot(controller, base_metrics, ensemble_metrics)
+        controller.mark_completed()
+        print(json.dumps(controller.status_summary(), indent=2))
+        return 0
+    except PauseRequested as exc:
+        controller.mark_paused(str(exc))
+        print(json.dumps(controller.status_summary(), indent=2))
+        return 0
+    except StopRequested as exc:
+        controller.mark_stopped(str(exc))
+        print(json.dumps(controller.status_summary(), indent=2))
+        return 0
+    except KeyboardInterrupt:
+        controller.mark_paused("keyboard interrupt")
+        print(json.dumps(controller.status_summary(), indent=2))
+        return 1
+    except Exception as exc:
+        controller.mark_failed(str(exc))
+        raise
+
+
+def _run_status(args: argparse.Namespace) -> int:
+    controller = TrainingRunController(RUNS_DIR, args.run_name)
+    print(json.dumps(controller.status_summary(), indent=2))
+    return 0
+
+
+def _run_pause(args: argparse.Namespace) -> int:
+    controller = TrainingRunController(RUNS_DIR, args.run_name)
+    controller.request_pause()
+    controller.mark_paused("pause requested from CLI")
+    print(json.dumps(controller.status_summary(), indent=2))
+    return 0
+
+
+def _run_stop(args: argparse.Namespace) -> int:
+    controller = TrainingRunController(RUNS_DIR, args.run_name)
+    controller.request_stop()
+    controller.mark_stopped("stop requested from CLI")
+    print(json.dumps(controller.status_summary(), indent=2))
+    return 0
+
+
+def build_parser() -> argparse.ArgumentParser:
+    parser = argparse.ArgumentParser(description="Retrain the six base Parkinson's models with resumable checkpoints.")
+    subparsers = parser.add_subparsers(dest="command", required=True)
+
+    train_parser = subparsers.add_parser("train", help="Start a new training run.")
+    train_parser.add_argument("--run-name", default="full_retrain", help="Name of the training run.")
+    train_parser.add_argument("--models", default="all", help="Comma-separated list of models or 'all'.")
+    train_parser.add_argument("--epochs", type=int, default=30, help="Max epochs per transformer trial.")
+    train_parser.add_argument("--patience", type=int, default=8, help="Early-stopping patience for transformers.")
+    train_parser.add_argument("--traditional-trials", type=int, default=6, help="Number of traditional-model trials per model.")
+    train_parser.add_argument("--transformer-trials", type=int, default=6, help="Number of transformer trials per model.")
+    train_parser.add_argument("--use-rag", dest="use_rag", action="store_true", help="Build RAG contexts for transformer training.")
+    train_parser.add_argument("--no-rag", dest="use_rag", action="store_false", help="Disable RAG context enrichment for transformer training.")
+    train_parser.add_argument("--gpu-profile", default="auto", choices=["auto", "rtx-a4000", "high-vram", "compat"], help="CUDA loader/batch preset for transformer training.")
+    train_parser.add_argument("--transformer-loss", default=DEFAULT_TRANSFORMER_LOSS, choices=["cross_entropy", "focal"], help="Loss function for transformer fine-tuning.")
+    train_parser.add_argument("--focal-gamma", type=float, default=DEFAULT_FOCAL_GAMMA, help="Gamma value for focal loss when --transformer-loss focal is used.")
+    train_parser.add_argument("--allow-cpu-transformers", action="store_true", help="Allow transformer training on CPU when CUDA is unavailable.")
+    train_parser.add_argument("--skip-ensemble", action="store_true", help="Skip ensemble retraining after the six base models.")
+    train_parser.add_argument("--resume", action="store_true", help="Resume an existing run from saved state.")
+    train_parser.add_argument("--dry-run", action="store_true", help="Initialize the run manifest without training.")
+    train_parser.set_defaults(func=_run_training, use_rag=True)
+
+    resume_parser = subparsers.add_parser("resume", help="Resume a paused training run.")
+    resume_parser.add_argument("--run-name", default="full_retrain", help="Name of the training run.")
+    resume_parser.add_argument("--models", default="all", help="Comma-separated list of models or 'all'.")
+    resume_parser.add_argument("--epochs", type=int, default=30, help="Max epochs per transformer trial.")
+    resume_parser.add_argument("--patience", type=int, default=8, help="Early-stopping patience for transformers.")
+    resume_parser.add_argument("--traditional-trials", type=int, default=6, help="Number of traditional-model trials per model.")
+    resume_parser.add_argument("--transformer-trials", type=int, default=6, help="Number of transformer trials per model.")
+    resume_parser.add_argument("--use-rag", dest="use_rag", action="store_true", help="Build RAG contexts for transformer training.")
+    resume_parser.add_argument("--no-rag", dest="use_rag", action="store_false", help="Disable RAG context enrichment for transformer training.")
+    resume_parser.add_argument("--gpu-profile", default="auto", choices=["auto", "rtx-a4000", "high-vram", "compat"], help="CUDA loader/batch preset for transformer training.")
+    resume_parser.add_argument("--transformer-loss", default=DEFAULT_TRANSFORMER_LOSS, choices=["cross_entropy", "focal"], help="Loss function for transformer fine-tuning.")
+    resume_parser.add_argument("--focal-gamma", type=float, default=DEFAULT_FOCAL_GAMMA, help="Gamma value for focal loss when --transformer-loss focal is used.")
+    resume_parser.add_argument("--allow-cpu-transformers", action="store_true", help="Allow transformer training on CPU when CUDA is unavailable.")
+    resume_parser.add_argument("--skip-ensemble", action="store_true", help="Skip ensemble retraining after the six base models.")
+    resume_parser.add_argument("--dry-run", action="store_true", help="Load and print the current run state without training.")
+    resume_parser.set_defaults(func=_run_training, resume=True, use_rag=True)
+
+    pause_parser = subparsers.add_parser("pause", help="Request a graceful pause for a running training job.")
+    pause_parser.add_argument("--run-name", default="full_retrain", help="Name of the training run.")
+    pause_parser.set_defaults(func=_run_pause)
+
+    stop_parser = subparsers.add_parser("stop", help="Request a graceful stop for a running training job.")
+    stop_parser.add_argument("--run-name", default="full_retrain", help="Name of the training run.")
+    stop_parser.set_defaults(func=_run_stop)
+
+    status_parser = subparsers.add_parser("status", help="Show the current run manifest.")
+    status_parser.add_argument("--run-name", default="full_retrain", help="Name of the training run.")
+    status_parser.set_defaults(func=_run_status)
+
+    return parser
+
+
+def main(argv: Optional[Sequence[str]] = None) -> int:
+    parser = build_parser()
+    args = parser.parse_args(argv)
+    return int(args.func(args))
+
+
+if __name__ == "__main__":
+    raise SystemExit(main())

src/train_multimodal.py

@@ -0,0 +1,136 @@
+"""
+Train and evaluate multimodal machine learning models for Parkinson's disease classification.
+This script combines traditional ML, transformer models, and ensemble methods.
+"""
+
+import os
+import sys
+import numpy as np
+import pandas as pd
+import matplotlib.pyplot as plt
+import seaborn as sns
+from sklearn.metrics import classification_report, confusion_matrix
+import warnings
+warnings.filterwarnings('ignore')
+
+# Add src directory to path
+sys.path.append(os.path.join(os.path.dirname(__file__)))
+sys.path.append(os.path.join(os.path.dirname(__file__), 'models'))
+
+from data_preprocessing import DataPreprocessor
+from models.multimodal_ml import MultimodalEnsemble, AdvancedFeatureEngineering, create_multimodal_pipeline
+
+
+def main():
+    """Main function to run multimodal ML pipeline."""
+    print("Multimodal Machine Learning for Parkinson's Disease Classification")
+    print("=" * 70)
+
+    # --- Load and prepare data ---
+    preprocessor = DataPreprocessor()
+    base_dir = os.path.dirname(os.path.dirname(os.path.abspath(__file__)))
+    file_paths = [
+        os.path.join(base_dir, "PPMI_Curated_Data_Cut_Public_20241211.csv"),
+        os.path.join(base_dir, "PPMI_Curated_Data_Cut_Public_20250321.csv"),
+        os.path.join(base_dir, "PPMI_Curated_Data_Cut_Public_20250714.csv"),
+    ]
+
+    print("\nLoading and preparing data with patient-level split...")
+    X_train, X_test, y_train, y_test = preprocessor.prepare_data(
+        file_paths,
+        test_size=0.2,
+        use_patient_split=True,
+    )
+    print(f"  Train: {X_train.shape}, Test: {X_test.shape}")
+
+    # Convert to numpy arrays if they are DataFrames
+    if isinstance(X_train, pd.DataFrame):
+        X_train_np = X_train.values
+        X_test_np = X_test.values
+    else:
+        X_train_np = X_train
+        X_test_np = X_test
+
+    if isinstance(y_train, pd.Series):
+        y_train_np = y_train.values
+        y_test_np = y_test.values
+    else:
+        y_train_np = y_train
+        y_test_np = y_test
+
+    # Create multimodal pipeline
+    print("\nCreating multimodal ML pipeline...")
+    ensemble, results = create_multimodal_pipeline(X_train, X_test, y_train_np, y_test_np)
+
+    # Detailed evaluation of ensemble
+    if ensemble.ensemble_model is not None:
+        print("\nDetailed Ensemble Evaluation:")
+        print("-" * 40)
+
+        ensemble_results = ensemble.evaluate_ensemble(X_test, y_test_np)
+
+        print(f"Ensemble Accuracy: {ensemble_results['accuracy']:.4f}")
+        print("\nClassification Report:")
+        print(ensemble_results['classification_report'])
+
+        # Plot confusion matrix
+        os.makedirs('notebooks', exist_ok=True)
+        plt.figure(figsize=(10, 8))
+        sns.heatmap(ensemble_results['confusion_matrix'],
+                   annot=True, fmt='d', cmap='Blues',
+                   xticklabels=['HC', 'PD', 'SWEDD', 'PRODROMAL'],
+                   yticklabels=['HC', 'PD', 'SWEDD', 'PRODROMAL'])
+        plt.title('Multimodal Ensemble - Confusion Matrix')
+        plt.ylabel('True Label')
+        plt.xlabel('Predicted Label')
+        plt.tight_layout()
+        plt.savefig('notebooks/multimodal_confusion_matrix.png', dpi=300, bbox_inches='tight')
+        plt.close()
+
+        print("Confusion matrix saved to notebooks/multimodal_confusion_matrix.png")
+
+    # Advanced analysis
+    print("\nAdvanced Multimodal Analysis:")
+    print("-" * 40)
+
+    # Feature importance analysis (if available)
+    if hasattr(ensemble.ensemble_model, 'coef_'):
+        feature_importance = np.abs(ensemble.ensemble_model.coef_).mean(axis=0)
+        print("Top 10 most important ensemble features:")
+        top_indices = np.argsort(feature_importance)[-10:][::-1]
+        for i, idx in enumerate(top_indices):
+            print(f"{i+1:2d}. Feature {idx}: {feature_importance[idx]:.4f}")
+
+    # Model diversity analysis
+    print("\nModel Diversity Analysis:")
+    if len(ensemble.traditional_models) > 0 and len(ensemble.transformer_models) > 0:
+        trad_preds, _ = ensemble.get_traditional_predictions(X_test)
+        trans_preds, _ = ensemble.get_transformer_predictions(X_test_np)
+
+        if trad_preds and trans_preds:
+            trad_pred = list(trad_preds.values())[0]
+            trans_pred = list(trans_preds.values())[0]
+
+            agreement = np.mean(trad_pred == trans_pred)
+            print(f"Agreement between traditional and transformer models: {agreement:.4f}")
+
+    # Performance summary
+    print("\nFinal Performance Summary:")
+    print("-" * 40)
+    best_model = max(results.items(), key=lambda x: x[1])
+    print(f"Best performing model: {best_model[0]} (Accuracy: {best_model[1]:.4f})")
+
+    if 'Ensemble' in results:
+        ensemble_acc = results['Ensemble']
+        individual_accs = [acc for name, acc in results.items() if name != 'Ensemble']
+        if individual_accs:
+            avg_individual = np.mean(individual_accs)
+            improvement = ensemble_acc - avg_individual
+            print(f"Ensemble improvement over average individual model: {improvement:.4f}")
+
+    print("\nMultimodal ML pipeline completed successfully!")
+    print("Results and visualizations saved to notebooks/ directory")
+
+
+if __name__ == "__main__":
+    main()

src/train_traditional_models.py

@@ -0,0 +1,187 @@
+"""
+Train and evaluate traditional ML models (LightGBM, XGBoost, SVM) on PPMI dataset
+with patient-level split to prevent data leakage.
+"""
+from data_preprocessing import DataPreprocessor
+from models.traditional_ml import TraditionalMLModels
+import pandas as pd
+import numpy as np
+import matplotlib.pyplot as plt
+import seaborn as sns
+import os
+from lightgbm import LGBMClassifier
+from xgboost import XGBClassifier
+from sklearn.svm import SVC
+
+
+def plot_confusion_matrices(results, save_dir='../notebooks'):
+    """Plot confusion matrices for all models."""
+    os.makedirs(save_dir, exist_ok=True)
+
+    for model_name, result in results.items():
+        plt.figure(figsize=(8, 6))
+        sns.heatmap(result['confusion_matrix'], annot=True, fmt='d', cmap='Blues')
+        plt.title(f'Confusion Matrix - {model_name}')
+        plt.xlabel('Predicted')
+        plt.ylabel('True')
+        plt.savefig(os.path.join(save_dir, f'{model_name}_confusion_matrix.png'))
+        plt.close()
+
+
+def plot_roc_curves(results, y_test, save_dir='../notebooks'):
+    """Plot ROC curves for all models using one-vs-rest approach."""
+    from sklearn.metrics import roc_curve, auc
+    from sklearn.preprocessing import label_binarize
+
+    n_classes = len(np.unique(y_test))
+    y_test_bin = label_binarize(y_test, classes=range(n_classes))
+
+    fig, axes = plt.subplots(2, 2, figsize=(15, 12))
+    axes = axes.ravel()
+
+    for i in range(n_classes):
+        for model_name, result in results.items():
+            if 'probabilities' not in result:
+                continue
+            y_prob = result['probabilities'][:, i]
+            fpr, tpr, _ = roc_curve(y_test_bin[:, i], y_prob)
+            roc_auc = auc(fpr, tpr)
+
+            axes[i].plot(fpr, tpr, label=f'{model_name} (AUC = {roc_auc:.2f})')
+            axes[i].plot([0, 1], [0, 1], 'k--')
+            axes[i].set_xlim([0.0, 1.0])
+            axes[i].set_ylim([0.0, 1.05])
+            axes[i].set_xlabel('False Positive Rate')
+            axes[i].set_ylabel('True Positive Rate')
+            axes[i].set_title(f'ROC Curve - Class {i}')
+            axes[i].legend(loc="lower right")
+
+    plt.tight_layout()
+    plt.savefig(os.path.join(save_dir, 'roc_curves.png'))
+    plt.close()
+
+
+def main():
+    print("=" * 80)
+    print("TRAINING TRADITIONAL ML MODELS WITH PATIENT-LEVEL SPLIT")
+    print("=" * 80)
+
+    # Initialize preprocessor
+    preprocessor = DataPreprocessor()
+
+    # Use all available datasets
+    base_dir = os.path.dirname(os.path.dirname(os.path.abspath(__file__)))
+    file_paths = [
+        os.path.join(base_dir, "PPMI_Curated_Data_Cut_Public_20241211.csv"),
+        os.path.join(base_dir, "PPMI_Curated_Data_Cut_Public_20250321.csv"),
+        os.path.join(base_dir, "PPMI_Curated_Data_Cut_Public_20250714.csv"),
+    ]
+
+    print("\n" + "=" * 80)
+    print("STEP 1: LOADING AND PREPARING DATA")
+    print("=" * 80)
+
+    X_train, X_test, y_train, y_test = preprocessor.prepare_data(
+        file_paths,
+        test_size=0.2,
+        use_patient_split=True,
+    )
+
+    print("\n" + "=" * 80)
+    print("STEP 2: CALCULATING CLASS WEIGHTS")
+    print("=" * 80)
+
+    from sklearn.utils.class_weight import compute_class_weight
+
+    class_weights = compute_class_weight('balanced', classes=np.unique(y_train), y=y_train)
+    class_weight_dict = dict(zip(np.unique(y_train), class_weights))
+    print("\nClass weights:", class_weight_dict)
+
+    print("\n" + "=" * 80)
+    print("STEP 3: INITIALIZING MODELS")
+    print("=" * 80)
+
+    ml_models = TraditionalMLModels()
+
+    ml_models.models['lightgbm'] = LGBMClassifier(
+        random_state=42,
+        class_weight=class_weight_dict,
+        objective='multiclass',
+        num_class=len(class_weight_dict),
+        n_estimators=200,
+        learning_rate=0.05,
+        max_depth=7,
+    )
+    ml_models.models['xgboost'] = XGBClassifier(
+        random_state=42,
+        objective='multi:softmax',
+        num_class=len(class_weight_dict),
+        n_estimators=200,
+        learning_rate=0.05,
+        max_depth=6,
+    )
+    ml_models.models['svm'] = SVC(
+        random_state=42,
+        class_weight=class_weight_dict,
+        probability=True,
+        decision_function_shape='ovr',
+        kernel='rbf',
+        C=10.0,
+        gamma='scale',
+    )
+
+    print("\n" + "=" * 80)
+    print("STEP 4: TRAINING MODELS")
+    print("=" * 80)
+    print("\nTraining models...")
+    cv_scores = ml_models.train_all_models(X_train, y_train)
+
+    print("\n" + "=" * 80)
+    print("STEP 5: TRAINING SCORES")
+    print("=" * 80)
+    print("\nTraining scores:")
+    for model_name, score in cv_scores.items():
+        print(f"  {model_name}: {score:.4f}")
+
+    print("\n" + "=" * 80)
+    print("STEP 6: EVALUATING ON TEST SET")
+    print("=" * 80)
+    print("\nEvaluating models on test set...")
+    results = ml_models.evaluate_all_models(X_test, y_test)
+
+    print("\n" + "=" * 80)
+    print("STEP 7: FINAL RESULTS")
+    print("=" * 80)
+    for model_name, result in results.items():
+        print(f"\n{model_name.upper()} Results:")
+        print(f"  Accuracy: {result['accuracy']:.4f}")
+        print(f"\nClassification Report:")
+        print(result['classification_report'])
+
+    print("\n" + "=" * 80)
+    print("STEP 8: GENERATING VISUALIZATIONS")
+    print("=" * 80)
+    plot_confusion_matrices(results)
+    plot_roc_curves(results, y_test)
+    print("\nVisualizations saved to notebooks/")
+
+    print("\n" + "=" * 80)
+    print("STEP 9: SAVING MODELS")
+    print("=" * 80)
+    print("\nSaving models...")
+    ml_models.save_models()
+
+    print("\n" + "=" * 80)
+    print("TRAINING COMPLETE!")
+    print("=" * 80)
+    print(f"\n[OK] Models trained with patient-level split (no leakage)")
+    print(f"[OK] {len(results)} models saved to models/saved/")
+    print(f"[OK] Visualizations saved to notebooks/")
+    print(f"[OK] Training data: {X_train.shape[0]} samples")
+    print(f"[OK] Test data: {X_test.shape[0]} samples")
+    print(f"[OK] Features: {X_train.shape[1]}")
+    print("\n" + "=" * 80)
+
+
+if __name__ == "__main__":
+    main()

src/train_transformer_models.py

@@ -0,0 +1,863 @@
+"""
+Training script for transformer-based models (BioGPT, Clinical-T5, PubMedBERT)
+on the PPMI dataset with RAG integration, CUDA acceleration, and leak-free patient split.
+
+Usage:
+    cd src
+    python train_transformer_models.py
+"""
+
+import sys
+import argparse
+import os
+from dataclasses import dataclass
+from pathlib import Path
+sys.path.append(os.path.dirname(os.path.abspath(__file__)))
+
+import torch
+import numpy as np
+from torch.utils.data import DataLoader, Subset
+from sklearn.utils.class_weight import compute_class_weight
+from sklearn.metrics import (
+    classification_report, confusion_matrix,
+    f1_score, precision_score, recall_score, roc_auc_score
+)
+from sklearn.model_selection import StratifiedShuffleSplit
+import matplotlib.pyplot as plt
+import seaborn as sns
+import pandas as pd
+import joblib
+import time
+import json
+from tqdm import tqdm
+
+from data_preprocessing import DataPreprocessor
+from models.transformer_models import TabularDataset
+from models.medical_transformers import (
+    BioMistralClassifier as BioGPTForTabular,
+    ClinicalT5Classifier as ClinicalT5ForTabular,
+    PubMedBERTClassifier as PubMedBERTForTabular,
+)
+from document_manager import DocumentManager
+
+
+# ---------------------------------------------------------------------------
+# Paths
+# ---------------------------------------------------------------------------
+ROOT = Path(__file__).resolve().parents[1]
+LEAK_FREE_SPLIT_PATH = ROOT / "evaluation_results" / "leak_free_split.npz"
+LEAK_FREE_META_PATH = ROOT / "evaluation_results" / "leak_free_split_meta.joblib"
+MODEL_DIR = ROOT / "models" / "saved"
+RESULTS_DIR = ROOT / "evaluation_results"
+PLOTS_DIR = ROOT / "evaluation_results" / "transformer_plots"
+
+# [CONFIG] Set to True if you want to use RAG (slower start), False for faster training
+USE_RAG = True
+REQUIRE_CUDA = os.getenv("PD_ALLOW_CPU_TRANSFORMERS", "0") != "1"
+
+
+
+# ---------------------------------------------------------------------------
+# Helpers
+# ---------------------------------------------------------------------------
+def _load_or_create_leak_free_split(preprocessor, file_paths):
+    """Load the cached leak-free split or regenerate it if missing."""
+    if LEAK_FREE_SPLIT_PATH.exists() and LEAK_FREE_META_PATH.exists():
+        split = np.load(LEAK_FREE_SPLIT_PATH)
+        meta = joblib.load(LEAK_FREE_META_PATH)
+        feature_names = meta.get("feature_names") if isinstance(meta, dict) else None
+        class_mapping = meta.get("class_mapping") if isinstance(meta, dict) else None
+        print("[DATA] Loaded cached leak-free split from evaluation_results.")
+        return (
+            split["X_train"], split["X_test"],
+            split["y_train"], split["y_test"],
+            feature_names, class_mapping,
+        )
+
+    print("[DATA] Leak-free split not found – regenerating via DataPreprocessor ...")
+    X_train, X_test, y_train, y_test = preprocessor.prepare_data(
+        file_paths, test_size=0.2, use_patient_split=True,
+    )
+    feature_names = preprocessor.get_feature_names()
+    class_mapping = preprocessor.get_class_mapping()
+
+    LEAK_FREE_SPLIT_PATH.parent.mkdir(parents=True, exist_ok=True)
+    np.savez(
+        LEAK_FREE_SPLIT_PATH,
+        X_train=X_train, X_test=X_test,
+        y_train=y_train, y_test=y_test,
+    )
+    joblib.dump(
+        {"feature_names": feature_names, "class_mapping": class_mapping},
+        LEAK_FREE_META_PATH,
+    )
+    print(f"[DATA] Saved fresh leak-free split → {LEAK_FREE_SPLIT_PATH}")
+    return X_train, X_test, y_train, y_test, feature_names, class_mapping
+
+
+def _stratified_val_indices(y_train, val_fraction=0.15, seed=42):
+    """Return (train_idx, val_idx) using a stratified split so every class
+    is proportionally represented in the validation set."""
+    sss = StratifiedShuffleSplit(n_splits=1, test_size=val_fraction, random_state=seed)
+    train_idx, val_idx = next(sss.split(np.zeros(len(y_train)), y_train))
+    return train_idx, val_idx
+
+
+def _prepare_batch(batch, device):
+    """Move tensors to the target device and keep optional RAG contexts aligned."""
+    if len(batch) == 3:
+        data, targets, contexts = batch
+    else:
+        data, targets = batch
+        contexts = None
+    data = data.to(device, non_blocking=True)
+    targets = targets.to(device, non_blocking=True)
+    if contexts is not None:
+        contexts = list(contexts)
+    return data, targets, contexts
+
+
+def _build_context_cache(features, build_fn, split_name="dataset", cache_path=None):
+    """Pre-compute RAG contexts with parallel processing and caching."""
+    if not USE_RAG:
+        return [""] * len(features)
+
+    if cache_path and os.path.exists(cache_path):
+        print(f"  [RAG] Loading cached contexts from {cache_path}")
+        return joblib.load(cache_path)
+
+    print(f"  [RAG] Generating contexts for {len(features)} samples (Parallel)...")
+    from joblib import Parallel, delayed
+    
+    # Run in parallel to speed up regex/cosine-sim
+    contexts = Parallel(n_jobs=-1, verbose=5)(
+        delayed(build_fn)(row) for row in features
+    )
+    
+    if cache_path:
+        joblib.dump(contexts, cache_path)
+        print(f"  [RAG] Saved contexts to {cache_path}")
+        
+    return contexts
+
+
+
+def _print_gpu_info(device):
+    """Print GPU diagnostics."""
+    if device.type != "cuda":
+        return
+    print(f"  GPU Name       : {torch.cuda.get_device_name(0)}")
+    print(f"  CUDA Version   : {torch.version.cuda}")
+    cap = torch.cuda.get_device_capability(0)
+    print(f"  Compute Cap.   : {cap[0]}.{cap[1]}")
+    mem_total = torch.cuda.get_device_properties(0).total_memory / 1024**3
+    print(f"  Total VRAM     : {mem_total:.1f} GB")
+    print(f"  cuDNN Enabled  : {torch.backends.cudnn.enabled}")
+    print(f"  cuDNN Benchmark: {torch.backends.cudnn.benchmark}")
+
+
+def _ensure_transformer_cuda() -> None:
+    if not torch.cuda.is_available():
+        if REQUIRE_CUDA:
+            raise RuntimeError(
+                "CUDA is required for transformer training in this accuracy-oriented configuration. "
+                "Install a CUDA-enabled PyTorch build and use a GPU, or set PD_ALLOW_CPU_TRANSFORMERS=1 to explicitly allow CPU fallback."
+            )
+        return
+    torch.set_float32_matmul_precision("highest")
+    if hasattr(torch.backends, "cuda") and hasattr(torch.backends.cuda, "matmul"):
+        torch.backends.cuda.matmul.allow_tf32 = False
+    if hasattr(torch.backends, "cudnn"):
+        torch.backends.cudnn.allow_tf32 = False
+        torch.backends.cudnn.benchmark = True
+
+
+@dataclass(frozen=True)
+class GPUExecutionProfile:
+    name: str
+    train_batch_by_model: dict
+    eval_batch_by_model: dict
+    grad_accum_by_model: dict
+    num_workers: int
+    prefetch_factor: int
+    persistent_workers: bool
+    notes: str
+
+
+def _detect_gpu_execution_profile():
+    if not torch.cuda.is_available():
+        return GPUExecutionProfile(
+            name="cpu",
+            train_batch_by_model={"pubmedbert": 4, "biogpt": 2, "clinical_t5": 2},
+            eval_batch_by_model={"pubmedbert": 8, "biogpt": 4, "clinical_t5": 4},
+            grad_accum_by_model={"pubmedbert": 12, "biogpt": 16, "clinical_t5": 16},
+            num_workers=0,
+            prefetch_factor=2,
+            persistent_workers=False,
+            notes="CPU fallback profile.",
+        )
+
+    gpu_name = torch.cuda.get_device_name(0).lower()
+    memory_gb = torch.cuda.get_device_properties(0).total_memory / 1024**3
+
+    if "a4000" in gpu_name or memory_gb >= 15.0:
+        return GPUExecutionProfile(
+            name="rtx-a4000",
+            train_batch_by_model={"pubmedbert": 16, "biogpt": 10, "clinical_t5": 10},
+            eval_batch_by_model={"pubmedbert": 48, "biogpt": 24, "clinical_t5": 24},
+            grad_accum_by_model={"pubmedbert": 4, "biogpt": 6, "clinical_t5": 6},
+            num_workers=4,
+            prefetch_factor=2,
+            persistent_workers=True,
+            notes="Optimized for RTX A4000 / ~16 GB VRAM.",
+        )
+
+    if memory_gb >= 11.0:
+        return GPUExecutionProfile(
+            name="high-vram",
+            train_batch_by_model={"pubmedbert": 12, "biogpt": 8, "clinical_t5": 8},
+            eval_batch_by_model={"pubmedbert": 32, "biogpt": 16, "clinical_t5": 16},
+            grad_accum_by_model={"pubmedbert": 6, "biogpt": 8, "clinical_t5": 8},
+            num_workers=2,
+            prefetch_factor=2,
+            persistent_workers=True,
+            notes="Generic 12 GB+ CUDA profile.",
+        )
+
+    return GPUExecutionProfile(
+        name="compat",
+        train_batch_by_model={"pubmedbert": 8, "biogpt": 6, "clinical_t5": 6},
+        eval_batch_by_model={"pubmedbert": 16, "biogpt": 8, "clinical_t5": 8},
+        grad_accum_by_model={"pubmedbert": 8, "biogpt": 10, "clinical_t5": 10},
+        num_workers=0,
+        prefetch_factor=2,
+        persistent_workers=False,
+        notes="Compatibility profile for lower-VRAM GPUs.",
+    )
+
+
+def _build_loader_kwargs(device, profile):
+    kwargs = {
+        "pin_memory": device.type == "cuda",
+        "num_workers": profile.num_workers if device.type == "cuda" else 0,
+    }
+    if kwargs["num_workers"] > 0:
+        kwargs["persistent_workers"] = profile.persistent_workers
+        kwargs["prefetch_factor"] = profile.prefetch_factor
+    return kwargs
+
+def _parse_selected_models(raw: str):
+    raw = (raw or 'all').strip().lower()
+    if raw in ('all', '*'):
+        return None
+    alias = {
+        'pubmed': 'pubmedbert',
+        'pubmedbert': 'pubmedbert',
+        'biogpt': 'biogpt',
+        'bio': 'biogpt',
+        'clinical': 'clinical_t5',
+        'clinical_t5': 'clinical_t5',
+        't5': 'clinical_t5',
+    }
+    out = []
+    for part in [p.strip() for p in raw.split(',') if p.strip()]:
+        out.append(alias.get(part, part))
+    return set(out) if out else None
+
+
+
+# ---------------------------------------------------------------------------
+# Training loop
+# ---------------------------------------------------------------------------
+def train_one_model(
+    model, optimizer, scheduler, criterion, scaler,
+    train_loader, val_loader, device, model_name,
+    num_epochs=25, patience=8, grad_accum_steps=2,
+    checkpoint_dir=None,
+):
+    """Train a single model with mixed precision, gradient accumulation, and
+    early stopping.  Returns the best model state dict and training history."""
+
+    checkpoint_path = checkpoint_dir / f"{model_name}_ckpt.pth" if checkpoint_dir else None
+    history = {"train_loss": [], "val_loss": [], "val_acc": [], "val_f1": [], "lr": []}
+    best_val_loss = float("inf")
+    early_stop_counter = 0
+    start_epoch = 0
+
+    # Resume from checkpoint if available
+    if checkpoint_path and checkpoint_path.exists():
+        print(f"  [CKPT] Found checkpoint at {checkpoint_path}")
+        try:
+            ckpt = torch.load(checkpoint_path, map_location=device, weights_only=False)
+            model.load_state_dict(ckpt["model_state_dict"])
+            optimizer.load_state_dict(ckpt["optimizer_state_dict"])
+            start_epoch = ckpt["epoch"] + 1
+            best_val_loss = ckpt["best_val_loss"]
+            history = ckpt.get("history", history)
+            print(f"  [CKPT] Resuming from epoch {start_epoch} (best val loss {best_val_loss:.4f})")
+        except Exception as e:
+            print(f"  [CKPT] Could not load checkpoint: {e}. Starting fresh.")
+            start_epoch = 0
+
+    use_amp = device.type == "cuda"
+
+    for epoch in range(start_epoch, num_epochs):
+        t0 = time.time()
+        model.train()
+        running_loss = 0.0
+        optimizer.zero_grad(set_to_none=True)
+
+        # Progress bar for training
+        pbar = tqdm(enumerate(train_loader), total=len(train_loader), 
+                    desc=f"Epoch {epoch+1}/{num_epochs}", unit="batch", ncols=100, ascii=True)
+        
+        for batch_idx, batch in pbar:
+            data, targets, contexts = _prepare_batch(batch, device)
+
+            with torch.amp.autocast(device_type=device.type, enabled=use_amp):
+                outputs = model(data, contexts)
+                loss = criterion(outputs, targets) / grad_accum_steps
+
+            scaler.scale(loss).backward()
+
+            if (batch_idx + 1) % grad_accum_steps == 0 or (batch_idx + 1) == len(train_loader):
+                scaler.unscale_(optimizer)
+                torch.nn.utils.clip_grad_norm_(model.parameters(), max_norm=1.0)
+                scaler.step(optimizer)
+                scaler.update()
+                optimizer.zero_grad(set_to_none=True)
+
+            current_loss = loss.item() * grad_accum_steps
+            running_loss += current_loss
+            
+            # Update progress bar every few batches to reduce overhead
+            if batch_idx % 10 == 0:
+                pbar.set_postfix(loss=f"{current_loss:.4f}", lr=f"{optimizer.param_groups[0]['lr']:.2e}")
+
+        avg_train_loss = running_loss / len(train_loader)
+        history["train_loss"].append(avg_train_loss)
+        history["lr"].append(optimizer.param_groups[0]["lr"])
+
+        # ---- Validation ----
+        model.eval()
+        val_loss = 0.0
+        all_preds, all_targets = [], []
+
+        with torch.no_grad():
+            for batch in tqdm(val_loader, desc=f"Val  {epoch+1}/{num_epochs}", unit="batch", leave=False, ncols=100, ascii=True):
+                data, targets, contexts = _prepare_batch(batch, device)
+                with torch.amp.autocast(device_type=device.type, enabled=use_amp):
+                    outputs = model(data, contexts)
+                    loss = criterion(outputs, targets)
+                val_loss += loss.item()
+                _, predicted = outputs.max(1)
+                all_preds.extend(predicted.cpu().numpy())
+                all_targets.extend(targets.cpu().numpy())
+
+        avg_val_loss = val_loss / len(val_loader)
+        val_acc = 100 * np.mean(np.array(all_preds) == np.array(all_targets))
+        val_f1 = f1_score(all_targets, all_preds, average="weighted")
+
+        history["val_loss"].append(avg_val_loss)
+        history["val_acc"].append(val_acc)
+        history["val_f1"].append(val_f1)
+        scheduler.step(avg_val_loss)
+
+        elapsed = time.time() - t0
+        gpu_mem = torch.cuda.memory_allocated(0) / 1024**2 if device.type == "cuda" else 0
+
+        print(
+            f"  Epoch {epoch+1:02d}/{num_epochs} │ "
+            f"Train Loss {avg_train_loss:.4f} │ Val Loss {avg_val_loss:.4f} │ "
+            f"Val Acc {val_acc:.2f}% │ Val F1 {val_f1:.4f} │ "
+            f"LR {optimizer.param_groups[0]['lr']:.2e} │ "
+            f"GPU {gpu_mem:.0f}MB │ {elapsed:.1f}s"
+        )
+
+        # ---- Checkpointing & early stopping ----
+        # Save every epoch so interrupted runs can resume from latest epoch.
+        if checkpoint_path:
+            torch.save({
+                "epoch": epoch,
+                "model_state_dict": model.state_dict(),
+                "optimizer_state_dict": optimizer.state_dict(),
+                "best_val_loss": best_val_loss,
+                "history": history,
+            }, checkpoint_path)
+
+        if avg_val_loss < best_val_loss:
+            best_val_loss = avg_val_loss
+            if checkpoint_path:
+                torch.save({
+                    "epoch": epoch,
+                    "model_state_dict": model.state_dict(),
+                    "optimizer_state_dict": optimizer.state_dict(),
+                    "best_val_loss": best_val_loss,
+                    "history": history,
+                }, checkpoint_path)
+            early_stop_counter = 0
+            print(f"  [*] New best (val loss {best_val_loss:.4f})")
+        else:
+            early_stop_counter += 1
+            if early_stop_counter >= patience:
+                print(f"  [X] Early stopping after {epoch+1} epochs")
+                break
+
+    # Load best weights
+    if checkpoint_path and checkpoint_path.exists():
+        ckpt = torch.load(checkpoint_path, map_location=device, weights_only=False)
+        model.load_state_dict(ckpt["model_state_dict"])
+        checkpoint_path.unlink()
+
+    return model, history, best_val_loss
+
+
+# ---------------------------------------------------------------------------
+# Evaluation
+# ---------------------------------------------------------------------------
+def evaluate_on_test(model, test_loader, criterion, device, model_name, class_names):
+    """Full evaluation on the held-out test set."""
+    model.eval()
+    use_amp = device.type == "cuda"
+    all_preds, all_targets, all_probs = [], [], []
+    test_loss = 0.0
+
+    with torch.no_grad():
+        for batch in tqdm(test_loader, desc=f"Evaluating {model_name}", unit="batch"):
+            data, targets, contexts = _prepare_batch(batch, device)
+            with torch.amp.autocast(device_type=device.type, enabled=use_amp):
+                outputs = model(data, contexts)
+                loss = criterion(outputs, targets)
+            test_loss += loss.item()
+            probs = torch.softmax(outputs, dim=1)
+            _, predicted = outputs.max(1)
+            all_preds.extend(predicted.cpu().numpy())
+            all_targets.extend(targets.cpu().numpy())
+            all_probs.extend(probs.cpu().numpy())
+
+    all_preds = np.array(all_preds)
+    all_targets = np.array(all_targets)
+    all_probs = np.array(all_probs)
+
+    accuracy = np.mean(all_preds == all_targets)
+    f1 = f1_score(all_targets, all_preds, average="weighted")
+    precision = precision_score(all_targets, all_preds, average="weighted")
+    recall = recall_score(all_targets, all_preds, average="weighted")
+
+    try:
+        auroc = roc_auc_score(all_targets, all_probs, multi_class="ovr", average="weighted")
+    except Exception:
+        auroc = 0.0
+
+    report = classification_report(all_targets, all_preds, target_names=class_names)
+    cm = confusion_matrix(all_targets, all_preds)
+
+    print(f"\n{'='*70}")
+    print(f"  {model_name.upper()} — TEST SET RESULTS")
+    print(f"{'='*70}")
+    print(f"  Accuracy  : {accuracy:.4f}  ({accuracy*100:.2f}%)")
+    print(f"  F1 Score  : {f1:.4f}")
+    print(f"  Precision : {precision:.4f}")
+    print(f"  Recall    : {recall:.4f}")
+    print(f"  AUROC     : {auroc:.4f}")
+    print(f"\n{report}")
+
+    return {
+        "accuracy": accuracy, "f1": f1, "precision": precision,
+        "recall": recall, "auroc": auroc,
+        "classification_report": report, "confusion_matrix": cm,
+        "predictions": all_preds, "targets": all_targets, "probabilities": all_probs,
+    }
+
+
+# ---------------------------------------------------------------------------
+# Plotting
+# ---------------------------------------------------------------------------
+def save_plots(results, history_dict, class_names, plots_dir):
+    """Save confusion matrices, training curves, and comparison charts."""
+    plots_dir = Path(plots_dir)
+    plots_dir.mkdir(parents=True, exist_ok=True)
+
+    for name, res in results.items():
+        # Confusion matrix
+        plt.figure(figsize=(8, 6))
+        sns.heatmap(res["confusion_matrix"], annot=True, fmt="d", cmap="Blues",
+                    xticklabels=class_names, yticklabels=class_names)
+        plt.title(f"{name} — Confusion Matrix (Leak-Free Split)")
+        plt.xlabel("Predicted")
+        plt.ylabel("True")
+        plt.tight_layout()
+        plt.savefig(plots_dir / f"{name}_confusion_matrix.png", dpi=200)
+        plt.close()
+
+        # Training curves
+        hist = history_dict[name]
+        fig, axes = plt.subplots(1, 3, figsize=(18, 5))
+
+        axes[0].plot(hist["train_loss"], label="Train", color="#3498db")
+        axes[0].plot(hist["val_loss"], label="Val", color="#e74c3c")
+        axes[0].set_title(f"{name} — Loss")
+        axes[0].set_xlabel("Epoch")
+        axes[0].legend()
+        axes[0].grid(True, alpha=0.3)
+
+        axes[1].plot(hist["val_acc"], color="#2ecc71")
+        axes[1].set_title(f"{name} — Val Accuracy (%)")
+        axes[1].set_xlabel("Epoch")
+        axes[1].grid(True, alpha=0.3)
+
+        axes[2].plot(hist["lr"], color="#9b59b6")
+        axes[2].set_title(f"{name} — Learning Rate")
+        axes[2].set_xlabel("Epoch")
+        axes[2].grid(True, alpha=0.3)
+
+        plt.tight_layout()
+        plt.savefig(plots_dir / f"{name}_training_curves.png", dpi=200)
+        plt.close()
+
+    # Comparison bar chart
+    model_names = list(results.keys())
+    metric_names = ["accuracy", "f1", "precision", "recall", "auroc"]
+    fig, axes = plt.subplots(1, len(metric_names), figsize=(5 * len(metric_names), 5))
+    colors = ["#3498db", "#2ecc71", "#e74c3c"]
+
+    for ax, metric in zip(axes, metric_names):
+        values = [results[m][metric] for m in model_names]
+        bars = ax.bar(model_names, values, color=colors[:len(model_names)])
+        ax.set_title(metric.upper())
+        ax.set_ylim(0, 1.05)
+        for bar, val in zip(bars, values):
+            ax.text(bar.get_x() + bar.get_width() / 2, bar.get_height() + 0.01,
+                    f"{val:.3f}", ha="center", va="bottom", fontsize=9)
+
+    plt.tight_layout()
+    plt.savefig(plots_dir / "transformer_comparison.png", dpi=200)
+    plt.close()
+
+    print(f"[PLOT] Saved all plots → {plots_dir}")
+
+
+# ---------------------------------------------------------------------------
+# Main
+# ---------------------------------------------------------------------------
+def main():
+    print("=" * 70)
+    print("  TRANSFORMER MODEL TRAINING — LEAK-FREE SPLIT + CUDA")
+    print("=" * 70)
+
+    # ---- Seed everything ----
+    torch.manual_seed(42)
+    np.random.seed(42)
+    _ensure_transformer_cuda()
+    if torch.cuda.is_available():
+        torch.cuda.manual_seed_all(42)
+
+    # ---- Device ----
+    device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
+    print(f"\n[DEVICE] Using: {device}")
+    if device.type == "cuda":
+        torch.backends.cudnn.enabled = True
+        torch.backends.cudnn.benchmark = True
+        torch.cuda.empty_cache()
+        _print_gpu_info(device)
+    else:
+        print("[WARNING] CUDA not available -- CPU fallback is enabled by PD_ALLOW_CPU_TRANSFORMERS=1")
+
+    # ---- Data ----
+    preprocessor = DataPreprocessor()
+    base_dir = str(ROOT)
+    file_paths = [
+        os.path.join(base_dir, "PPMI_Curated_Data_Cut_Public_20240129.csv"),
+        os.path.join(base_dir, "PPMI_Curated_Data_Cut_Public_20241211.csv"),
+        os.path.join(base_dir, "PPMI_Curated_Data_Cut_Public_20250321.csv"),
+        os.path.join(base_dir, "PPMI_Curated_Data_Cut_Public_20250714.csv"),
+    ]
+
+    X_train, X_test, y_train, y_test, feature_names, class_mapping = \
+        _load_or_create_leak_free_split(preprocessor, file_paths)
+
+    X_train = np.asarray(X_train, dtype=np.float32)
+    X_test = np.asarray(X_test, dtype=np.float32)
+    y_train = np.asarray(y_train, dtype=np.int64)
+    y_test = np.asarray(y_test, dtype=np.int64)
+
+    print(f"[DATA] Train: {X_train.shape}  Test: {X_test.shape}")
+    print(f"[DATA] Classes: {len(np.unique(y_train))}  Distribution: {dict(zip(*np.unique(y_train, return_counts=True)))}")
+
+    # ---- Class weights ----
+    cw = compute_class_weight("balanced", classes=np.unique(y_train), y=y_train)
+    class_weights_tensor = torch.FloatTensor(cw).to(device)
+    print(f"[DATA] Class weights: {dict(zip(np.unique(y_train), np.round(cw, 3)))}")
+
+    # ---- Stratified validation split (preserves class ratios) ----
+    train_idx, val_idx = _stratified_val_indices(y_train, val_fraction=0.15)
+    print(f"[DATA] Stratified split: {len(train_idx)} train / {len(val_idx)} val")
+
+    # ---- Feature names ----
+    if feature_names is None:
+        feature_names = preprocessor.get_feature_names()
+
+    # ---- RAG context ----
+    docs_path = str(ROOT / "medical_docs")
+    doc_manager = DocumentManager(docs_dir=docs_path)
+    doc_count = doc_manager.get_document_count()
+    print(f"[RAG] Loaded {doc_count.get('total', doc_count)} documents for context enrichment")
+
+    def get_rag_context(sample_features):
+        feature_desc = {name: float(val) for name, val in zip(feature_names, sample_features)}
+        query_parts = []
+        symptoms = {
+            "tremor": feature_desc.get("sym_tremor", 0),
+            "rigidity": feature_desc.get("sym_rigid", 0),
+            "bradykinesia": feature_desc.get("sym_brady", 0),
+            "postural instability": feature_desc.get("sym_posins", 0),
+        }
+        for symptom, severity in symptoms.items():
+            if severity > 0:
+                query_parts.append(f"{symptom} severity:{severity}")
+        moca = feature_desc.get("moca", 30)
+        if moca < 26:
+            query_parts.append("cognitive impairment")
+        age = feature_desc.get("age", 0)
+        if age:
+            query_parts.append(f"age {int(age)}")
+        if feature_desc.get("fampd", 0) > 0:
+            query_parts.append("family history Parkinson's disease")
+        query = "Parkinson's disease " + " ".join(query_parts)
+        passages = doc_manager.extract_relevant_passages(query, top_k=2)
+        if not passages:
+            return ""
+        return " ".join(
+            f"From '{p['doc_title']}' {p['text'][:300]}..." for p in passages
+        )
+
+    print(f"\n[RAG] RAG Enabled: {USE_RAG}")
+    if USE_RAG:
+        print("[RAG] Pre-computing context for train + test splits ...")
+    
+    train_cache = RESULTS_DIR / "rag_contexts_train.pkl"
+    test_cache = RESULTS_DIR / "rag_contexts_test.pkl"
+    
+    train_contexts = _build_context_cache(X_train, get_rag_context, "train", train_cache)
+    test_contexts = _build_context_cache(X_test, get_rag_context, "test", test_cache)
+
+
+    # ---- Datasets ----
+    full_train_ds = TabularDataset(X_train, y_train, feature_names, contexts=train_contexts)
+    test_ds = TabularDataset(X_test, y_test, feature_names, contexts=test_contexts)
+
+    train_subset = Subset(full_train_ds, train_idx)
+    val_subset = Subset(full_train_ds, val_idx)
+
+    gpu_profile = _detect_gpu_execution_profile()
+    loader_kwargs = _build_loader_kwargs(device, gpu_profile)
+    print(f"\n[GPU PROFILE] {gpu_profile.name} -> {gpu_profile.notes}")
+
+    # ---- Model definitions (lazy — created one at a time to fit in available VRAM) ----
+    input_dim = X_train.shape[1]
+    num_classes = len(np.unique(y_train))
+    class_names = ["HC", "PD", "SWEDD", "PRODROMAL"]
+
+    print(f"\n[MODEL] Input dim: {input_dim}  Num classes: {num_classes}")
+
+    # Each entry: (display_name, save_name, model_factory)
+    # Models are created lazily inside the loop to avoid GPU OOM.
+    selected_models = _parse_selected_models(os.getenv("PD_TRAIN_MODELS", "all"))
+
+    model_configs = [
+        (
+            "PubMedBERT (Encoder-Only)", "pubmedbert",
+            lambda: PubMedBERTForTabular(input_dim, num_classes, dropout=0.10, freeze_bert=False),
+            {"lr": 1.5e-5, "weight_decay": 0.02},
+        ),
+        (
+            "BioGPT", "biogpt",
+            lambda: BioGPTForTabular(input_dim, num_classes, dropout=0.12, train_decoder_layers=8),
+            {"lr": 2e-5, "weight_decay": 0.02},
+        ),
+        (
+            "Clinical-T5", "clinical_t5",
+            lambda: ClinicalT5ForTabular(input_dim, num_classes, dropout=0.10, freeze_encoder=False),
+            {"lr": 1.5e-5, "weight_decay": 0.02},
+        ),
+    ]
+
+    if selected_models is not None:
+        model_configs = [cfg for cfg in model_configs if cfg[1] in selected_models]
+        print(f"[MODEL] Filter active -> {sorted(selected_models)}")
+    if not model_configs:
+        raise ValueError("No valid models selected for training.")
+
+    # ---- Training config ----
+    NUM_EPOCHS = 30
+    PATIENCE = 10
+    DEFAULT_GRAD_ACCUM = 8
+
+    criterion = torch.nn.CrossEntropyLoss(weight=class_weights_tensor)
+    checkpoint_dir = MODEL_DIR / "_checkpoints"
+    checkpoint_dir.mkdir(parents=True, exist_ok=True)
+
+    all_results = {}
+    all_histories = {}
+
+    for display_name, save_name, model_factory, opt_kwargs in model_configs:
+        final_path = MODEL_DIR / f"{display_name}_best.pth"
+        train_bs = gpu_profile.train_batch_by_model.get(save_name, 8)
+        eval_bs = gpu_profile.eval_batch_by_model.get(save_name, max(train_bs * 2, 8))
+        grad_accum = gpu_profile.grad_accum_by_model.get(save_name, DEFAULT_GRAD_ACCUM)
+
+        train_loader = DataLoader(train_subset, batch_size=train_bs, shuffle=True, **loader_kwargs)
+        val_loader = DataLoader(val_subset, batch_size=eval_bs, shuffle=False, **loader_kwargs)
+        test_loader = DataLoader(test_ds, batch_size=eval_bs, shuffle=False, **loader_kwargs)
+
+        print(f"\n[LOADER] {display_name}: train_bs={train_bs} eval_bs={eval_bs} grad_accum={grad_accum} workers={loader_kwargs.get('num_workers', 0)}")
+
+        # ---- Skip if already trained ----
+        if final_path.exists():
+            print(f"\n{'='*70}")
+            print(f"  SKIPPING: {display_name}  (already trained)")
+            print(f"  Loading saved weights from: {final_path}")
+            print(f"{'='*70}")
+            try:
+                # Create the model on CPU first, load weights, then move to GPU
+                print(f"\n[MODEL] Initializing {display_name} for evaluation ...")
+                model = model_factory()
+                model.load_state_dict(torch.load(final_path, map_location="cpu", weights_only=True))
+                model.to(device)
+
+                # Evaluate on test set
+                result = evaluate_on_test(
+                    model, test_loader, criterion, device, display_name, class_names,
+                )
+                all_results[display_name] = result
+                all_histories[display_name] = {"train_loss": [], "val_loss": [], "val_acc": [], "val_f1": [], "lr": []}
+
+                # Free GPU memory before next model
+                del model
+                if device.type == "cuda":
+                    torch.cuda.empty_cache()
+
+                continue
+            except RuntimeError as e:
+                print(f"  [WARN] Cannot load saved weights (architecture changed?): {e}")
+                print(f"  [WARN] Deleting stale checkpoint and retraining ...")
+                final_path.unlink(missing_ok=True)
+
+        # ---- Create model fresh for training ----
+        print(f"\n[MODEL] Initializing {display_name} for training ...")
+        if device.type == "cuda":
+            torch.cuda.empty_cache()
+
+        model = model_factory().to(device)
+        optimizer = torch.optim.AdamW(
+            filter(lambda p: p.requires_grad, model.parameters()),
+            **opt_kwargs,
+        )
+        scheduler = torch.optim.lr_scheduler.ReduceLROnPlateau(
+            optimizer, mode="min", factor=0.5, patience=3, min_lr=1e-7,
+        )
+
+        trainable_params = sum(p.numel() for p in model.parameters() if p.requires_grad)
+        total_params = sum(p.numel() for p in model.parameters())
+
+        print(f"\n{'='*70}")
+        print(f"  TRAINING: {display_name}")
+        print(f"  Trainable params: {trainable_params:,} / {total_params:,}")
+        print(f"  Epochs: {NUM_EPOCHS}  Patience: {PATIENCE}  Grad Accum: {grad_accum}")
+        print(f"{'='*70}")
+
+        if device.type == "cuda":
+            torch.cuda.reset_peak_memory_stats()
+            torch.cuda.empty_cache()
+
+        scaler = torch.amp.GradScaler(device=device.type, enabled=device.type == "cuda")
+
+        trained_model, history, best_val = train_one_model(
+            model, optimizer, scheduler, criterion, scaler,
+            train_loader, val_loader, device, save_name,
+            num_epochs=NUM_EPOCHS, patience=PATIENCE, grad_accum_steps=grad_accum,
+            checkpoint_dir=checkpoint_dir,
+        )
+
+        # Save final model
+        torch.save(trained_model.state_dict(), final_path)
+        alias_paths = {
+            "pubmedbert": [
+                MODEL_DIR / "pubmedbert_transformer.pth",
+                MODEL_DIR / "pubmedbert.pth",
+            ],
+            "biogpt": [
+                MODEL_DIR / "biogpt_transformer.pth",
+                MODEL_DIR / "biogpt.pth",
+                MODEL_DIR / "biomistral.pth",
+            ],
+            "clinical_t5": [
+                MODEL_DIR / "clinical_t5_transformer.pth",
+                MODEL_DIR / "clinicalt5_transformer.pth",
+                MODEL_DIR / "clinical_t5.pth",
+            ],
+        }
+        for alias_path in alias_paths.get(save_name, []):
+            torch.save(trained_model.state_dict(), alias_path)
+        print(f"  [SAVE] Model saved → {final_path}")
+
+        if device.type == "cuda":
+            peak = torch.cuda.max_memory_allocated(0) / 1024**3
+            print(f"  [GPU] Peak VRAM usage: {peak:.2f} GB")
+
+        # Evaluate on test set
+        result = evaluate_on_test(
+            trained_model, test_loader, criterion, device, display_name, class_names,
+        )
+        all_results[display_name] = result
+        all_histories[display_name] = history
+
+        # Free GPU memory before next model
+        del model, trained_model, optimizer, scheduler, scaler
+        if device.type == "cuda":
+            torch.cuda.empty_cache()
+
+    # NOTE: We intentionally keep the checkpoint dir so that interrupted
+    # training can resume from the last saved epoch on the next run.
+
+    # ---- Save metrics to JSON + CSV ----
+    RESULTS_DIR.mkdir(parents=True, exist_ok=True)
+    summary_rows = []
+    for name, res in all_results.items():
+        summary_rows.append({
+            "Model": name,
+            "Type": "Transformer",
+            "Accuracy": round(res["accuracy"], 4),
+            "F1_Score": round(res["f1"], 4),
+            "Precision": round(res["precision"], 4),
+            "Recall": round(res["recall"], 4),
+            "AUROC": round(res["auroc"], 4),
+        })
+    df = pd.DataFrame(summary_rows)
+    df.to_csv(RESULTS_DIR / "transformer_metrics_latest.csv", index=False)
+    with open(RESULTS_DIR / "transformer_metrics_latest.json", "w") as f:
+        json.dump(summary_rows, f, indent=2)
+    print(f"\n[SAVE] Metrics  → {RESULTS_DIR / 'transformer_metrics_latest.csv'}")
+
+    # ---- Plots ----
+    save_plots(all_results, all_histories, class_names, PLOTS_DIR)
+
+    # ---- Final summary ----
+    print(f"\n{'='*70}")
+    print("  FINAL COMPARISON")
+    print(f"{'='*70}")
+    print(f"{'Model':<30} {'Accuracy':>10} {'F1':>10} {'Precision':>10} {'Recall':>10} {'AUROC':>10}")
+    print("-" * 80)
+    for row in summary_rows:
+        print(f"{row['Model']:<30} {row['Accuracy']:>10.4f} {row['F1_Score']:>10.4f} "
+              f"{row['Precision']:>10.4f} {row['Recall']:>10.4f} {row['AUROC']:>10.4f}")
+
+    best_by_f1 = max(all_results, key=lambda k: all_results[k]["f1"])
+    print(f"\n  [*] Best model (by F1): {best_by_f1} -- F1 {all_results[best_by_f1]['f1']:.4f}")
+    print(f"\n{'='*70}")
+    print("  TRAINING COMPLETE!")
+    print(f"{'='*70}\n")
+
+
+if __name__ == "__main__":
+    main()

src/training_runtime.py

@@ -0,0 +1,215 @@
+from __future__ import annotations
+
+import json
+from dataclasses import dataclass
+from datetime import datetime, timezone
+from pathlib import Path
+from typing import Any, Dict, Iterable, Optional
+
+
+def _iso_now() -> str:
+    return datetime.now(timezone.utc).isoformat()
+
+
+def _atomic_write_json(path: Path, payload: Dict[str, Any]) -> None:
+    path.parent.mkdir(parents=True, exist_ok=True)
+    temp_path = path.with_suffix(path.suffix + ".tmp")
+    temp_path.write_text(json.dumps(payload, indent=2, sort_keys=True), encoding="utf-8")
+    temp_path.replace(path)
+
+
+def _load_json(path: Path, default: Optional[Dict[str, Any]] = None) -> Dict[str, Any]:
+    if not path.exists():
+        return dict(default or {})
+    with path.open("r", encoding="utf-8") as handle:
+        return json.load(handle)
+
+
+class PauseRequested(RuntimeError):
+    """Raised when a training run should pause after saving progress."""
+
+
+class StopRequested(RuntimeError):
+    """Raised when a training run should stop gracefully."""
+
+
+@dataclass(frozen=True)
+class TrainingPaths:
+    root: Path
+    run_name: str
+
+    @property
+    def run_dir(self) -> Path:
+        return self.root / self.run_name
+
+    @property
+    def state_path(self) -> Path:
+        return self.run_dir / "run_state.json"
+
+    @property
+    def pause_flag(self) -> Path:
+        return self.run_dir / "pause.flag"
+
+    @property
+    def stop_flag(self) -> Path:
+        return self.run_dir / "stop.flag"
+
+    @property
+    def checkpoints_dir(self) -> Path:
+        return self.run_dir / "checkpoints"
+
+    @property
+    def metrics_dir(self) -> Path:
+        return self.run_dir / "metrics"
+
+    @property
+    def logs_dir(self) -> Path:
+        return self.run_dir / "logs"
+
+
+class TrainingRunController:
+    """Shared manifest/checkpoint controller for resumable training runs."""
+
+    def __init__(self, base_dir: Path | str, run_name: str):
+        self.paths = TrainingPaths(Path(base_dir), run_name)
+        self.paths.run_dir.mkdir(parents=True, exist_ok=True)
+        self.paths.checkpoints_dir.mkdir(parents=True, exist_ok=True)
+        self.paths.metrics_dir.mkdir(parents=True, exist_ok=True)
+        self.paths.logs_dir.mkdir(parents=True, exist_ok=True)
+        self.state = _load_json(
+            self.paths.state_path,
+            {
+                "run_name": run_name,
+                "status": "created",
+                "created_at": _iso_now(),
+                "updated_at": _iso_now(),
+                "selected_models": [],
+                "config": {},
+                "current": {},
+                "models": {},
+                "events": [],
+            },
+        )
+        self._flush()
+
+    def _flush(self) -> None:
+        self.state["updated_at"] = _iso_now()
+        _atomic_write_json(self.paths.state_path, self.state)
+
+    def initialize(self, selected_models: Iterable[str], config: Dict[str, Any], resume: bool = False) -> None:
+        if not resume or not self.state.get("selected_models"):
+            self.state["selected_models"] = list(selected_models)
+            self.state["config"] = config
+        self._flush()
+
+    def get_status(self) -> str:
+        return str(self.state.get("status", "created"))
+
+    def mark_running(self, stage: str, model_name: Optional[str] = None, extra: Optional[Dict[str, Any]] = None) -> None:
+        self.state["status"] = "running"
+        current = {"stage": stage}
+        if model_name:
+            current["model_name"] = model_name
+        if extra:
+            current.update(extra)
+        self.state["current"] = current
+        self._append_event("running", current)
+        self._flush()
+
+    def mark_paused(self, reason: str = "pause requested") -> None:
+        self.state["status"] = "paused"
+        self._append_event("paused", {"reason": reason})
+        self._flush()
+
+    def mark_stopped(self, reason: str = "stop requested") -> None:
+        self.state["status"] = "stopped"
+        self._append_event("stopped", {"reason": reason})
+        self._flush()
+
+    def mark_completed(self) -> None:
+        self.state["status"] = "completed"
+        self.state["completed_at"] = _iso_now()
+        self._append_event("completed", {})
+        self.clear_pause()
+        self.clear_stop()
+        self._flush()
+
+    def mark_failed(self, error_message: str) -> None:
+        self.state["status"] = "failed"
+        self.state["last_error"] = error_message
+        self._append_event("failed", {"error": error_message})
+        self._flush()
+
+    def _append_event(self, event_type: str, payload: Dict[str, Any]) -> None:
+        events = self.state.setdefault("events", [])
+        events.append({"at": _iso_now(), "type": event_type, **payload})
+        if len(events) > 100:
+            del events[:-100]
+
+    def update_model_state(self, model_name: str, **fields: Any) -> None:
+        models = self.state.setdefault("models", {})
+        model_state = models.setdefault(model_name, {})
+        model_state.update(fields)
+        model_state["updated_at"] = _iso_now()
+        self._flush()
+
+    def append_trial_result(self, model_name: str, result: Dict[str, Any]) -> None:
+        models = self.state.setdefault("models", {})
+        model_state = models.setdefault(model_name, {})
+        trial_results = model_state.setdefault("trial_results", [])
+        trial_results.append(result)
+        model_state["updated_at"] = _iso_now()
+        self._flush()
+
+    def write_metrics_file(self, filename: str, payload: Dict[str, Any]) -> Path:
+        target = self.paths.metrics_dir / filename
+        _atomic_write_json(target, payload)
+        return target
+
+    def request_pause(self) -> None:
+        self.paths.pause_flag.parent.mkdir(parents=True, exist_ok=True)
+        self.paths.pause_flag.write_text("pause\n", encoding="utf-8")
+
+    def clear_pause(self) -> None:
+        if self.paths.pause_flag.exists():
+            self.paths.pause_flag.unlink()
+
+    def request_stop(self) -> None:
+        self.paths.stop_flag.parent.mkdir(parents=True, exist_ok=True)
+        self.paths.stop_flag.write_text("stop\n", encoding="utf-8")
+
+    def clear_stop(self) -> None:
+        if self.paths.stop_flag.exists():
+            self.paths.stop_flag.unlink()
+
+    def raise_if_requested(self) -> None:
+        if self.paths.stop_flag.exists():
+            raise StopRequested("Stop requested via stop.flag")
+        if self.paths.pause_flag.exists():
+            raise PauseRequested("Pause requested via pause.flag")
+
+    def load_checkpoint_state(self, model_name: str, default: Optional[Dict[str, Any]] = None) -> Dict[str, Any]:
+        path = self.paths.checkpoints_dir / f"{model_name}.json"
+        return _load_json(path, default)
+
+    def save_checkpoint_state(self, model_name: str, payload: Dict[str, Any]) -> Path:
+        path = self.paths.checkpoints_dir / f"{model_name}.json"
+        _atomic_write_json(path, payload)
+        return path
+
+    def status_summary(self) -> Dict[str, Any]:
+        return {
+            "run_name": self.state.get("run_name"),
+            "status": self.state.get("status"),
+            "selected_models": self.state.get("selected_models", []),
+            "current": self.state.get("current", {}),
+            "models": self.state.get("models", {}),
+            "paths": {
+                "run_dir": str(self.paths.run_dir),
+                "state_path": str(self.paths.state_path),
+                "pause_flag": str(self.paths.pause_flag),
+                "stop_flag": str(self.paths.stop_flag),
+                "checkpoints_dir": str(self.paths.checkpoints_dir),
+                "metrics_dir": str(self.paths.metrics_dir),
+            },
+        }

src/treatment_model.py

@@ -0,0 +1,144 @@
+"""
+Treatment Response Model — LEDD → UPDRS Dose-Response Curve.
+
+Extracts PPMI treatment data (patients on PD treatment with ON/OFF
+UPDRS3 scores), fits a log-decay relationship between LEDD dose and
+UPDRS3 improvement, and predicts treatment effect for new patients.
+"""
+
+from __future__ import annotations
+
+import logging
+import math
+from typing import Any, Dict, Optional
+
+import numpy as np
+
+logger = logging.getLogger(__name__)
+
+# Default coefficients for the dose-response curve when PPMI fit fails.
+# UPDRS_reduction = a * log(1 + LEDD) + b * disease_duration + c
+_DEFAULT_COEFS = {"a": 3.5, "b": -0.3, "c": -2.0}
+
+
+class TreatmentModel:
+    """Dose-response model: LEDD → UPDRS3 reduction."""
+
+    def __init__(self) -> None:
+        self.fitted = False
+        self.coefs: Dict[str, float] = dict(_DEFAULT_COEFS)
+        self.r_squared: Optional[float] = None
+
+    # ------------------------------------------------------------------
+    # Public API
+    # ------------------------------------------------------------------
+    def fit(self, csv_path: str) -> "TreatmentModel":
+        """Fit the dose-response model from PPMI CSV data."""
+        try:
+            import pandas as pd
+            df = pd.read_csv(csv_path, low_memory=False)
+            return self._fit_from_dataframe(df)
+        except Exception as exc:
+            logger.warning("TreatmentModel.fit failed: %s — using defaults", exc)
+            self.fitted = False
+            return self
+
+    def _fit_from_dataframe(self, df: "pd.DataFrame") -> "TreatmentModel":
+        import pandas as pd
+
+        required = {"PDTRTMNT", "LEDD", "updrs3_score", "updrs3_score_on"}
+        if not required.issubset(set(df.columns)):
+            logger.warning("Missing columns for treatment model; using defaults.")
+            return self
+
+        # Filter: patients on treatment with valid ON/OFF UPDRS3
+        sub = df[df["PDTRTMNT"] == 1].copy()
+        for col in ["LEDD", "updrs3_score", "updrs3_score_on"]:
+            sub[col] = pd.to_numeric(sub[col], errors="coerce")
+        sub = sub.dropna(subset=["LEDD", "updrs3_score", "updrs3_score_on"])
+        sub = sub[sub["LEDD"] > 0]
+
+        if len(sub) < 20:
+            logger.warning("Too few treatment records (%d); using defaults.", len(sub))
+            return self
+
+        # Compute UPDRS reduction (OFF - ON; positive = improvement)
+        sub["updrs_reduction"] = sub["updrs3_score"] - sub["updrs3_score_on"]
+
+        # Prepare features: log(1 + LEDD), duration
+        sub["log_ledd"] = np.log1p(sub["LEDD"].values)
+        if "duration_yrs" in sub.columns:
+            sub["dur"] = pd.to_numeric(sub["duration_yrs"], errors="coerce").fillna(0)
+        else:
+            sub["dur"] = 0.0
+
+        y = sub["updrs_reduction"].values
+        X = np.column_stack([
+            sub["log_ledd"].values,
+            sub["dur"].values,
+            np.ones(len(sub)),
+        ])
+
+        # Ordinary least squares
+        try:
+            coefs, residuals, _, _ = np.linalg.lstsq(X, y, rcond=None)
+        except np.linalg.LinAlgError:
+            logger.warning("TreatmentModel OLS failed; using defaults.")
+            return self
+
+        self.coefs = {
+            "a": round(float(coefs[0]), 4),
+            "b": round(float(coefs[1]), 4),
+            "c": round(float(coefs[2]), 4),
+        }
+
+        # R²
+        y_pred = X @ coefs
+        ss_res = np.sum((y - y_pred) ** 2)
+        ss_tot = np.sum((y - np.mean(y)) ** 2)
+        self.r_squared = round(1 - ss_res / max(ss_tot, 1e-10), 4)
+
+        self.fitted = True
+        logger.info(
+            "TreatmentModel fitted: coefs=%s, R²=%.4f, n=%d",
+            self.coefs, self.r_squared, len(sub),
+        )
+        return self
+
+    def predict_treatment_effect(
+        self,
+        ledd: Optional[float],
+        duration_years: float = 0.0,
+    ) -> float:
+        """
+        Predict UPDRS3 reduction for given LEDD and disease duration.
+
+        Returns a positive number meaning points of improvement.
+        """
+        if ledd is None or ledd <= 0:
+            return 0.0
+
+        a = self.coefs["a"]
+        b = self.coefs["b"]
+        c = self.coefs["c"]
+
+        reduction = a * math.log1p(ledd) + b * duration_years + c
+        # Clamp to non-negative (treatment can't make things worse in this model)
+        return round(max(0.0, reduction), 2)
+
+    def apply_treatment_effect(
+        self,
+        updrs3_off: float,
+        ledd: Optional[float],
+        duration_years: float = 0.0,
+    ) -> float:
+        """Return predicted ON-medication UPDRS3 score."""
+        reduction = self.predict_treatment_effect(ledd, duration_years)
+        return round(max(0.0, updrs3_off - reduction), 2)
+
+    def to_dict(self) -> Dict[str, Any]:
+        return {
+            "fitted": self.fitted,
+            "coefs": self.coefs,
+            "r_squared": self.r_squared,
+        }

src/twin_engine.py

@@ -0,0 +1,689 @@
+from __future__ import annotations
+
+import logging
+import math
+from copy import deepcopy
+from datetime import datetime, timezone
+from typing import Any, Dict, List, Optional
+from uuid import uuid4
+
+from twin_schema import (
+    DigitalTwin,
+    TwinForecastPoint,
+    TwinSimulation,
+    TwinSnapshot,
+    TwinState,
+    TwinStaticProfile,
+)
+from twin_store import TwinStore
+
+logger = logging.getLogger(__name__)
+
+# Lazy singleton for the predictor bridge
+_bridge_instance = None
+
+
+def _get_bridge():
+    global _bridge_instance
+    if _bridge_instance is None:
+        try:
+            from twin_predictor_bridge import TwinPredictorBridge
+            _bridge_instance = TwinPredictorBridge()
+            status = _bridge_instance.get_status()
+            logger.info("TwinPredictorBridge initialized: %s", status)
+            if status.get("silhouette_score") is not None:
+                print(f"[TWIN] Progression silhouette score: {status['silhouette_score']:.3f}")
+            if status.get("treatment_r_squared") is not None:
+                print(f"[TWIN] Treatment model R²: {status['treatment_r_squared']:.4f}")
+        except Exception as exc:
+            logger.warning("Failed to init TwinPredictorBridge: %s", exc)
+            _bridge_instance = None
+    return _bridge_instance
+
+
+CLASS_NAMES = ["Healthy Control", "Parkinson's Disease", "SWEDD", "Prodromal PD"]
+FORECAST_HORIZONS_MONTHS = [3, 6, 12]
+
+
+def _iso_now() -> str:
+    return datetime.now(timezone.utc).replace(microsecond=0).isoformat().replace("+00:00", "Z")
+
+
+def _safe_float(value: Any) -> Optional[float]:
+    if value in (None, ""):
+        return None
+    if isinstance(value, bool):
+        return float(int(value))
+    if isinstance(value, (int, float)):
+        if isinstance(value, float) and math.isnan(value):
+            return None
+        return float(value)
+    try:
+        stripped = str(value).strip()
+        if stripped == "":
+            return None
+        parsed = float(stripped)
+    except (TypeError, ValueError):
+        return None
+    if math.isnan(parsed):
+        return None
+    return parsed
+
+
+def _coerce_text(value: Any) -> Optional[str]:
+    if value is None:
+        return None
+    text = str(value).strip()
+    return text or None
+
+
+def _clamp(value: Optional[float], low: float, high: float) -> Optional[float]:
+    if value is None:
+        return None
+    return max(low, min(high, value))
+
+
+def _round_optional(value: Optional[float], digits: int = 2) -> Optional[float]:
+    if value is None:
+        return None
+    return round(value, digits)
+
+
+def _parse_date(date_string: Optional[str]) -> Optional[datetime]:
+    if not date_string:
+        return None
+    for fmt in ("%Y-%m-%d", "%Y-%m-%dT%H:%M:%S", "%Y-%m-%dT%H:%M:%SZ"):
+        try:
+            return datetime.strptime(date_string[: len(fmt)], fmt)
+        except ValueError:
+            continue
+    return None
+
+
+def _scale(value: Optional[float], maximum: float) -> Optional[float]:
+    if value is None or maximum <= 0:
+        return None
+    return _clamp(value / maximum, 0, 1)
+
+
+def _inverse_scale(value: Optional[float], maximum: float) -> Optional[float]:
+    if value is None or maximum <= 0:
+        return None
+    return _clamp((maximum - value) / maximum, 0, 1)
+
+
+def _mean_defined(values: List[Optional[float]]) -> float:
+    defined = [value for value in values if value is not None]
+    if not defined:
+        return 0.0
+    return sum(defined) / len(defined)
+
+
+class DigitalTwinEngine:
+    def __init__(self, store: Optional[TwinStore] = None, db_path: Optional[str] = None):
+        self.store = store or TwinStore(db_path=db_path)
+        self.bridge = _get_bridge()
+
+    def list_twins(self) -> List[Dict[str, Any]]:
+        return self.store.list_twins()
+
+    def get_twin(self, twin_id: str) -> Optional[Dict[str, Any]]:
+        return self.store.get_twin(twin_id)
+
+    def create_twin(
+        self,
+        patient_data: Dict[str, Any],
+        patient_label: Optional[str] = None,
+        source_patno: Optional[int] = None,
+        predictor: Optional[Any] = None,
+    ) -> Dict[str, Any]:
+        created_at = _iso_now()
+        twin_id = f"twin_{uuid4().hex[:12]}"
+        profile = self._build_profile(
+            twin_id=twin_id,
+            patient_data=patient_data,
+            patient_label=patient_label,
+            source_patno=source_patno,
+            created_at=created_at,
+        )
+        snapshot = self._build_snapshot(patient_data, snapshot_index=0)
+        prediction_summary = self._predict_current_state(patient_data, predictor)
+        bridge_result = self._bridge_predict(patient_data, [snapshot.to_dict()])
+        state = self._build_state(profile, [snapshot], prediction_summary, bridge_result)
+        forecast = self._build_forecast(snapshot, state, bridge_result)
+        twin = DigitalTwin(
+            profile=profile,
+            snapshots=[snapshot],
+            current_state=state,
+            forecast=forecast,
+            prediction_summary=prediction_summary,
+        )
+        return self.store.upsert_twin(twin)
+
+    def add_snapshot(
+        self,
+        twin_id: str,
+        patient_data: Dict[str, Any],
+        predictor: Optional[Any] = None,
+    ) -> Optional[Dict[str, Any]]:
+        stored = self.store.get_twin(twin_id)
+        if stored is None:
+            return None
+
+        snapshots = [self._snapshot_from_dict(item) for item in stored["snapshots"]]
+        profile = self._profile_from_dict(stored["profile"])
+        snapshots.append(self._build_snapshot(patient_data, snapshot_index=len(snapshots)))
+        prediction_summary = self._predict_current_state(patient_data, predictor)
+        snap_dicts = [s.to_dict() for s in snapshots]
+        bridge_result = self._bridge_predict(patient_data, snap_dicts)
+        state = self._build_state(profile, snapshots, prediction_summary, bridge_result)
+        forecast = self._build_forecast(snapshots[-1], state, bridge_result)
+        twin = DigitalTwin(
+            profile=profile,
+            snapshots=snapshots,
+            current_state=state,
+            forecast=forecast,
+            prediction_summary=prediction_summary,
+        )
+        return self.store.upsert_twin(twin)
+
+    def simulate(
+        self,
+        twin_id: str,
+        overrides: Dict[str, Any],
+        scenario_name: Optional[str] = None,
+        predictor: Optional[Any] = None,
+    ) -> Optional[Dict[str, Any]]:
+        stored = self.store.get_twin(twin_id)
+        if stored is None or not stored["snapshots"]:
+            return None
+
+        profile = self._profile_from_dict(stored["profile"])
+        history = [self._snapshot_from_dict(item) for item in stored["snapshots"]]
+        base_snapshot = self._snapshot_from_dict(stored["snapshots"][-1])
+        raw_inputs = deepcopy(base_snapshot.raw_inputs)
+        raw_inputs.update(overrides)
+        simulated_snapshot = self._build_snapshot(
+            raw_inputs,
+            snapshot_index=len(stored["snapshots"]),
+            default_event_id="SIM",
+        )
+        prediction_summary = self._predict_current_state(raw_inputs, predictor)
+        all_snaps = history + [simulated_snapshot]
+        snap_dicts = [s.to_dict() for s in all_snaps]
+        bridge_result = self._bridge_predict(raw_inputs, snap_dicts)
+        state = self._build_state(profile, all_snaps, prediction_summary, bridge_result)
+        forecast = self._build_forecast(simulated_snapshot, state, bridge_result)
+        simulation = TwinSimulation(
+            scenario_name=(scenario_name or "Scenario").strip() or "Scenario",
+            overrides=overrides,
+            simulated_snapshot=simulated_snapshot,
+            state=state,
+            forecast=forecast,
+        )
+        return simulation.to_dict()
+
+    def _build_profile(
+        self,
+        twin_id: str,
+        patient_data: Dict[str, Any],
+        patient_label: Optional[str],
+        source_patno: Optional[int],
+        created_at: str,
+    ) -> TwinStaticProfile:
+        resolved_label = str(patient_label or patient_data.get("patient_id") or twin_id)
+        return TwinStaticProfile(
+            twin_id=twin_id,
+            patient_label=resolved_label,
+            source_patno=source_patno,
+            created_at=created_at,
+            enrollment_cohort=str(patient_data.get("COHORT") or "Unknown"),
+            subgroup=_coerce_text(patient_data.get("subgroup")),
+            sex=_safe_float(patient_data.get("SEX")),
+            education_years=_safe_float(patient_data.get("EDUCYRS")),
+            race=_safe_float(patient_data.get("race")),
+            family_pd=_safe_float(patient_data.get("fampd")),
+            family_pd_bin=_safe_float(patient_data.get("fampd_bin")),
+            bmi=_safe_float(patient_data.get("BMI")),
+            age_diag=_safe_float(patient_data.get("agediag")),
+            age_onset=_safe_float(patient_data.get("ageonset")),
+            dominant_side=_safe_float(patient_data.get("DOMSIDE")),
+        )
+
+    def _build_snapshot(
+        self,
+        patient_data: Dict[str, Any],
+        snapshot_index: int,
+        default_event_id: str = "MANUAL",
+    ) -> TwinSnapshot:
+        visit_date = _coerce_text(patient_data.get("visit_date")) or datetime.now().strftime("%Y-%m-%d")
+        return TwinSnapshot(
+            snapshot_id=f"snap_{uuid4().hex[:12]}",
+            event_id=_coerce_text(patient_data.get("EVENT_ID")) or f"{default_event_id}_{snapshot_index + 1}",
+            visit_date=visit_date,
+            year_index=_safe_float(patient_data.get("YEAR")),
+            age_at_visit=_safe_float(patient_data.get("age_at_visit") or patient_data.get("age")),
+            duration_years=_safe_float(patient_data.get("duration_yrs")),
+            treatment_flag=_safe_float(patient_data.get("PDTRTMNT")),
+            ledd=_safe_float(patient_data.get("LEDD")),
+            motor={
+                "sym_tremor": _safe_float(patient_data.get("sym_tremor")),
+                "sym_rigid": _safe_float(patient_data.get("sym_rigid")),
+                "sym_brady": _safe_float(patient_data.get("sym_brady")),
+                "sym_posins": _safe_float(patient_data.get("sym_posins")),
+                "hy": _safe_float(patient_data.get("hy")),
+                "hy_on": _safe_float(patient_data.get("hy_on")),
+                "pigd": _safe_float(patient_data.get("pigd")),
+                "td_pigd": _safe_float(patient_data.get("td_pigd")),
+                "updrs1_score": _safe_float(patient_data.get("updrs1_score")),
+                "updrs2_score": _safe_float(patient_data.get("updrs2_score")),
+                "updrs3_score": _safe_float(patient_data.get("updrs3_score")),
+                "updrs3_score_on": _safe_float(patient_data.get("updrs3_score_on")),
+                "updrs4_score": _safe_float(patient_data.get("updrs4_score")),
+                "updrs_totscore": _safe_float(patient_data.get("updrs_totscore")),
+                "updrs_totscore_on": _safe_float(patient_data.get("updrs_totscore_on")),
+            },
+            cognition={
+                "moca": _safe_float(patient_data.get("moca")),
+                "bjlot": _safe_float(patient_data.get("bjlot")),
+                "clockdraw": _safe_float(patient_data.get("clockdraw")),
+                "hvlt_immediaterecall": _safe_float(patient_data.get("hvlt_immediaterecall")),
+                "hvlt_retention": _safe_float(patient_data.get("hvlt_retention")),
+                "hvlt_discrimination": _safe_float(patient_data.get("hvlt_discrimination")),
+                "lexical": _safe_float(patient_data.get("lexical")),
+                "lns": _safe_float(patient_data.get("lns")),
+            },
+            non_motor={
+                "ess": _safe_float(patient_data.get("ess")),
+                "rem": _safe_float(patient_data.get("rem")),
+                "gds": _safe_float(patient_data.get("gds")),
+                "stai": _safe_float(patient_data.get("stai")),
+                "quip_any": _safe_float(patient_data.get("quip_any")),
+                "NP1COG": _safe_float(patient_data.get("NP1COG")),
+                "NP1DPRS": _safe_float(patient_data.get("NP1DPRS")),
+                "NP1ANXS": _safe_float(patient_data.get("NP1ANXS")),
+                "NP1APAT": _safe_float(patient_data.get("NP1APAT")),
+                "NP1FATG": _safe_float(patient_data.get("NP1FATG")),
+            },
+            autonomic={
+                "scopa": _safe_float(patient_data.get("scopa")),
+                "orthostasis": _safe_float(patient_data.get("orthostasis")),
+            },
+            biomarkers={
+                "abeta": _safe_float(patient_data.get("abeta")),
+                "tau": _safe_float(patient_data.get("tau")),
+                "ptau": _safe_float(patient_data.get("ptau")),
+                "asyn": _safe_float(patient_data.get("asyn")),
+                "nfl_serum": _safe_float(patient_data.get("nfl_serum")),
+                "NFL_CSF": _safe_float(patient_data.get("NFL_CSF")),
+            },
+            imaging={
+                "MIA_CAUDATE_mean": _safe_float(patient_data.get("MIA_CAUDATE_mean")),
+                "MIA_PUTAMEN_mean": _safe_float(patient_data.get("MIA_PUTAMEN_mean")),
+                "MIA_STRIATUM_mean": _safe_float(patient_data.get("MIA_STRIATUM_mean")),
+            },
+            raw_inputs=deepcopy(patient_data),
+        )
+
+    def _predict_current_state(
+        self,
+        patient_data: Dict[str, Any],
+        predictor: Optional[Any],
+    ) -> Dict[str, Any]:
+        if predictor is not None:
+            required_fields = [
+                "age",
+                "SEX",
+                "EDUCYRS",
+                "BMI",
+                "sym_tremor",
+                "sym_rigid",
+                "sym_brady",
+                "sym_posins",
+            ]
+            if all(_safe_float(patient_data.get(field)) is not None for field in required_fields):
+                try:
+                    prediction = predictor.predict_patient(patient_data)
+                    class_index = int(prediction["ensemble_prediction"])
+                    return {
+                        "prediction": CLASS_NAMES[class_index],
+                        "confidence": round(float(prediction.get("confidence") or 0.0), 3),
+                        "probabilities": {
+                            CLASS_NAMES[idx]: round(float(prob), 4)
+                            for idx, prob in enumerate(prediction.get("ensemble_probabilities", []))
+                        },
+                        "source": "assessment_model",
+                    }
+                except Exception:
+                    pass
+
+        motor_score = sum(
+            value or 0.0
+            for value in (
+                _safe_float(patient_data.get("sym_tremor")),
+                _safe_float(patient_data.get("sym_rigid")),
+                _safe_float(patient_data.get("sym_brady")),
+                _safe_float(patient_data.get("sym_posins")),
+            )
+        )
+        rem = _safe_float(patient_data.get("rem")) or 0.0
+        moca = _safe_float(patient_data.get("moca"))
+
+        if motor_score <= 1.0 and rem == 0 and (moca is None or moca >= 27):
+            prediction = "Healthy Control"
+            confidence = 0.58
+        elif motor_score >= 5.0:
+            prediction = "Parkinson's Disease"
+            confidence = 0.62
+        elif rem == 1 or (moca is not None and moca < 25):
+            prediction = "Prodromal PD"
+            confidence = 0.56
+        else:
+            prediction = "SWEDD"
+            confidence = 0.52
+
+        return {
+            "prediction": prediction,
+            "confidence": confidence,
+            "probabilities": {prediction: confidence},
+            "source": "heuristic_fallback",
+        }
+
+    def _bridge_predict(
+        self,
+        patient_data: Dict[str, Any],
+        snapshots: List[Dict[str, Any]],
+    ) -> Dict[str, Any]:
+        """Run the TwinPredictorBridge (ML + fallback)."""
+        if self.bridge is None:
+            self.bridge = _get_bridge()
+        if self.bridge is not None:
+            try:
+                return self.bridge.predict(patient_data, snapshots)
+            except Exception as exc:
+                logger.warning("Bridge predict failed: %s", exc)
+        return {}
+
+    def _build_state(
+        self,
+        profile: TwinStaticProfile,
+        snapshots: List[TwinSnapshot],
+        prediction_summary: Dict[str, Any],
+        bridge_result: Optional[Dict[str, Any]] = None,
+    ) -> TwinState:
+        latest = snapshots[-1]
+        motor_index = self._motor_burden_index(latest)
+        cognitive_index = self._cognitive_burden_index(latest)
+        non_motor_index = self._non_motor_burden_index(latest)
+        progression_velocity = self._progression_velocity(snapshots)
+        treatment_response_proxy = self._treatment_response_proxy(latest)
+        br = bridge_result or {}
+        evidence = self._build_evidence(
+            profile=profile,
+            snapshot=latest,
+            motor_index=motor_index,
+            cognitive_index=cognitive_index,
+            non_motor_index=non_motor_index,
+            progression_velocity=progression_velocity,
+            prediction_summary=prediction_summary,
+            bridge_result=br,
+        )
+
+        return TwinState(
+            current_cohort_estimate=prediction_summary.get("prediction", "Unknown"),
+            prediction_source=prediction_summary.get("source", "heuristic"),
+            confidence=float(br.get("confidence") or prediction_summary.get("confidence") or 0.0),
+            motor_burden_index=_round_optional(motor_index),
+            cognitive_burden_index=_round_optional(cognitive_index),
+            non_motor_burden_index=_round_optional(non_motor_index),
+            progression_velocity=_round_optional(progression_velocity),
+            treatment_response_proxy=_round_optional(treatment_response_proxy),
+            computed_at=_iso_now(),
+            cluster_id=br.get("cluster_id"),
+            cluster_label=br.get("cluster_label"),
+            treatment_effect=_round_optional(br.get("treatment_effect")),
+            ci_lower=_round_optional(br.get("ci_lower")),
+            ci_upper=_round_optional(br.get("ci_upper")),
+            evidence=evidence,
+        )
+
+    def _build_forecast(
+        self,
+        snapshot: TwinSnapshot,
+        state: TwinState,
+        bridge_result: Optional[Dict[str, Any]] = None,
+    ) -> List[TwinForecastPoint]:
+        motor_index = state.motor_burden_index or 0.0
+        cognitive_index = state.cognitive_burden_index or 0.0
+        non_motor_index = state.non_motor_burden_index or 0.0
+        duration_years = snapshot.duration_years or 0.0
+
+        current_updrs3 = snapshot.motor.get("updrs3_score")
+        if current_updrs3 is None:
+            current_updrs3 = 8 + motor_index * 22
+
+        current_total = snapshot.motor.get("updrs_totscore")
+        if current_total is None:
+            current_total = current_updrs3 * 1.7 + 8
+
+        current_moca = snapshot.cognition.get("moca")
+        if current_moca is None:
+            current_moca = 30 - cognitive_index * 8
+
+        current_hy = snapshot.motor.get("hy")
+        if current_hy is None:
+            current_hy = 1 + motor_index * 2.2
+
+        # Use cluster-weighted profiles from the bridge if available
+        br = bridge_result or {}
+        profile = br.get("progression_profile", {})
+        cluster_label = br.get("cluster_label", "moderate")
+        treatment_effect = br.get("treatment_effect", 0.0) or 0.0
+
+        if profile:
+            yearly_updrs3_gain = profile.get("updrs3_gain", 3.5)
+            yearly_moca_loss = profile.get("moca_loss", 0.7)
+            yearly_hy_gain = profile.get("hy_gain", 0.25)
+        else:
+            velocity = state.progression_velocity or 0.0
+            yearly_updrs3_gain = 1.5 + motor_index * 3.5 + duration_years * 0.2 + velocity * 2.0
+            yearly_moca_loss = 0.4 + cognitive_index * 0.9 + velocity * 0.2
+            yearly_hy_gain = 0.15 + motor_index * 0.35 + velocity * 0.05
+
+        yearly_total_gain = yearly_updrs3_gain * 1.8 + non_motor_index * 1.2
+
+        # Apply treatment effect: reduce UPDRS gains
+        treatment_offset = min(treatment_effect, yearly_updrs3_gain * 0.8)
+
+        forecast: List[TwinForecastPoint] = []
+        for months in FORECAST_HORIZONS_MONTHS:
+            years = months / 12.0
+            accel = 1.0 + duration_years * 0.02
+            raw_updrs3 = current_updrs3 + (yearly_updrs3_gain * accel - treatment_offset) * years
+            predicted_updrs3 = _round_optional(max(0, raw_updrs3))
+            predicted_total = _round_optional(max(0, current_total + (yearly_total_gain * accel - treatment_offset * 1.5) * years))
+            predicted_moca = _round_optional(_clamp(current_moca - yearly_moca_loss * years, 0, 30))
+            predicted_hy = _round_optional(_clamp(current_hy + yearly_hy_gain * years, 0, 5))
+            risk_level = self._risk_level(predicted_updrs3, predicted_moca, state.current_cohort_estimate)
+            forecast.append(
+                TwinForecastPoint(
+                    horizon_months=months,
+                    predicted_updrs3=predicted_updrs3,
+                    predicted_total_updrs=predicted_total,
+                    predicted_moca=predicted_moca,
+                    predicted_hy=predicted_hy,
+                    risk_level=risk_level,
+                    uncertainty={
+                        "updrs3_pm": _round_optional(1.5 + months * 0.4),
+                        "total_updrs_pm": _round_optional(3.0 + months * 0.8),
+                        "moca_pm": _round_optional(0.4 + months * 0.08),
+                    },
+                )
+            )
+
+        return forecast
+
+    def _motor_burden_index(self, snapshot: TwinSnapshot) -> float:
+        components = [
+            _scale(snapshot.motor.get("sym_tremor"), 4),
+            _scale(snapshot.motor.get("sym_rigid"), 4),
+            _scale(snapshot.motor.get("sym_brady"), 4),
+            _scale(snapshot.motor.get("sym_posins"), 4),
+            _scale(snapshot.motor.get("updrs3_score"), 60),
+            _scale(snapshot.motor.get("updrs_totscore"), 120),
+            _scale(snapshot.motor.get("hy"), 5),
+        ]
+        return _mean_defined(components)
+
+    def _cognitive_burden_index(self, snapshot: TwinSnapshot) -> float:
+        components = [
+            _inverse_scale(snapshot.cognition.get("moca"), 30),
+            _inverse_scale(snapshot.cognition.get("bjlot"), 30),
+            _inverse_scale(snapshot.cognition.get("clockdraw"), 4),
+            _inverse_scale(snapshot.cognition.get("hvlt_immediaterecall"), 36),
+            _inverse_scale(snapshot.cognition.get("lns"), 21),
+        ]
+        return _mean_defined(components)
+
+    def _non_motor_burden_index(self, snapshot: TwinSnapshot) -> float:
+        components = [
+            _scale(snapshot.non_motor.get("ess"), 24),
+            _scale(snapshot.non_motor.get("gds"), 15),
+            _scale((_safe_float(snapshot.non_motor.get("stai")) or 20) - 20, 60),
+            _scale(snapshot.non_motor.get("rem"), 1),
+            _scale(snapshot.non_motor.get("quip_any"), 1),
+            _scale(snapshot.autonomic.get("scopa"), 39),
+            _scale(snapshot.autonomic.get("orthostasis"), 1),
+            _scale(snapshot.non_motor.get("NP1DPRS"), 4),
+            _scale(snapshot.non_motor.get("NP1ANXS"), 4),
+            _scale(snapshot.non_motor.get("NP1APAT"), 4),
+            _scale(snapshot.non_motor.get("NP1FATG"), 4),
+        ]
+        return _mean_defined(components)
+
+    def _progression_velocity(self, snapshots: List[TwinSnapshot]) -> Optional[float]:
+        if len(snapshots) < 2:
+            return None
+
+        first = snapshots[0]
+        last = snapshots[-1]
+        delta_years: Optional[float] = None
+
+        first_date = _parse_date(first.visit_date)
+        last_date = _parse_date(last.visit_date)
+        if first_date is not None and last_date is not None and last_date > first_date:
+            delta_years = (last_date - first_date).days / 365.25
+
+        # Fall back to YEAR index when visits share the same date.
+        if delta_years is None or delta_years <= 0:
+            first_year = _safe_float(first.year_index)
+            last_year = _safe_float(last.year_index)
+            if first_year is not None and last_year is not None and last_year > first_year:
+                delta_years = last_year - first_year
+
+        # Final fallback to disease duration deltas.
+        if delta_years is None or delta_years <= 0:
+            first_duration = _safe_float(first.duration_years)
+            last_duration = _safe_float(last.duration_years)
+            if (
+                first_duration is not None
+                and last_duration is not None
+                and last_duration > first_duration
+            ):
+                delta_years = last_duration - first_duration
+
+        if delta_years is None or delta_years <= 0:
+            return None
+
+        first_composite = (
+            self._motor_burden_index(first)
+            + self._cognitive_burden_index(first)
+            + self._non_motor_burden_index(first)
+        ) / 3.0
+        last_composite = (
+            self._motor_burden_index(last)
+            + self._cognitive_burden_index(last)
+            + self._non_motor_burden_index(last)
+        ) / 3.0
+        return (last_composite - first_composite) / delta_years
+
+    def _treatment_response_proxy(self, snapshot: TwinSnapshot) -> Optional[float]:
+        updrs_off = snapshot.motor.get("updrs3_score")
+        updrs_on = snapshot.motor.get("updrs3_score_on")
+        if updrs_off is not None and updrs_on is not None:
+            return updrs_off - updrs_on
+
+        hy_off = snapshot.motor.get("hy")
+        hy_on = snapshot.motor.get("hy_on")
+        if hy_off is not None and hy_on is not None:
+            return hy_off - hy_on
+        return None
+
+    def _build_evidence(
+        self,
+        profile: TwinStaticProfile,
+        snapshot: TwinSnapshot,
+        motor_index: float,
+        cognitive_index: float,
+        non_motor_index: float,
+        progression_velocity: Optional[float],
+        prediction_summary: Dict[str, Any],
+        bridge_result: Optional[Dict[str, Any]] = None,
+    ) -> List[str]:
+        br = bridge_result or {}
+        evidence = [
+            f"Current cohort estimate uses {prediction_summary.get('source', 'heuristic')} inference.",
+            "Forecasts use cluster-weighted trajectory prediction (v2) and should be treated as decision support only.",
+        ]
+
+        data_source = br.get("data_source")
+        if data_source:
+            evidence.append(f"Cohort split enforced at inference with source: {data_source}.")
+
+        cluster_label = br.get("cluster_label")
+        if cluster_label:
+            evidence.append(f"Assigned progression cluster: {cluster_label} progressor.")
+
+        treatment_effect = br.get("treatment_effect")
+        if treatment_effect and treatment_effect > 0:
+            evidence.append(f"Estimated treatment effect (LEDD): {treatment_effect:.1f} UPDRS3 point reduction.")
+
+        ci_lo = br.get("ci_lower")
+        ci_hi = br.get("ci_upper")
+        if ci_lo is not None and ci_hi is not None:
+            evidence.append(f"Risk confidence interval (bootstrap 100): [{ci_lo:.2f}, {ci_hi:.2f}].")
+
+        if snapshot.ledd is not None:
+            evidence.append(f"Medication context captured via LEDD {snapshot.ledd:.1f}.")
+        if profile.subgroup:
+            evidence.append(f"Patient subgroup context: {profile.subgroup}.")
+        if motor_index >= 0.55:
+            evidence.append("Motor burden is elevated relative to the entered symptom profile.")
+        if cognitive_index >= 0.4:
+            evidence.append("Cognitive burden suggests closer monitoring of executive and memory measures.")
+        if non_motor_index >= 0.45:
+            evidence.append("Non-motor burden is material and likely to affect quality of life trajectory.")
+        if progression_velocity is not None:
+            evidence.append(f"Estimated progression velocity across snapshots: {progression_velocity:.2f} burden units/year.")
+
+        return evidence
+
+    def _risk_level(
+        self,
+        predicted_updrs3: Optional[float],
+        predicted_moca: Optional[float],
+        cohort_estimate: str,
+    ) -> str:
+        if cohort_estimate == "Parkinson's Disease" and (predicted_updrs3 or 0) >= 20:
+            return "high"
+        if predicted_moca is not None and predicted_moca < 24:
+            return "high"
+        if (predicted_updrs3 or 0) >= 10:
+            return "medium"
+        return "low"
+
+    def _profile_from_dict(self, payload: Dict[str, Any]) -> TwinStaticProfile:
+        return TwinStaticProfile(**payload)
+
+    def _snapshot_from_dict(self, payload: Dict[str, Any]) -> TwinSnapshot:
+        return TwinSnapshot(**payload)

src/twin_predictor_bridge.py

@@ -0,0 +1,269 @@
+"""
+TwinPredictorBridge — ML / Heuristic Fallback Router.
+
+Wraps progression_model, treatment_model, and risk_stratifier behind
+a single predict() method. If any ML model fails, it falls back to
+heuristic defaults. Also enforces PPMI / patient-entered cohort
+separation and provides a MODELS_LOADED status flag.
+"""
+
+from __future__ import annotations
+
+import logging
+import os
+from pathlib import Path
+from typing import Any, Dict, List, Optional, Tuple
+
+logger = logging.getLogger(__name__)
+
+
+class TwinPredictorBridge:
+    """Thin bridge routing to ML models with heuristic fallback."""
+
+    def __init__(self, ppmi_csv_path: Optional[str] = None) -> None:
+        self._ppmi_csv = ppmi_csv_path or self._find_ppmi_csv()
+        self.progression = None
+        self.treatment = None
+        self.risk = None
+        self.models_loaded = False
+        self._load_models()
+
+    # ------------------------------------------------------------------
+    # Public API
+    # ------------------------------------------------------------------
+    def predict(
+        self,
+        patient_data: Dict[str, Any],
+        snapshots: Optional[List[Dict[str, Any]]] = None,
+    ) -> Dict[str, Any]:
+        """
+        Run all three sub-models and return unified result.
+
+        Returns dict with: cluster_id, cluster_label, treatment_effect,
+        risk_category, confidence, ci_lower, ci_upper, progression_profile.
+        """
+        snaps = snapshots or []
+        data_source = self._infer_data_source(patient_data, snaps)
+
+        # 1. Progression cluster
+        cluster_id, cluster_label = self._assign_cluster(snaps, patient_data)
+        profile = self._get_cluster_profile(cluster_label)
+
+        # 2. Treatment effect
+        treatment_effect = self._predict_treatment(patient_data)
+
+        # 3. Risk stratification
+        risk = self._stratify_risk(patient_data)
+
+        return {
+            "cluster_id": cluster_id,
+            "cluster_label": cluster_label,
+            "progression_profile": profile,
+            "treatment_effect": treatment_effect,
+            "risk_category": risk.get("category", "Unknown"),
+            "confidence": risk.get("confidence", 0.0),
+            "ci_lower": risk.get("ci_lower", 0.0),
+            "ci_upper": risk.get("ci_upper", 0.0),
+            "risk_details": risk,
+            "models_loaded": self.models_loaded,
+            "cohort_split_enforced": True,
+            "data_source": data_source,
+        }
+
+    def get_status(self) -> Dict[str, Any]:
+        """Return model-loading status for the /api/health endpoint."""
+        return {
+            "models_loaded": self.models_loaded,
+            "cohort_split_enforced": True,
+            "progression_fitted": self.progression is not None and self.progression.fitted,
+            "treatment_fitted": self.treatment is not None and self.treatment.fitted,
+            "risk_available": self.risk is not None,
+            "ppmi_csv": self._ppmi_csv or "not found",
+            "silhouette_score": (
+                self.progression.silhouette_score_
+                if self.progression and self.progression.silhouette_score_ is not None
+                else None
+            ),
+            "treatment_r_squared": (
+                self.treatment.r_squared
+                if self.treatment and self.treatment.r_squared is not None
+                else None
+            ),
+        }
+
+    # ------------------------------------------------------------------
+    # Model loading
+    # ------------------------------------------------------------------
+    def _load_models(self) -> None:
+        loaded_count = 0
+
+        # Progression model
+        try:
+            from progression_model import ProgressionModel
+            self.progression = ProgressionModel()
+            if self._ppmi_csv:
+                self.progression.fit(self._ppmi_csv)
+            loaded_count += 1
+            if self.progression.silhouette_score_ is not None:
+                logger.info(
+                    "Progression silhouette score: %.3f",
+                    self.progression.silhouette_score_,
+                )
+        except Exception as exc:
+            logger.warning("Failed to load progression model: %s", exc)
+            self.progression = None
+
+        # Treatment model
+        try:
+            from treatment_model import TreatmentModel
+            self.treatment = TreatmentModel()
+            if self._ppmi_csv:
+                self.treatment.fit(self._ppmi_csv)
+            loaded_count += 1
+            if self.treatment.r_squared is not None:
+                logger.info("Treatment R²: %.4f", self.treatment.r_squared)
+        except Exception as exc:
+            logger.warning("Failed to load treatment model: %s", exc)
+            self.treatment = None
+
+        # Risk stratifier
+        try:
+            from risk_stratifier import RiskStratifier
+            self.risk = RiskStratifier(n_bootstrap=100)
+            loaded_count += 1
+        except Exception as exc:
+            logger.warning("Failed to load risk stratifier: %s", exc)
+            self.risk = None
+
+        self.models_loaded = bool(
+            self.progression is not None
+            and self.progression.fitted
+            and self.treatment is not None
+            and self.treatment.fitted
+            and self.risk is not None
+        )
+        logger.info("TwinPredictorBridge: %d/3 models loaded", loaded_count)
+
+    # ------------------------------------------------------------------
+    # Sub-model routing with fallback
+    # ------------------------------------------------------------------
+    def _assign_cluster(
+        self,
+        snapshots: List[Dict[str, Any]],
+        patient_data: Dict[str, Any],
+    ) -> Tuple[str, str]:
+        try:
+            if self.progression is not None:
+                return self.progression.assign_cluster(snapshots, patient_data)
+        except Exception as exc:
+            logger.warning("Progression cluster fallback: %s", exc)
+        return self._heuristic_cluster(patient_data)
+
+    def _get_cluster_profile(self, cluster_label: str) -> Dict[str, float]:
+        try:
+            if self.progression is not None:
+                return self.progression.get_cluster_profile(cluster_label)
+        except Exception:
+            pass
+        from progression_model import DEFAULT_CLUSTER_PROFILES
+        return DEFAULT_CLUSTER_PROFILES.get(cluster_label, DEFAULT_CLUSTER_PROFILES["moderate"])
+
+    def _predict_treatment(self, patient_data: Dict[str, Any]) -> float:
+        try:
+            if self.treatment is not None:
+                ledd = patient_data.get("ledd") or patient_data.get("LEDD")
+                dur = patient_data.get("duration_years") or patient_data.get("duration_yrs") or 0.0
+                if ledd is not None:
+                    return self.treatment.predict_treatment_effect(float(ledd), float(dur))
+        except Exception as exc:
+            logger.warning("Treatment model fallback: %s", exc)
+        return self._heuristic_treatment(patient_data)
+
+    def _stratify_risk(self, patient_data: Dict[str, Any]) -> Dict[str, Any]:
+        try:
+            if self.risk is not None:
+                return self.risk.stratify(patient_data)
+        except Exception as exc:
+            logger.warning("Risk stratifier fallback: %s", exc)
+        return self._heuristic_risk(patient_data)
+
+    # ------------------------------------------------------------------
+    # Heuristic fallbacks
+    # ------------------------------------------------------------------
+    @staticmethod
+    def _heuristic_cluster(patient_data: Dict[str, Any]) -> Tuple[str, str]:
+        updrs3 = patient_data.get("updrs3_score") or 0
+        try:
+            updrs3 = float(updrs3)
+        except (TypeError, ValueError):
+            updrs3 = 0.0
+        if updrs3 >= 25:
+            return ("2", "fast")
+        if updrs3 >= 12:
+            return ("1", "moderate")
+        return ("0", "slow")
+
+    @staticmethod
+    def _heuristic_treatment(patient_data: Dict[str, Any]) -> float:
+        updrs_off = patient_data.get("updrs3_score")
+        updrs_on = patient_data.get("updrs3_score_on")
+        if updrs_off is not None and updrs_on is not None:
+            try:
+                return max(0.0, float(updrs_off) - float(updrs_on))
+            except (TypeError, ValueError):
+                pass
+        return 0.0
+
+    @staticmethod
+    def _heuristic_risk(patient_data: Dict[str, Any]) -> Dict[str, Any]:
+        sym = sum(
+            float(patient_data.get(s) or 0)
+            for s in ("sym_tremor", "sym_rigid", "sym_brady", "sym_posins")
+        )
+        if sym >= 5:
+            cat, conf = "PD", 0.62
+        elif sym >= 2:
+            cat, conf = "Prodromal PD", 0.50
+        else:
+            cat, conf = "HC", 0.55
+        return {
+            "category": cat,
+            "confidence": conf,
+            "ci_lower": round(conf - 0.08, 4),
+            "ci_upper": round(conf + 0.08, 4),
+        }
+
+    # ------------------------------------------------------------------
+    @staticmethod
+    def _infer_data_source(
+        patient_data: Dict[str, Any],
+        snapshots: List[Dict[str, Any]],
+    ) -> str:
+        """
+        Resolve cohort source used for prediction-time separation.
+
+        This is metadata-only enforcement for inference: models are always
+        trained from PPMI CSV and never refit from patient-entered payloads.
+        """
+        source_tag = str(patient_data.get("source") or patient_data.get("data_source") or "").strip().lower()
+        if source_tag in {"ppmi", "ppmi_validation", "ppmi_train"}:
+            return "ppmi_validation"
+
+        if patient_data.get("source_patno") is not None or patient_data.get("PATNO") is not None:
+            return "ppmi_validation"
+
+        for snap in snapshots:
+            raw = snap.get("raw_inputs") if isinstance(snap, dict) else None
+            if isinstance(raw, dict) and (raw.get("source_patno") is not None or raw.get("PATNO") is not None):
+                return "ppmi_validation"
+
+        return "patient_entered_validation"
+
+    # ------------------------------------------------------------------
+    @staticmethod
+    def _find_ppmi_csv() -> Optional[str]:
+        project_root = Path(__file__).resolve().parent.parent
+        candidates = sorted(project_root.glob("PPMI_Curated_Data_Cut_Public_*.csv"), reverse=True)
+        if candidates:
+            return str(candidates[0])
+        return None

src/twin_schema.py

@@ -0,0 +1,120 @@
+from __future__ import annotations
+
+from dataclasses import asdict, dataclass, field
+from typing import Any, Dict, List, Optional
+
+
+@dataclass
+class TwinStaticProfile:
+    twin_id: str
+    patient_label: str
+    source_patno: Optional[int]
+    created_at: str
+    enrollment_cohort: str
+    subgroup: Optional[str]
+    sex: Optional[float]
+    education_years: Optional[float]
+    race: Optional[float]
+    family_pd: Optional[float]
+    family_pd_bin: Optional[float]
+    bmi: Optional[float]
+    age_diag: Optional[float]
+    age_onset: Optional[float]
+    dominant_side: Optional[float]
+
+    def to_dict(self) -> Dict[str, Any]:
+        return asdict(self)
+
+
+@dataclass
+class TwinSnapshot:
+    snapshot_id: str
+    event_id: str
+    visit_date: str
+    year_index: Optional[float]
+    age_at_visit: Optional[float]
+    duration_years: Optional[float]
+    treatment_flag: Optional[float]
+    ledd: Optional[float]
+    motor: Dict[str, Optional[float]]
+    cognition: Dict[str, Optional[float]]
+    non_motor: Dict[str, Optional[float]]
+    autonomic: Dict[str, Optional[float]]
+    biomarkers: Dict[str, Optional[float]] = field(default_factory=dict)
+    imaging: Dict[str, Optional[float]] = field(default_factory=dict)
+    raw_inputs: Dict[str, Any] = field(default_factory=dict)
+
+    def to_dict(self) -> Dict[str, Any]:
+        return asdict(self)
+
+
+@dataclass
+class TwinState:
+    current_cohort_estimate: str
+    prediction_source: str
+    confidence: float
+    motor_burden_index: Optional[float]
+    cognitive_burden_index: Optional[float]
+    non_motor_burden_index: Optional[float]
+    progression_velocity: Optional[float]
+    treatment_response_proxy: Optional[float]
+    computed_at: str
+    cluster_id: Optional[str] = None
+    cluster_label: Optional[str] = None
+    treatment_effect: Optional[float] = None
+    ci_lower: Optional[float] = None
+    ci_upper: Optional[float] = None
+    evidence: List[str] = field(default_factory=list)
+
+    def to_dict(self) -> Dict[str, Any]:
+        return asdict(self)
+
+
+@dataclass
+class TwinForecastPoint:
+    horizon_months: int
+    predicted_updrs3: Optional[float]
+    predicted_total_updrs: Optional[float]
+    predicted_moca: Optional[float]
+    predicted_hy: Optional[float]
+    risk_level: str
+    uncertainty: Dict[str, Optional[float]]
+
+    def to_dict(self) -> Dict[str, Any]:
+        return asdict(self)
+
+
+@dataclass
+class TwinSimulation:
+    scenario_name: str
+    overrides: Dict[str, Any]
+    simulated_snapshot: TwinSnapshot
+    state: TwinState
+    forecast: List[TwinForecastPoint]
+
+    def to_dict(self) -> Dict[str, Any]:
+        return {
+            "scenario_name": self.scenario_name,
+            "overrides": self.overrides,
+            "simulated_snapshot": self.simulated_snapshot.to_dict(),
+            "state": self.state.to_dict(),
+            "forecast": [point.to_dict() for point in self.forecast],
+        }
+
+
+@dataclass
+class DigitalTwin:
+    profile: TwinStaticProfile
+    snapshots: List[TwinSnapshot]
+    current_state: TwinState
+    forecast: List[TwinForecastPoint]
+    prediction_summary: Dict[str, Any] = field(default_factory=dict)
+
+    def to_dict(self) -> Dict[str, Any]:
+        return {
+            "profile": self.profile.to_dict(),
+            "snapshots": [snapshot.to_dict() for snapshot in self.snapshots],
+            "current_state": self.current_state.to_dict(),
+            "forecast": [point.to_dict() for point in self.forecast],
+            "prediction_summary": self.prediction_summary,
+        }

src/twin_store.py

@@ -0,0 +1,212 @@
+from __future__ import annotations
+
+import json
+import sqlite3
+from contextlib import closing
+from pathlib import Path
+from typing import Any, Dict, List, Optional
+
+from twin_schema import DigitalTwin
+
+
+class TwinStore:
+    def __init__(self, db_path: Optional[str] = None):
+        if db_path is None:
+            project_root = Path(__file__).resolve().parent.parent
+            db_path = str(project_root / "data" / "digital_twins.sqlite3")
+
+        self.db_path = Path(db_path)
+        self.db_path.parent.mkdir(parents=True, exist_ok=True)
+        self._init_db()
+
+    def _connect(self) -> sqlite3.Connection:
+        connection = sqlite3.connect(str(self.db_path))
+        connection.row_factory = sqlite3.Row
+        return connection
+
+    def _init_db(self) -> None:
+        with closing(self._connect()) as conn:
+            conn.execute(
+                """
+                CREATE TABLE IF NOT EXISTS twins (
+                    twin_id TEXT PRIMARY KEY,
+                    patient_label TEXT NOT NULL,
+                    source_patno INTEGER,
+                    created_at TEXT NOT NULL,
+                    updated_at TEXT NOT NULL,
+                    current_cohort_estimate TEXT NOT NULL,
+                    confidence REAL NOT NULL,
+                    snapshot_count INTEGER NOT NULL,
+                    profile_json TEXT NOT NULL,
+                    current_state_json TEXT NOT NULL,
+                    forecast_json TEXT NOT NULL,
+                    prediction_json TEXT NOT NULL
+                )
+                """
+            )
+            conn.execute(
+                """
+                CREATE TABLE IF NOT EXISTS twin_snapshots (
+                    snapshot_id TEXT PRIMARY KEY,
+                    twin_id TEXT NOT NULL,
+                    snapshot_index INTEGER NOT NULL,
+                    created_at TEXT NOT NULL,
+                    snapshot_json TEXT NOT NULL,
+                    FOREIGN KEY (twin_id) REFERENCES twins(twin_id) ON DELETE CASCADE
+                )
+                """
+            )
+            conn.execute(
+                "CREATE INDEX IF NOT EXISTS idx_twin_snapshots_twin_id ON twin_snapshots(twin_id)"
+            )
+            conn.commit()
+
+    def upsert_twin(self, twin: DigitalTwin) -> Dict[str, Any]:
+        payload = twin.to_dict()
+        profile = payload["profile"]
+        state = payload["current_state"]
+        forecast = payload["forecast"]
+        prediction = payload.get("prediction_summary", {})
+        snapshots = payload["snapshots"]
+
+        with closing(self._connect()) as conn:
+            existing = conn.execute(
+                "SELECT twin_id FROM twins WHERE twin_id = ?",
+                (profile["twin_id"],),
+            ).fetchone()
+
+            values = (
+                profile["twin_id"],
+                profile.get("patient_label") or profile["twin_id"],
+                profile.get("source_patno"),
+                profile["created_at"],
+                state["computed_at"],
+                state["current_cohort_estimate"],
+                float(state.get("confidence") or 0.0),
+                len(snapshots),
+                json.dumps(profile),
+                json.dumps(state),
+                json.dumps(forecast),
+                json.dumps(prediction),
+            )
+
+            if existing:
+                conn.execute(
+                    """
+                    UPDATE twins
+                    SET patient_label = ?, source_patno = ?, created_at = ?, updated_at = ?,
+                        current_cohort_estimate = ?, confidence = ?, snapshot_count = ?,
+                        profile_json = ?, current_state_json = ?, forecast_json = ?, prediction_json = ?
+                    WHERE twin_id = ?
+                    """,
+                    (
+                        values[1],
+                        values[2],
+                        values[3],
+                        values[4],
+                        values[5],
+                        values[6],
+                        values[7],
+                        values[8],
+                        values[9],
+                        values[10],
+                        values[11],
+                        values[0],
+                    ),
+                )
+                conn.execute(
+                    "DELETE FROM twin_snapshots WHERE twin_id = ?",
+                    (profile["twin_id"],),
+                )
+            else:
+                conn.execute(
+                    """
+                    INSERT INTO twins (
+                        twin_id, patient_label, source_patno, created_at, updated_at,
+                        current_cohort_estimate, confidence, snapshot_count,
+                        profile_json, current_state_json, forecast_json, prediction_json
+                    )
+                    VALUES (?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?)
+                    """,
+                    values,
+                )
+
+            for index, snapshot in enumerate(snapshots):
+                conn.execute(
+                    """
+                    INSERT INTO twin_snapshots (
+                        snapshot_id, twin_id, snapshot_index, created_at, snapshot_json
+                    )
+                    VALUES (?, ?, ?, ?, ?)
+                    """,
+                    (
+                        snapshot["snapshot_id"],
+                        profile["twin_id"],
+                        index,
+                        snapshot.get("visit_date") or profile["created_at"],
+                        json.dumps(snapshot),
+                    ),
+                )
+            conn.commit()
+
+        stored = self.get_twin(profile["twin_id"])
+        if stored is None:
+            raise RuntimeError("Failed to persist digital twin")
+        return stored
+
+    def list_twins(self) -> List[Dict[str, Any]]:
+        with closing(self._connect()) as conn:
+            rows = conn.execute(
+                """
+                SELECT twin_id, patient_label, source_patno, created_at, updated_at,
+                       current_cohort_estimate, confidence, snapshot_count
+                FROM twins
+                ORDER BY updated_at DESC
+                """
+            ).fetchall()
+
+        return [dict(row) for row in rows]
+
+    def get_twin(self, twin_id: str) -> Optional[Dict[str, Any]]:
+        with closing(self._connect()) as conn:
+            row = conn.execute(
+                """
+                SELECT twin_id, patient_label, source_patno, created_at, updated_at,
+                       current_cohort_estimate, confidence, snapshot_count,
+                       profile_json, current_state_json, forecast_json, prediction_json
+                FROM twins
+                WHERE twin_id = ?
+                """,
+                (twin_id,),
+            ).fetchone()
+
+            if row is None:
+                return None
+
+            snapshots = conn.execute(
+                """
+                SELECT snapshot_json
+                FROM twin_snapshots
+                WHERE twin_id = ?
+                ORDER BY snapshot_index ASC
+                """,
+                (twin_id,),
+            ).fetchall()
+
+        return {
+            "profile": json.loads(row["profile_json"]),
+            "current_state": json.loads(row["current_state_json"]),
+            "forecast": json.loads(row["forecast_json"]),
+            "prediction_summary": json.loads(row["prediction_json"]),
+            "snapshots": [json.loads(snapshot["snapshot_json"]) for snapshot in snapshots],
+            "summary": {
+                "twin_id": row["twin_id"],
+                "patient_label": row["patient_label"],
+                "source_patno": row["source_patno"],
+                "created_at": row["created_at"],
+                "updated_at": row["updated_at"],
+                "current_cohort_estimate": row["current_cohort_estimate"],
+                "confidence": row["confidence"],
+                "snapshot_count": row["snapshot_count"],
+            },
+        }