classifier_last_token.py

"""
CHR vs. HC linear probe on a frozen HF decoder, LAST-TOKEN pooled.

Same pipeline as classifier_mean_pool.py, but each patient is summarized by the
hidden state of the LAST real token of their concatenated transcript instead of
the mean over all tokens. In a causal decoder the final token is the only
position that has attended to the entire sequence, so its representation is the
natural whole-sequence summary. This is the frozen-backbone adaptation of the
central finding in:

  Suganthan et al., "Adapting Decoder-Based Language Models for Diverse Encoder
  Downstream Tasks," arXiv:2503.02656 (2025).

--------------------------------------------------------------------------------
Suggested improvements distilled from that paper (and how each maps here)
--------------------------------------------------------------------------------
The paper adapts a decoder (Gemma) into an encoder and ablates pooling +
attention masking on GLUE / MS MARCO. Their headline results (Gemma-9B, GLUE):
bidirectional attention 90.9 vs causal 87.5; last-token pooling 90.9 > mean 89.7
> first-K; dropout optimum at 10%. THE CAVEAT: every bidirectional number comes
from FINE-TUNING the backbone, whereas this project keeps the backbone FROZEN
and trains only a light head. So their findings split into what survives a
freeze and what does not:

  1. Last-token ("Last-K", K=1) pooling  [IMPLEMENTED HERE]
     Well-founded for a frozen causal decoder: the paper shows last-token beats
     mean, and *specifically* wins under causal masking (the regime we are stuck
     in). Replaces masked-mean pooling. Highest-value, cheapest change.

  2. Dropout ~10% on the head  [see classifier_last_token_dropout.py]
     The paper's tuned regularization optimum (on attention softmax + FFN
     outputs). Head-only, needs no re-extraction, and directly targets the head
     instability we saw (FFN heads collapsing at the smallest and largest
     backbones).

  3. Last-K token pooling, K > 1  [see classifier_last_k.py]
     Mean of the final K real tokens rather than just the last one. Tests
     whether a short tail window carries more signal than the single final
     token. Contingent: only interesting if last-token already helps.

  4. Bidirectional attention during extraction  [NOT IMPLEMENTED - speculative]
     The paper's single biggest lever (+3.4), but it comes from fine-tuning.
     Feeding a frozen, causally-trained model a bidirectional mask at inference
     shows it an attention pattern its weights never saw; it may help or hurt.
     Worth a future ablation (a new cache keyed by attn="bidirectional"), but
     out of scope for this pooling-only pass.

  5. De-scoped by design:
       - First-K pooling: on a causal decoder the first tokens have attended to
         almost nothing, so it is theoretically wrong here.
       - KV-Probe attention variant: our attention_pool heads are already the
         paper's Query-Probe (a learned query over the sequence), and it was not
         the paper's winner.
       - Padding side: the paper found it does not matter after adaptation; we
         mask/index padding out regardless (and force right-padding for
         deterministic last-token indexing).

--------------------------------------------------------------------------------
Pipeline (identical to mean_pool except the pooling reduction):
  1. Concatenate each patient's participant utterances into one string.
  2. Run the FROZEN decoder and take the LAST real token's hidden state
     -> one (d,) vector per patient.
  3. Cache embeddings to disk, keyed by model/pooling/max_length (its own cache,
     never colliding with the mean-pool cache).
  4. Stratified 80/20 split; train Linear(d -> 1) head (class-weighted BCE).
  5. Tune the decision threshold on TRAIN, then report test metrics and write
     results to results/<model>/last_token/.
"""
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)
POOLING = "last"           # last real token of the transcript

# 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. "transcript_lasttoken" keeps this
# cache separate from the mean-pool "transcript" cache; the dropout variant
# reuses the same pooling, so it shares this cache (identical embeddings).
CACHE_DIR = embedding_cache_dir(MODEL_ID, "transcript_lasttoken", 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": "last_token_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 last-token embeddings (one-time cost)...")
    embedder = HFEmbedder(MODEL_ID, device, MAX_LENGTH, load_mode=LOAD_MODE, pooling=POOLING)
    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_lasttoken", 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": "last_token_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, "last_token", MAX_LENGTH, SEED),
        model_name="last_token", 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()