""" src/nlp/classifier.py ──────────────────────────────────────────────────────────────── Clinical text classifier — fine-tunes Bio_ClinicalBERT for severity classification (routine / urgent / critical). Design ────── The classifier is task-agnostic: the ``task`` config parameter controls which label set is used. Swapping to readmission risk (Phase 2) is a one-line change in config.py. Two modes of operation: Training — fine-tunes the base model on MTSamples with weak-supervision severity labels. Saves the fine-tuned weights to data/models/. Inference — loads the fine-tuned weights and classifies a single note or a batch. Training takes ~20 minutes on CPU for MTSamples (~4,000 notes after filtering). On a free Colab GPU it takes under 5 minutes. The fine-tuned model is saved and reused on subsequent runs. The training loop is written from scratch using the HuggingFace Trainer API. This is more explicit than using a Pipeline object and gives full control over evaluation and checkpoint saving. ──────────────────────────────────────────────────────────────── """ from __future__ import annotations import json import os from pathlib import Path import numpy as np import pandas as pd from src.utils.config import ( ClassifierConfig, ModelConfig, TrainingConfig, force_offline_hf_env, ) from src.utils.logger import get_logger logger = get_logger(__name__) def _from_pretrained_offline_first(loader_cls, model_name_or_path, **kwargs): """Load a HuggingFace tokenizer/model from the local cache first. HuggingFace libraries make a handful of network round-trips (HEAD requests checking for config/adapter file updates) even when a model is fully cached locally — a transient network blip during one of those checks crashes the whole load. Trying ``local_files_only=True`` first avoids the network entirely when the model is already cached, and only falls back to a live download when nothing is cached yet. Args: loader_cls: A HuggingFace ``...from_pretrained`` class, e.g. ``AutoTokenizer`` or ``AutoModelForSequenceClassification``. model_name_or_path: HuggingFace Hub model ID or local path. **kwargs: Forwarded to ``from_pretrained``. """ try: with force_offline_hf_env(): return loader_cls.from_pretrained( model_name_or_path, local_files_only=True, **kwargs ) except Exception: logger.info( "%s not found in local cache — downloading: %s", loader_cls.__name__, model_name_or_path, ) return loader_cls.from_pretrained(model_name_or_path, **kwargs) # ── Output dataclass ────────────────────────────────────────────── class ClassificationResult: """The output of classifying a single clinical note. Attributes: label : Predicted class label (e.g. ``"urgent"``). confidence : Softmax probability of the predicted class. probabilities: Full probability distribution over all classes. task : Which classification task produced this result. """ def __init__( self, label: str, confidence: float, probabilities: dict[str, float], task: str = ClassifierConfig.task, ) -> None: self.label = label self.confidence = confidence self.probabilities = probabilities self.task = task def to_dict(self) -> dict: """Serialise to a plain dictionary for JSON responses.""" return { "label": self.label, "confidence": round(self.confidence, 3), "probabilities": { k: round(v, 3) for k, v in self.probabilities.items() }, "task": self.task, } def __repr__(self) -> str: return ( f"ClassificationResult(label={self.label!r}, " f"confidence={self.confidence:.3f})" ) # ── Dataset ─────────────────────────────────────────────────────── class ClinicalNoteDataset: """PyTorch Dataset wrapper for tokenised clinical notes. Handles tokenisation and encoding in __getitem__ rather than upfront, which keeps memory usage manageable for large datasets. Args: texts : List of clinical note strings. labels : Integer label indices parallel to texts. tokenizer : HuggingFace tokenizer instance. max_length: Maximum token length (default 512 for BERT). """ def __init__( self, texts: list[str], labels: list[int], tokenizer, max_length: int = TrainingConfig.max_length, ) -> None: self.texts = texts self.labels = labels self.tokenizer = tokenizer self.max_length = max_length def __len__(self) -> int: return len(self.texts) def __getitem__(self, idx: int) -> dict: """Return tokenised input for one example. Args: idx: Index into the dataset. Returns: Dict with ``input_ids``, ``attention_mask``, ``token_type_ids``, and ``labels`` tensors. """ import torch encoding = self.tokenizer( self.texts[idx], max_length = self.max_length, padding = "max_length", truncation = True, return_tensors = "pt", ) return { "input_ids": encoding["input_ids"].squeeze(), "attention_mask": encoding["attention_mask"].squeeze(), "labels": torch.tensor(self.labels[idx], dtype=torch.long), } # ── Classifier ──────────────────────────────────────────────────── class ClinicalClassifier: """Fine-tune and serve Bio_ClinicalBERT for clinical text classification. Args: task : Classification task. One of ``"severity"`` (Phase 1) or ``"readmission"`` (Phase 2). model_name : HuggingFace model ID for the base model. output_dir : Where to save fine-tuned weights. Example:: # Train clf = ClinicalClassifier(task="severity") clf.train(train_df) # Infer clf = ClinicalClassifier(task="severity") clf.load() result = clf.predict("Patient admitted to ICU following cardiac arrest.") print(result.label, result.confidence) # → critical 0.94 """ def __init__( self, task: str = ClassifierConfig.task, model_name: str = ModelConfig.classifier_model, output_dir: Path | None = None, ) -> None: self._task = task self._model_name = model_name self._output_dir = output_dir or ModelConfig.fine_tuned_dir self._labels = ClassifierConfig.active_labels(self._task) self._label2id = {lbl: i for i, lbl in enumerate(self._labels)} self._id2label = dict(enumerate(self._labels)) self._model = None self._tokenizer = None logger.info( "ClinicalClassifier init: task=%s, model=%s, labels=%s", self._task, self._model_name, self._labels, ) # ── Model I/O ───────────────────────────────────────────────── def _load_tokenizer(self): """Load the tokenizer, downloading it if necessary.""" if self._tokenizer is not None: return self._tokenizer from transformers import AutoTokenizer logger.info("Loading tokenizer: %s", self._model_name) self._tokenizer = _from_pretrained_offline_first( AutoTokenizer, self._model_name ) return self._tokenizer def load(self) -> None: """Load fine-tuned weights, preferring local disk over HF Hub over the untrained base model, in that order. Loads from ``output_dir`` if it contains a saved checkpoint. Otherwise falls back to ``ModelConfig.classifier_hf_hub_fallback`` — a copy of the fine-tuned weights on HuggingFace Hub, needed because a fresh deployment container won't have ``output_dir`` locally (the weights are too large to commit to git). Only if *that* also fails does this fall back to the untrained base ``model_name`` — silently doing so without trying the HF Hub copy first would produce confident-looking but meaningless severity predictions with no error raised. Call this before predict() if you have already trained. Raises: FileNotFoundError: If no source has a valid model and HuggingFace cannot fetch the base either. """ from transformers import AutoModelForSequenceClassification if self._output_dir.exists(): source = str(self._output_dir) elif ModelConfig.classifier_hf_hub_fallback: source = ModelConfig.classifier_hf_hub_fallback logger.warning( "No local checkpoint at %s — loading fine-tuned weights " "from HF Hub fallback (%s) instead of the untrained " "base model.", self._output_dir, source, ) else: source = self._model_name logger.warning( "No local checkpoint and no HF Hub fallback configured " "— loading UNTRAINED base model %s. Severity " "predictions will be meaningless.", source, ) logger.info("Loading classifier from: %s", source) self._model = _from_pretrained_offline_first( AutoModelForSequenceClassification, source, num_labels = len(self._labels), id2label = self._id2label, label2id = self._label2id, ) self._model.eval() self._load_tokenizer() logger.info("Classifier loaded OK") def save(self) -> None: """Save the fine-tuned model and tokenizer to disk. Also saves a label mapping JSON so the model can be reloaded without the original config. """ if self._model is None or self._tokenizer is None: raise RuntimeError("No model to save — train or load first.") self._output_dir.mkdir(parents=True, exist_ok=True) self._model.save_pretrained(self._output_dir) self._tokenizer.save_pretrained(self._output_dir) # Save label mapping for reloading without config label_map_path = self._output_dir / "label_map.json" label_map_path.write_text( json.dumps({"id2label": self._id2label, "task": self._task}, indent=2) ) logger.info("Model saved to %s", self._output_dir) # ── Training ────────────────────────────────────────────────── def train( self, df: pd.DataFrame, resume: bool = True, weighted_loss: bool = False, seed: int | None = None, ) -> dict[str, float]: """Fine-tune the base model on labelled clinical notes. Splits data into train / validation / test sets, trains for the configured number of epochs, evaluates on the validation set after each epoch, and saves the best checkpoint. Args: df: DataFrame with ``transcription`` and ``severity`` columns (or ``readmission`` for Phase 2). Run the ETL pipeline first to produce this. resume: If ``True`` (default) and a previously fine-tuned checkpoint already exists at ``output_dir``, continue training from those weights instead of the base pretrained model. This is a warm start, not a true resume — optimizer/scheduler state and epoch count are not restored, so the LR schedule still runs the full configured epoch count, but a re-run after an interrupted session does not throw away prior fine-tuning progress. Pass ``False`` to force training from the base model regardless of any existing checkpoint. weighted_loss: If ``True``, weight the cross-entropy loss inversely to each class's frequency in the training split (``n_samples / (n_classes * n_samples_c)``, the standard "balanced" formula). Use this when one class is under-represented (e.g. "critical" at ~15% of severity data) and the model is under-predicting it. seed: Random seed for the new classification head's weight initialisation and data loader shuffling. Defaults to ``TrainingConfig.random_seed``. Without this, the head sitting on top of the pretrained base gets a fresh random init every run -- on a small dataset with few epochs that alone can swing minority-class metrics significantly (observed: critical-class recall varying 0.71-0.83 across otherwise-identical unseeded runs). Fixing the seed makes a single run reproducible; it does not by itself make results *better* -- for that, train several seeds and keep the best (see scripts/train_severity_classifier.py --n-seeds). Returns: Dict with final evaluation metrics: ``val_accuracy``, ``val_f1``, ``test_accuracy``, ``test_f1``, ``per_class`` (precision/recall/f1/support per label), ``confusion_matrix``, and ``history`` (per-epoch validation metrics). Raises: KeyError: If the required columns are not present. """ label_col = self._task if label_col not in df.columns: raise KeyError( f"Column '{label_col}' not found. " f"Available: {list(df.columns)}" ) logger.info("Starting fine-tuning: task=%s", self._task) import random import numpy as np import torch from sklearn.model_selection import train_test_split from torch.optim import AdamW from torch.utils.data import DataLoader from transformers import ( AutoModelForSequenceClassification, get_linear_schedule_with_warmup, ) tokenizer = self._load_tokenizer() cfg = TrainingConfig # ── Seeding ──────────────────────────────────────────────── # Covers the classification head's random init (torch) and # DataLoader shuffling (also torch's RNG). random/numpy are # seeded too in case any dependency defers to them internally. seed = seed if seed is not None else cfg.random_seed random.seed(seed) np.random.seed(seed) torch.manual_seed(seed) torch.cuda.manual_seed_all(seed) logger.info("Training seed: %d", seed) # ── Encode labels ───────────────────────────────────────── valid_mask = df[label_col].isin(self._labels) df = df[valid_mask].copy() texts = df["transcription"].tolist() labels = [self._label2id[lbl] for lbl in df[label_col]] logger.info( "Training data: %d notes, label dist=%s", len(texts), {lbl: labels.count(i) for lbl, i in self._label2id.items()}, ) # ── Train / val / test split ────────────────────────────── # First cut off the test set, then split remaining into train/val x_trainval, x_test, y_trainval, y_test = train_test_split( texts, labels, test_size = cfg.test_split, random_state = cfg.random_seed, stratify = labels, ) val_ratio = cfg.val_split / (1 - cfg.test_split) x_train, x_val, y_train, y_val = train_test_split( x_trainval, y_trainval, test_size = val_ratio, random_state = cfg.random_seed, stratify = y_trainval, ) logger.info( "Split: train=%d val=%d test=%d", len(x_train), len(x_val), len(x_test), ) # ── Datasets and loaders ────────────────────────────────── train_ds = ClinicalNoteDataset(x_train, y_train, tokenizer) val_ds = ClinicalNoteDataset(x_val, y_val, tokenizer) test_ds = ClinicalNoteDataset(x_test, y_test, tokenizer) train_loader = DataLoader(train_ds, batch_size=cfg.batch_size, shuffle=True) val_loader = DataLoader(val_ds, batch_size=cfg.batch_size) test_loader = DataLoader(test_ds, batch_size=cfg.batch_size) # ── Model ───────────────────────────────────────────────── # Warm-start from an existing checkpoint if one is present — # avoids discarding prior fine-tuning progress on a re-run # after an interrupted/restarted session. checkpoint_exists = (self._output_dir / "config.json").exists() resuming = resume and checkpoint_exists model_source = str(self._output_dir) if resuming else self._model_name if resume and not checkpoint_exists: logger.info( "No existing checkpoint at %s — training from base model %s", self._output_dir, self._model_name, ) elif resuming: logger.info("Resuming fine-tuning from checkpoint: %s", self._output_dir) self._model = _from_pretrained_offline_first( AutoModelForSequenceClassification, model_source, num_labels = len(self._labels), id2label = self._id2label, label2id = self._label2id, ) device = torch.device("cuda" if torch.cuda.is_available() else "cpu") self._model.to(device) if device.type == "cpu": torch.set_num_threads(os.cpu_count()) logger.info( "Training on device: %s (%d threads)", device, torch.get_num_threads(), ) total_steps = len(train_loader) * cfg.epochs warmup_steps = total_steps // 10 # 10% warmup optimizer = AdamW(self._model.parameters(), lr=cfg.learning_rate) scheduler = get_linear_schedule_with_warmup( optimizer, num_warmup_steps = warmup_steps, num_training_steps = total_steps, ) # ── Class weights ───────────────────────────────────────── # Balanced formula: n_samples / (n_classes * n_samples_c). # Computed from the actual training split, not the full # dataset, since that's the distribution the model sees. loss_weight = None if weighted_loss: n_train, n_classes = len(y_train), len(self._labels) counts = [max(y_train.count(i), 1) for i in range(n_classes)] weights = [n_train / (n_classes * c) for c in counts] loss_weight = torch.tensor(weights, dtype=torch.float32, device=device) logger.info( "Weighted loss enabled: %s", dict(zip(self._labels, (round(w, 3) for w in weights), strict=False)), ) # ── Training loop ───────────────────────────────────────── # When resuming, evaluate the loaded checkpoint first so a # worse early epoch doesn't overwrite a better prior result. if resuming: baseline = self._evaluate(val_loader, device) best_val_f1, best_val_acc = baseline["f1"], baseline["accuracy"] logger.info( "Resumed checkpoint baseline: val_acc=%.4f val_f1=%.4f", best_val_acc, best_val_f1, ) else: best_val_f1 = 0.0 best_val_acc = 0.0 metrics_history = [] for epoch in range(1, cfg.epochs + 1): # Train self._model.train() epoch_loss = 0.0 for batch in train_loader: optimizer.zero_grad() batch = {k: v.to(device) for k, v in batch.items()} outputs = self._model(**batch) # Always compute loss manually via F.cross_entropy rather # than relying on outputs.loss (the model's internal loss # is unweighted) -- passing weight=None reproduces the # model's default behaviour exactly when weighted_loss # is off, so this is a single code path either way. loss = torch.nn.functional.cross_entropy( outputs.logits, batch["labels"], weight=loss_weight ) epoch_loss += loss.item() loss.backward() torch.nn.utils.clip_grad_norm_(self._model.parameters(), 1.0) optimizer.step() scheduler.step() avg_loss = epoch_loss / len(train_loader) # Validate val_metrics = self._evaluate(val_loader, device) logger.info( "Epoch %d/%d loss=%.4f val_acc=%.4f val_f1=%.4f", epoch, cfg.epochs, avg_loss, val_metrics["accuracy"], val_metrics["f1"], ) metrics_history.append({ "epoch": epoch, "loss": avg_loss, "accuracy": val_metrics["accuracy"], "f1": val_metrics["f1"], }) # Save the best checkpoint if val_metrics["f1"] > best_val_f1: best_val_f1 = val_metrics["f1"] best_val_acc = val_metrics["accuracy"] self.save() logger.info(" → New best model saved (val_f1=%.4f)", best_val_f1) # ── Final evaluation on test set ────────────────────────── logger.info("Loading best checkpoint for test evaluation...") self.load() # load() always places the model on CPU -- move it back to the # training device or GPU runs crash here with a device mismatch # against the CUDA-resident evaluation batches. self._model.to(device) test_metrics = self._evaluate(test_loader, device) from sklearn.metrics import classification_report, confusion_matrix test_report = classification_report( test_metrics["labels"], test_metrics["preds"], target_names = self._labels, labels = list(range(len(self._labels))), output_dict = True, zero_division = 0, ) per_class = { lbl: { "precision": test_report[lbl]["precision"], "recall": test_report[lbl]["recall"], "f1": test_report[lbl]["f1-score"], "support": test_report[lbl]["support"], } for lbl in self._labels } for lbl, m in per_class.items(): logger.info( " %-10s precision=%.3f recall=%.3f f1=%.3f (n=%d)", lbl, m["precision"], m["recall"], m["f1"], m["support"], ) cm = confusion_matrix( test_metrics["labels"], test_metrics["preds"], labels = list(range(len(self._labels))), ).tolist() results = { "val_accuracy": best_val_acc, "val_f1": best_val_f1, "test_accuracy": test_metrics["accuracy"], "test_f1": test_metrics["f1"], "per_class": per_class, "confusion_matrix": cm, "history": metrics_history, } logger.info( "Training complete — test_acc=%.4f test_f1=%.4f", results["test_accuracy"], results["test_f1"], ) self._save_metrics(results) return results def _save_metrics(self, results: dict) -> None: """Persist training metrics to output_dir/training_metrics.json. Read by the dashboard's Model Metrics page (via the API, not directly -- the dashboard process never has this file locally). Args: results: The dict returned by train(). """ self._output_dir.mkdir(parents=True, exist_ok=True) metrics_path = self._output_dir / "training_metrics.json" metrics_path.write_text(json.dumps(results, indent=2)) logger.info("Training metrics saved to %s", metrics_path) def _evaluate(self, data_loader, device) -> dict[str, float]: """Run inference on a DataLoader and return accuracy and F1. Args: data_loader: PyTorch DataLoader with labelled examples. device: torch.device for inference. Returns: Dict with ``"accuracy"`` and ``"f1"`` keys. """ import torch from sklearn.metrics import accuracy_score, f1_score self._model.eval() all_preds = [] all_labels = [] with torch.no_grad(): for batch in data_loader: labels = batch.pop("labels").numpy() batch = {k: v.to(device) for k, v in batch.items()} logits = self._model(**batch).logits preds = logits.argmax(dim=-1).cpu().numpy() all_preds.extend(preds) all_labels.extend(labels) return { "accuracy": accuracy_score(all_labels, all_preds), "f1": f1_score( all_labels, all_preds, average="weighted", zero_division=0, ), "labels": all_labels, "preds": all_preds, } # ── Inference ───────────────────────────────────────────────── def predict(self, text: str) -> ClassificationResult: """Classify a single clinical note. Args: text: Cleaned clinical note text. Pass raw text through :func:`src.utils.text_utils.prepare_for_inference` before calling this method. Returns: :class:`ClassificationResult` with the predicted label, confidence, and full probability distribution. Raises: RuntimeError: If the model has not been loaded yet. """ if self._model is None or self._tokenizer is None: raise RuntimeError( "Model not loaded. Call load() before predict()." ) import torch device = next(self._model.parameters()).device encoding = self._tokenizer( text, max_length = TrainingConfig.max_length, truncation = True, padding = "max_length", return_tensors = "pt", ) self._model.eval() with torch.no_grad(): logits = self._model( input_ids = encoding["input_ids"].to(device), attention_mask = encoding["attention_mask"].to(device), ).logits probs = torch.softmax(logits, dim=-1).squeeze().cpu().numpy() pred_idx = int(np.argmax(probs)) pred_label = self._id2label[pred_idx] confidence = float(probs[pred_idx]) prob_dict = { self._id2label[i]: float(p) for i, p in enumerate(probs) } return ClassificationResult( label = pred_label, confidence = confidence, probabilities = prob_dict, task = self._task, ) def predict_batch(self, texts: list[str]) -> list[ClassificationResult]: """Classify multiple clinical notes. More efficient than calling predict() in a loop because the tokeniser and model run in batches. Args: texts: List of cleaned clinical note strings. Returns: List of :class:`ClassificationResult` objects, one per input. """ if not texts: return [] # For simplicity, loop — a production implementation would # use DataLoader batching here. return [self.predict(t) for t in texts]