models/classifier.py

#!/usr/bin/env python3
"""DistilBERT + a linear classification head, in either ``lora`` or ``full`` mode.

Both modes use the identical architecture. The only difference is which
parameters carry ``requires_grad=True`` — which is precisely the comparison this
project is about.

Run ``python -m models.classifier`` to print both modes side by side and see the
parameter counts without training anything.
"""

from __future__ import annotations

import json
import logging
from pathlib import Path
from typing import Any, Literal

import torch
import torch.nn as nn
from transformers import AutoConfig, AutoModel

try:
    from models.lora import inject_lora, mark_only_lora_as_trainable, parameter_report
except ImportError:  # pragma: no cover - fallback for direct script execution
    from lora import inject_lora, mark_only_lora_as_trainable, parameter_report  # type: ignore

log = logging.getLogger("classifier")

Mode = Literal["lora", "full"]

DEFAULT_MODEL = "distilbert-base-uncased"

#: DistilBERT names its attention projections q_lin / k_lin / v_lin / out_lin.
#: Query and value only — see models/lora.py for why.
LORA_TARGETS: tuple[str, ...] = ("q_lin", "v_lin")


class TextClassifier(nn.Module):
    """Classifies a short text into one of ``num_labels`` classes.

    Architecture::

        input_ids (B, L)
          -> DistilBERT encoder            -> (B, L, 768)
          -> take position 0, the [CLS]    -> (B, 768)
          -> dropout
          -> Linear(768, num_labels)       -> (B, num_labels)   logits

    A deliberate simplification: HuggingFace's own
    ``DistilBertForSequenceClassification`` inserts a 768x768 ``pre_classifier``
    layer before the final one. That adds ~590k trainable parameters to the head —
    which would be **four times larger than all the LoRA adapters combined** and
    would completely distort the headline "1% of parameters" comparison. A single
    linear head (59k params at 77 classes) keeps the measurement honest.
    """

    def __init__(
        self,
        model_name: str = DEFAULT_MODEL,
        num_labels: int = 77,
        mode: Mode = "lora",
        lora_r: int = 8,
        lora_alpha: int = 16,
        lora_dropout: float = 0.0,
        head_dropout: float = 0.1,
        class_weights: torch.Tensor | None = None,
        label_names: list[str] | None = None,
    ) -> None:
        """
        Args:
            model_name: Pretrained checkpoint to load.
            num_labels: Number of classes. 77 for banking77.
            mode: ``"lora"`` freezes the base and adapts query/value projections.
                ``"full"`` trains every parameter.
            lora_r: LoRA rank (ignored in full mode).
            lora_alpha: LoRA scaling numerator (ignored in full mode).
            lora_dropout: Dropout inside the LoRA path (ignored in full mode).
            head_dropout: Dropout before the classification head.
            class_weights: Optional ``(num_labels,)`` tensor for
                ``CrossEntropyLoss``. See ``data.dataset.class_weights``.
            label_names: Ordered class names. Stored in the checkpoint so the
                demo and evaluator resolve predictions to names without needing
                the dataset — the single most likely source of silently
                mislabelled output if left to be re-derived.
        """
        super().__init__()
        if mode not in ("lora", "full"):
            raise ValueError(f"mode must be 'lora' or 'full', got {mode!r}")

        self.model_name = model_name
        self.num_labels = num_labels
        self.mode = mode
        self.label_names = list(label_names) if label_names else None
        self.lora_r = lora_r
        self.lora_alpha = lora_alpha

        config = AutoConfig.from_pretrained(model_name)
        self.encoder = AutoModel.from_pretrained(model_name)
        hidden_size = config.dim if hasattr(config, "dim") else config.hidden_size

        self.dropout = nn.Dropout(head_dropout)
        self.classifier = nn.Linear(hidden_size, num_labels)

        self.n_adapted = 0
        if mode == "lora":
            self.n_adapted = inject_lora(
                self.encoder,
                target_names=LORA_TARGETS,
                r=lora_r,
                alpha=lora_alpha,
                dropout=lora_dropout,
            )
            log.info("Injected LoRA (r=%d) into %d projections", lora_r, self.n_adapted)
            mark_only_lora_as_trainable(self, also_train=("classifier",))

        # register_buffer, not a plain attribute: buffers move with .to(device)
        # but are not parameters, so the optimizer ignores them.
        if class_weights is not None:
            self.register_buffer("class_weights", class_weights)
        else:
            self.class_weights = None

    def forward(
        self,
        input_ids: torch.Tensor,
        attention_mask: torch.Tensor,
        labels: torch.Tensor | None = None,
    ) -> dict[str, torch.Tensor | None]:
        """Run a forward pass and, if labels are given, compute the loss.

        Args:
            input_ids: ``(B, L)`` int64 token ids.
            attention_mask: ``(B, L)`` int64, 1 for real tokens and 0 for padding.
            labels: ``(B,)`` int64 class ids in ``[0, num_labels)``, or None at
                inference time.

        Returns:
            ``{"logits": (B, num_labels), "loss": scalar or None}``.
        """
        # (B, L) -> (B, L, 768). The attention mask keeps padded positions from
        # contributing to any token's representation.
        hidden = self.encoder(
            input_ids=input_ids, attention_mask=attention_mask
        ).last_hidden_state

        # DistilBERT has no pooler of its own, so pool manually: position 0 is
        # the [CLS] token, whose representation attends over the whole sequence
        # and is the conventional sentence-level summary.
        # (B, L, 768) -> (B, 768)
        pooled = hidden[:, 0]

        logits = self.classifier(self.dropout(pooled))  # (B, 768) -> (B, num_labels)

        loss = None
        if labels is not None:
            # CrossEntropyLoss takes RAW logits, not softmax probabilities — it
            # applies log_softmax internally. Feeding it softmax output is a
            # classic silent bug: it still trains, just badly.
            loss_fn = nn.CrossEntropyLoss(weight=self.class_weights)
            loss = loss_fn(logits, labels)

        return {"logits": logits, "loss": loss}

    @torch.no_grad()
    def predict(
        self, input_ids: torch.Tensor, attention_mask: torch.Tensor
    ) -> tuple[torch.Tensor, torch.Tensor]:
        """Return ``(predicted_ids (B,), probabilities (B, num_labels))``.

        ``@torch.no_grad()`` disables the autograd graph — at inference it is
        pure overhead in both memory and time.
        """
        self.eval()
        logits = self.forward(input_ids, attention_mask)["logits"]
        probs = torch.softmax(logits, dim=-1)
        return probs.argmax(dim=-1), probs

    def trainable_parameter_report(self) -> dict[str, Any]:
        """The project's key measurement: how much of the model is actually trained."""
        report = parameter_report(self)
        report.update(
            {
                "mode": self.mode,
                "lora_rank": self.lora_r if self.mode == "lora" else None,
                "adapted_layers": self.n_adapted,
            }
        )
        return report

    # ------------------------------------------------------------------
    # Checkpointing
    # ------------------------------------------------------------------

    def adapter_state_dict(self) -> dict[str, torch.Tensor]:
        """Only the tensors that actually changed during LoRA training.

        The frozen encoder is byte-identical to the public checkpoint, so saving
        it would be storing a copy of something already on the Hub. Keeping just
        the adapters and head takes the artifact from ~265MB to ~200KB, which is
        what makes the free Hugging Face Space deploy practical (see DEPLOY.md).
        """
        return {
            name: param.detach().cpu().clone()
            for name, param in self.state_dict().items()
            if "lora_A" in name or "lora_B" in name or name.startswith("classifier")
        }

    def save(self, path: str | Path) -> None:
        """Save a checkpoint: adapters only in LoRA mode, everything in full mode."""
        path = Path(path)
        path.parent.mkdir(parents=True, exist_ok=True)

        state = (
            self.adapter_state_dict()
            if self.mode == "lora"
            else {k: v.detach().cpu() for k, v in self.state_dict().items()}
        )
        torch.save(
            {
                "state_dict": state,
                "config": {
                    "model_name": self.model_name,
                    "num_labels": self.num_labels,
                    "mode": self.mode,
                    "lora_r": self.lora_r,
                    "lora_alpha": self.lora_alpha,
                    "label_names": self.label_names,
                },
            },
            path,
        )
        size_kb = path.stat().st_size / 1024
        log.info("Saved %s checkpoint to %s (%.0f KB)", self.mode, path, size_kb)

    @classmethod
    def load(cls, path: str | Path, device: torch.device | str = "cpu") -> "TextClassifier":
        """Rebuild a trained model from a checkpoint.

        The saved config records the mode and rank, so the architecture is
        reconstructed correctly before the weights are loaded in. Getting the
        rank wrong would produce a shape mismatch here rather than silently
        wrong predictions — which is the point of storing it.
        """
        # weights_only=False: the checkpoint holds a config dict alongside the
        # tensors. Only load checkpoints you produced yourself.
        payload = torch.load(path, map_location=device, weights_only=False)
        config = payload["config"]

        model = cls(
            model_name=config["model_name"],
            num_labels=config["num_labels"],
            mode=config["mode"],
            lora_r=config["lora_r"],
            lora_alpha=config["lora_alpha"],
            label_names=config.get("label_names"),
        )
        # strict=False for LoRA: the checkpoint intentionally omits the frozen
        # encoder, which came from the pretrained download instead.
        missing, unexpected = model.load_state_dict(
            payload["state_dict"], strict=(config["mode"] == "full")
        )
        if unexpected:
            raise ValueError(f"Checkpoint has unexpected keys: {unexpected[:5]}")
        if config["mode"] == "lora":
            mark_only_lora_as_trainable(model, also_train=("classifier",))

        return model.to(device)


