classifier_mean_pool.py

"""
CHR vs. HC linear probe on a frozen HF base model (mean-pool baseline).

This is the SIMPLE baseline: one vector per patient via masked mean-pooling
over the whole concatenated transcript. Split, training, threshold tuning,
metrics, reporting and saved artifacts are shared with the attention-pool
models (see shared/utils.py) so the variants are directly comparable.

Pipeline:
  1. Concatenate each patient's participant utterances into one string.
  2. Run the FROZEN base model over raw text and masked mean-pool the last
     hidden state -> one (d,) vector per patient.
  3. Cache embeddings to disk (invalidated if model/config/source dir change).
  4. Stratified shuffled 80/20 split.
  5. Train Linear(d -> 1) head (per-patient grad accumulation, class-weighted
     BCE), monitoring train/test loss and accuracy every epoch.
  6. Tune the decision threshold on TRAIN, then report test metrics and write
     results to results/<model>/mean_pool/.
"""
import os
os.environ['CUDA_DEVICE_ORDER'] = 'PCI_BUS_ID'

import numpy as np
import torch
import torch.nn as nn
from labl.transcripts import Transcript
from typing import Tuple
from dotenv import load_dotenv

from shared.utils import (
    set_seed,
    dir_fingerprint,
    HFEmbedder,
    parse_runtime_args,
    embedding_cache_dir,
    load_dense_cache,
    save_dense_cache,
    stratified_split,
    take,
    train_head,
    predict_probs,
    select_threshold,
    compute_metrics,
    format_report,
    results_dir,
    save_results,
)

load_dotenv()

# ----------------------------------------------------------------------
# Config
# ----------------------------------------------------------------------
GPU = 0
MODEL_ID = "google/gemma-4-E2B"
LOAD_MODE = "bf16"         # bf16 | 8bit | 4bit (quantize large backbones to fit)
MAX_LENGTH = 8192          # truncate long transcripts to avoid OOM
EXTRACT_BATCH = 8          # transcripts per forward pass during extraction
ACCUM_STEPS = 8           # patients accumulated per optimizer step (sim. batch size)
LR = 1e-3
EPOCHS = 100
TEST_SIZE = 0.20
SEED = 42
SESSION = "session001"     # only use baseline-session transcripts (one per patient)

# Threshold tuning: how to pick the CHR/HC decision cutoff from TRAIN probs.
#   "f1"               -> maximize CHR F1 (balanced precision/recall on CHR)
#   "youden"           -> maximize TPR - FPR (ROC operating point, cost-neutral)
#   "chr_recall_floor" -> maximize CHR recall s.t. HC recall >= HC_RECALL_FLOOR
THRESHOLD_CRITERION = "youden"
HC_RECALL_FLOOR = 0.50     # only used when criterion == "chr_recall_floor"

TRANSCRIPT_DIR = os.environ.get("TRANSCRIPT_DIR")
if not TRANSCRIPT_DIR:
    raise RuntimeError(
        "TRANSCRIPT_DIR is not set. Copy .env.example to .env and set it "
        "to the absolute path of the sorted transcripts directory."
    )
# Keyed by model / featurization / max_length so other models and scripts
# never share or invalidate this cache; anchored to the repo, not the CWD.
CACHE_DIR = embedding_cache_dir(MODEL_ID, "transcript", MAX_LENGTH)


# ----------------------------------------------------------------------
# Cache (one dense (N, d) tensor for the whole cohort)
# ----------------------------------------------------------------------
def cache_config() -> dict:
    return {"model_id": MODEL_ID, "load_mode": LOAD_MODE, "max_length": MAX_LENGTH,
            "pooling": "masked_mean_transcript", "session": SESSION}