def _demo() -> None:
    """Print both modes side by side. Run: ``python -m models.classifier``."""
    logging.basicConfig(level=logging.INFO, format="%(levelname)-7s %(message)s")
    for noisy in ("httpx", "urllib3", "filelock", "huggingface_hub"):
        logging.getLogger(noisy).setLevel(logging.WARNING)

    print("\n" + "=" * 68)
    print("  PARAMETER COMPARISON — no training, just architecture")
    print("=" * 68)

    try:
        from data.dataset import load_labels
        labels = load_labels()
    except Exception:  # noqa: BLE001 - demo works even before download.py runs
        labels = [f"class_{i}" for i in range(77)]
    print(f"  (using {len(labels)} classes)")

    rows = []
    for mode in ("full", "lora"):
        model = TextClassifier(mode=mode, num_labels=len(labels), label_names=labels)
        report = model.trainable_parameter_report()
        rows.append(report)
        label = f"{mode}" + (f" (r={report['lora_rank']})" if mode == "lora" else "")
        print(f"\n  {label}")
        print("  " + "-" * 56)
        print(f"    trainable : {report['trainable_params']:>12,}")
        print(f"    frozen    : {report['frozen_params']:>12,}")
        print(f"    total     : {report['total_params']:>12,}")
        print(f"    trainable : {report['trainable_pct']:>11.3f}% of total")
        if mode == "lora":
            print(f"    adapted   : {report['adapted_layers']} projections "
                  f"(query + value across 6 layers)")

    ratio = rows[0]["trainable_params"] / rows[1]["trainable_params"]
    print("\n  " + "-" * 56)
    print(f"  Full fine-tuning trains {ratio:,.0f}x more parameters than LoRA.")

    # Verify the zero-init claim rather than asserting it in a comment.
    lora_model = TextClassifier(mode="lora", num_labels=len(labels))
    first_adapter = next(
        m for m in lora_model.encoder.modules() if type(m).__name__ == "LoRALinear"
    )
    print(f"\n  Sanity check at initialization:")
    print(f"    lora_B all zeros : {bool((first_adapter.lora_B == 0).all())}  "
          f"(so B@A = 0, model == pretrained)")
    print(f"    lora_A all zeros : {bool((first_adapter.lora_A == 0).all())}  "
          f"(must be False, or no gradient could flow)")
    print("\n" + "=" * 68 + "\n")


if __name__ == "__main__":
    _demo()