def build_cache(device: str) -> Tuple[torch.Tensor, np.ndarray]:
    fingerprint = dir_fingerprint(TRANSCRIPT_DIR)

    cached = load_dense_cache(CACHE_DIR, cache_config(), fingerprint)
    if cached is not None:
        print("valid cache found, loading from disk")
        return cached

    print("cache missing or stale (config or source dir changed), rebuilding")
    Transcript.set_directory_path(TRANSCRIPT_DIR)
    texts, labels = [], []
    for t in Transcript.list_transcripts():
        # guard against non-baseline sessions
        if t.session != SESSION:
            continue
        utterances = [line.text for line in t.participant_lines if line.text.strip()]
        if not utterances:
            continue
        texts.append(" ".join(utterances))
        labels.append(1 if t.group_status.name == "CHR" else 0)

    y = np.array(labels, dtype=np.float32)
    print(f"loaded {len(texts)} transcripts  "
          f"(CHR={int(y.sum())}, HC={int((y == 0).sum())})")

    print("extracting embeddings (one-time cost)...")
    embedder = HFEmbedder(MODEL_ID, device, MAX_LENGTH, load_mode=LOAD_MODE)
    X = embedder.embed(texts, EXTRACT_BATCH)  # (N, d) on CPU

    save_dense_cache(CACHE_DIR, cache_config(), fingerprint, X, y)

    del embedder
    torch.cuda.empty_cache()
    return X, y


# ----------------------------------------------------------------------
# Classifier head
# ----------------------------------------------------------------------
class LinearHead(nn.Module):
    """Single linear layer mapping one patient's pooled embedding to one logit."""

    def __init__(self, hidden_size: int):
        super().__init__()
        self.fc = nn.Linear(hidden_size, 1)

    def forward(self, x: torch.Tensor) -> torch.Tensor:
        # x: (d,) one patient -> (1,) logit
        return self.fc(x)


# ----------------------------------------------------------------------
# Main
# ----------------------------------------------------------------------
def main():
    global MODEL_ID, LOAD_MODE, EXTRACT_BATCH, GPU, MAX_LENGTH, SEED, CACHE_DIR
    args = parse_runtime_args(MODEL_ID, EXTRACT_BATCH, GPU, MAX_LENGTH, SEED)
    MODEL_ID, LOAD_MODE, EXTRACT_BATCH, GPU, MAX_LENGTH, SEED = (
        args.model, args.load_mode, args.extract_batch, args.gpu,
        args.max_length, args.seed,
    )
    # model / max_length may have changed via CLI, so recompute the cache dir
    CACHE_DIR = embedding_cache_dir(MODEL_ID, "transcript", MAX_LENGTH)

    set_seed(SEED)
    os.environ['CUDA_VISIBLE_DEVICES'] = str(GPU)
    device = "cuda:0" if torch.cuda.is_available() else "cpu"

    X, y = build_cache(device)
    n_chr, n_hc = int(y.sum()), int((y == 0).sum())
    print(f"{X.shape[0]} patients  (CHR={n_chr}, HC={n_hc})  d={X.shape[1]}")

    train_idx, test_idx = stratified_split(y, TEST_SIZE, SEED)
    X_train, X_test = take(X, train_idx), take(X, test_idx)
    y_train, y_test = y[train_idx], y[test_idx]
    print(f"train={len(train_idx)}  test={len(test_idx)}")

    head = LinearHead(X.shape[1])
    history = train_head(
        head, X_train, y_train, X_test, y_test, device,
        epochs=EPOCHS, lr=LR, accum_steps=ACCUM_STEPS,
    )

    # derive threshold on TRAIN probs only, then freeze for test
    train_probs = predict_probs(head, X_train, device)
    tuned_t = select_threshold(train_probs, y_train, THRESHOLD_CRITERION, HC_RECALL_FLOOR)

    test_probs = predict_probs(head, X_test, device)
    m_default = compute_metrics(test_probs, y_test, 0.5)
    m_tuned = compute_metrics(test_probs, y_test, tuned_t)

    print("\n" + format_report(m_default, m_tuned, len(test_idx), THRESHOLD_CRITERION))

    config = {
        "model_id": MODEL_ID,
        "load_mode": LOAD_MODE,
        "max_length": MAX_LENGTH,
        "pooling": "masked_mean_transcript",
        "session": SESSION,
        "lr": LR,
        "epochs": EPOCHS,
        "accum_steps": ACCUM_STEPS,
        "test_size": TEST_SIZE,
        "seed": SEED,
        "n_chr": n_chr,
        "n_hc": n_hc,
    }
    save_results(
        results_dir(MODEL_ID, "mean_pool", MAX_LENGTH, SEED),
        model_name="mean_pool", config=config,
        history=history, m_default=m_default, m_tuned=m_tuned,
        n_train=len(train_idx), n_test=len(test_idx), criterion=THRESHOLD_CRITERION,
    )


if __name__ == "__main__":
    main()