Metadata-Version: 2.4
Name: plmsommelier
Version: 1.0.0
Summary: Pick the best layer of a protein language model for your dataset, and get back a truncated model.
Author: Roman Joeres
Author-email: Ilya Senatorov <il.senatorov@gmail.com>
License-Expression: MIT
Project-URL: Homepage, https://github.com/kalininalab/PLMSommelier
Project-URL: Repository, https://github.com/kalininalab/PLMSommelier
Project-URL: Issues, https://github.com/kalininalab/PLMSommelier/issues
Project-URL: Changelog, https://github.com/kalininalab/PLMSommelier/blob/main/CHANGELOG.md
Project-URL: Paper, https://arxiv.org/abs/2608.12090
Keywords: protein language model,layer selection,probing,esm,bioinformatics
Classifier: Development Status :: 4 - Beta
Classifier: Intended Audience :: Science/Research
Classifier: Programming Language :: Python :: 3
Classifier: Programming Language :: Python :: 3.10
Classifier: Programming Language :: Python :: 3.11
Classifier: Programming Language :: Python :: 3.12
Classifier: Programming Language :: Python :: 3.13
Classifier: Topic :: Scientific/Engineering :: Bio-Informatics
Classifier: Topic :: Scientific/Engineering :: Artificial Intelligence
Classifier: Typing :: Typed
Classifier: Operating System :: OS Independent
Requires-Python: >=3.10
Description-Content-Type: text/markdown
License-File: LICENSE
Requires-Dist: cyclopts>=3
Requires-Dist: numpy>=2.1
Requires-Dist: pandas>=2.2.3
Requires-Dist: scikit-learn>=1.5.2
Requires-Dist: torch>=2.5
Requires-Dist: transformers<6,>=5.9
Requires-Dist: sentencepiece>=0.2.1
Requires-Dist: protobuf>=3.20
Requires-Dist: tqdm
Provides-Extra: multimolecule
Requires-Dist: multimolecule>=0.2.1; extra == "multimolecule"
Provides-Extra: dev
Requires-Dist: pytest>=7; extra == "dev"
Requires-Dist: pytest-cov; extra == "dev"
Requires-Dist: ruff; extra == "dev"
Requires-Dist: pre-commit; extra == "dev"
Requires-Dist: build; extra == "dev"
Requires-Dist: twine; extra == "dev"
Dynamic: license-file

# PLMSommelier

[![PyPI](https://img.shields.io/pypi/v/plmsommelier)](https://pypi.org/project/plmsommelier/)
[![Python](https://img.shields.io/pypi/pyversions/plmsommelier)](https://pypi.org/project/plmsommelier/)
[![CI](https://github.com/kalininalab/PLMSommelier/actions/workflows/ci.yml/badge.svg)](https://github.com/kalininalab/PLMSommelier/actions/workflows/ci.yml)
[![License: MIT](https://img.shields.io/badge/license-MIT-blue.svg)](LICENSE)

**The last layer of a protein language model is almost never the best one. This
tool finds the layer that is, and hands you back a truncated model.**

Implements the tool described in [*Task- and dataset-specific information
in protein language models*](https://arxiv.org/abs/2608.12090), which probed 13 PLMs across 15
downstream tasks and found the deepest layer won in only **17.9%** of cases.

## Install

```bash
uv tool install plmsommelier   # isolated, no venv to manage yourself
# or: pipx install plmsommelier
# or: pip install plmsommelier   # into a virtualenv
```

<details>
<summary>CPU vs. GPU torch (Linux)</summary>

PLMSommelier itself is CUDA-agnostic -- it never links CUDA directly, so any
torch build works. But on Linux, plain `pip install plmsommelier` resolves
torch's *default* index, which is a CUDA build pulling several GB of `nvidia-*`
packages even on a machine with no GPU. Install torch yourself first if that's
not what you want:

```bash
# CPU-only (much smaller download; also what you want on a machine with no GPU)
pip install torch --index-url https://download.pytorch.org/whl/cpu
pip install plmsommelier

# GPU: pick the wheel matching your driver, then install as usual
pip install torch --index-url https://download.pytorch.org/whl/cu126   # CUDA 12
pip install torch --index-url https://download.pytorch.org/whl/cu130   # CUDA 13
pip install plmsommelier
```

macOS wheels are CPU-only already (Apple Silicon uses `mps` automatically at
runtime instead), so the plain `pip install plmsommelier` is correct there --
no separate torch step needed. If a `cuXXX` install fails or a model refuses
to use your GPU, check that your driver actually supports the CUDA version
that wheel targets; a driver too old for the chosen wheel is the usual cause.

</details>

## Quickstart

`examples/fluorescence_sample.csv` ships in the repo (a 500-row subsample of
the TAPE fluorescence benchmark -- see `examples/README.md`), so this runs in
about a minute on CPU with no data of your own needed:

```bash
plmsommelier suggest examples/fluorescence_sample.csv facebook/esm2_t6_8M_UR50D \
    --task regression --out ./my-esm-truncated
```

`data` and `model` are positional (`DATA MODEL`); `--data`/`--model` work
identically if you prefer named flags.

```
model    facebook/esm2_t6_8M_UR50D
dataset  fluorescence_sample  (regression, pearson, knn probe)
data     400 train / 100 val

best layer      0 of 6 (0% depth)   pearson = 0.4238
last layer      6               pearson = 0.2369
gain over last  +78.9%
seed agreement  100% (layers chosen across seeds: [0])
confidence      moderate   (plateau agreement 100%, separation 1.51x, seed spread 0% of depth)

layer performance (pearson, +/- 1 sd across 5 seeds):
    0 +0.4238 ########################################-----   +/-0.0537  <- best
    1 +0.3764 ####################################---------   +/-0.1012
    2 +0.3770 ####################################-----   +/-0.0634
    3 +0.1572 ###############------   +/-0.0736
    4 +0.0432 ####----   +/-0.0519
    5 +0.2014 ###################-------   +/-0.0824
    6 +0.2369 #######################------   +/-0.0710  (last)
```

(Layer `0` is the token embedding output before any transformer block --
a real, if unglamorous, answer on this tiny 6-layer model and 500-row sample;
larger models and datasets typically pick somewhere mid-stack. See the paper
for what "best layer" looks like at scale.)

`./my-esm-truncated` is a normal HuggingFace model directory holding only the
chosen number of blocks. It loads anywhere the original did, runs faster, and
scores at least as well on your task:

```python
from transformers import AutoModel, AutoTokenizer

model = AutoModel.from_pretrained("./my-esm-truncated")
tok = AutoTokenizer.from_pretrained("./my-esm-truncated")
```

T5-encoder checkpoints (Ankh, ProtT5, ProstT5, ...) load via
`AutoModelForTextEncoding` instead of `AutoModel`. You don't need to remember
which is which: the exact loading snippet for whatever model you just
truncated is always written into the generated `{out}/README.md`.

Any HuggingFace protein language model works — just pass its id to `model`.
Checkpoints that ship custom modeling code need `--trust-remote-code`. See
[Extending to a custom PLM](#extending-to-a-custom-plm) below for models that
need more than that.

## Python API

The CLI is a thin wrapper around five importable functions
(`plmsommelier/__init__.py`), wired together in exactly this order by
`cli.py::suggest_layer` if you'd rather read one function than five:

```python
from plmsommelier import embed_layers, load_dataset, load_model, save_truncated, select_layer

plm = load_model("facebook/esm2_t6_8M_UR50D")
ds = load_dataset("examples/fluorescence_sample.csv", task="regression")

train_embeddings = embed_layers(plm, ds.train_seqs)
val_embeddings = embed_layers(plm, ds.val_seqs)

result = select_layer(ds, train_embeddings, val_embeddings, model_name=plm.model_id)
print(result.summary())  # the block printed by the CLI
print(result.curve_plot())  # the ASCII per-layer curve
result.to_dict()  # what --json writes out

save_truncated(plm, result, "./my-esm-truncated")
```

`Result` (`plmsommelier.select.Result`) carries the full per-layer curve
(`result.curve`), the per-seed scores behind the stability check
(`result.seed_curves`, `result.layer_sigma`), and the `confidence` verdict --
see [Confidence and what to do about it](#confidence-and-what-to-do-about-it).

## Input format

A CSV with `sequence` and a label column:

| column | required | notes |
|---|---|---|
| `sequence` | yes | amino acid sequence |
| `label` / `labels` / `Y` | yes | the target (checked in that order, or pass `--label-col`) |
| `ID` | no | not used by `plmsommelier`; keep it for your own bookkeeping if you like |
| `split` | no | `train` / `valid` (or `val`); generated if absent |

`--task` (`regression`, `classification` or `multi-label`) is inferred from
the label column when that's unambiguous: multiple label columns ->
multi-label, non-numeric or exactly two distinct values -> classification, a
float dtype -> regression. A whole-number column with more than two values
could be class codes or an integer-valued regression target, and inference
refuses to guess there -- pass `--task` explicitly.

`--max-seq-len` (default 2000 residues) drops longer sequences before
splitting or subsampling, so a single outlier protein can't end up alone in
an unbounded batch -- attention memory is quadratic in length. Pass `0` to
disable it.

Device is auto-detected -- CUDA, else Apple Silicon's `mps`, else CPU -- and
can be forced with `--device cpu` / `--device cuda` / `--device mps`. A
visible CUDA device only counts if the installed torch build actually ships
kernels for it: an older card (e.g. anything pre-Turing) against a build that
dropped support for it would otherwise crash on the first forward pass with
`no kernel image is available for execution on the device` instead of just
running on the CPU. Auto-detect falls through with a warning in that case;
pass `--device cuda` explicitly to force it anyway.

## CLI reference

`plmsommelier suggest --help` is always authoritative -- every flag is
generated directly from `suggest_layer`'s Python signature and docstring, so
it can't drift out of sync with this table.

| flag | default | what it does |
|---|---|---|
| `data` / `--data` | *(required)* | CSV path (positional or named) |
| `model` / `--model` | *(required)* | HuggingFace model id or local path |
| `--task` | inferred | `regression` / `classification` / `multi-label` |
| `--label-col` | auto-detected | override label column detection |
| `--probe` | `knn` | `knn` or `lr` (linear probe) |
| `--k` | `10` | neighbours for the kNN probe |
| `--sample` | `5000` | row budget across train+val; `>=1` absolute, `(0,1)` a fraction |
| `--max-seq-len` | `2000` | drop longer sequences before splitting/subsampling; `0` disables |
| `--n-seeds` | `5` | resampling repeats for the confidence check; `0` disables it |
| `--seed` | `42` | random seed |
| `--tolerance` | `0.02` | relative score band treated as a tied "plateau" |
| `--device` | auto | `cpu` / `cuda` / `mps` |
| `--cache-dir` | HF default | HuggingFace cache directory |
| `--trust-remote-code` | `False` | needed for checkpoints with custom modeling code |
| `--out` | *(none)* | write the truncated model here |
| `--progress` / `--no-progress` | on | embedding progress bar |

Flags shared by every command, handled by the top-level launcher:

| flag | what it does |
|---|---|
| `--json PATH` | also write the result as JSON |
| `--quiet` | suppress the rendered summary/curve output |
| `--traceback` | let exceptions propagate instead of printing `error: ...` and exiting 1 |

## Extending to a custom PLM

There is no model registry -- `model` takes any HuggingFace id or local path,
and works with zero configuration for most architectures. Reach for the
options below only as far as your checkpoint actually needs, in order of
how rare each one is:

1. **Nothing.** `load_model` (`plmsommelier/model.py`) is architecture-agnostic:
   it finds the transformer blocks and final norm by attribute name, detects
   the sequence-length ceiling from the tokenizer/config/causal-mask buffer,
   and picks the right `Auto*` loading class. Most encoder PLMs on the Hub
   just work.

2. **`--trust-remote-code`** for checkpoints that ship custom modeling code
   in their own repo (an `auto_map` in `config.json`). Same flag `transformers`
   itself uses.

3. **A pretraining-convention entry in `_QUIRKS`** (`plmsommelier/model.py`),
   keyed on `config.model_type`, for models whose config doesn't record how
   they were actually pretrained -- HuggingFace has no field for "residues
   are space-separated tokens", for instance. Supported keys:

   | key | effect |
   |---|---|
   | `space_join` | join residues with spaces before tokenizing (ProtTrans-style) |
   | `residue_map` | a `str.translate` table applied to the sequence first (e.g. rare residues -> `X`) |
   | `prefix_text` | text prepended before tokenizing (e.g. a mode/direction token) |
   | `dtype` | force this torch dtype instead of the auto-detected default |
   | `encoder_only` | force `AutoModelForTextEncoding` (`True`) vs. `AutoModel`/`AutoModelForCausalLM` (`False`), overriding the `is_encoder_decoder` config check |

   T5/BERT/ALBERT (the ProtTrans family) are the worked example already in
   the table:

   ```python
   "t5": dict(
       space_join=True,
       residue_map=str.maketrans("UZOB", "XXXX"),
       dtype=torch.bfloat16,   # T5 activations overflow in fp16
       encoder_only=True,      # loads T5EncoderModel; the decoder is never built
   ),
   ```

   Add a family the same way: an entry keyed on its `model_type`
   (`AutoConfig.from_pretrained(your_model).model_type` tells you the key),
   with only the keys it actually needs.

4. **A `_REGISTRY_PACKAGES` entry** for architectures that live in a
   third-party package rather than in `transformers` itself -- these
   checkpoints carry no `auto_map`, so `--trust-remote-code` can't reach them;
   the package must be imported first so it registers itself with the `Auto*`
   classes. `multimolecule/proteinbert` is the current example: install
   `plmsommelier[multimolecule]`, and `model.py` imports `multimolecule`
   before loading whenever `config.model_type == "proteinbert"`. Add a new
   package the same way: `"model_type": ("import_name", "pip_extra_name")`,
   plus an extra in `pyproject.toml` if it isn't already installed alongside
   `plmsommelier`.

5. **No fork needed for a one-off.** `_QUIRKS` and `_REGISTRY_PACKAGES` are
   plain module-level dicts -- from your own script, `import plmsommelier.model
   as m; m._QUIRKS["your_model_type"] = {...}` before calling `load_model`
   works without touching this repo. Treat that as a private escape hatch,
   though (the leading underscore is deliberate) -- if it's a real,
   reusable family, a PR adding it to the table is the better home for it.

Whatever route you take, verify a new family the way `tests/test_model.py`
verifies the built-in ones: the truncation round-trip (`save_truncated` then
reload reproduces the same layer's output), sequence-length detection not
silently truncating valid input, and the final-norm invariant --
see [CONTRIBUTING.md](CONTRIBUTING.md)'s invariants list for what each of
those actually guards against.

### Tested models

Checkpoints below were loaded and probed end-to-end (`load_model` +
`embed_layers`) against the real weights, not just read from the code — either
by the repo's own real-weight test fixtures (`tests/test_model.py`) or by a
manual smoke test. Anything not listed still has a good chance of working —
`load_model` is architecture-agnostic — it just hasn't been verified here yet.

| Model | HuggingFace ID | Status |
|---|---|---|
| ESM-2 | `facebook/esm2_t6_8M_UR50D` (+ larger) | ✅ confirmed (test suite) |
| ESM-1b | `facebook/esm1b_t33_650M_UR50S` | ✅ confirmed (manual) |
| IgBert | `Exscientia/IgBert` | ✅ confirmed (manual) |
| ProGen2 | `hugohrban/progen2-small` (+ larger) | ✅ confirmed (test suite) |
| proteinbert (multimolecule) | `multimolecule/proteinbert` | ✅ confirmed (test suite) |
| RITA | `lightonai/RITA_s` (+ larger) | ❌ fails — tokenizer ships with no pad, eos, or unk token; `load_model` refuses to guess one |

ProtBert, ProtAlbert, ProtT5, Ankh, ProstT5, and ProtGPT2 are handled by name
in `_QUIRKS` but have no real-weight test coverage yet — treat them as
likely-to-work, not confirmed.

## Confidence and what to do about it

Every run redraws the training data (and, where there's enough of it, the
validation data) `--n-seeds` times and re-scores every layer, to check
whether the chosen layer survives resampling. That feeds a `confidence`
verdict -- `high`, `moderate`, or `low` -- printed as part of the summary and
recorded in `--json` output. `unmeasured` means the check wasn't run at all
(`--n-seeds 0`), which is different from `low`: `low` means "we checked, and
the pick isn't stable"; `unmeasured` means "we didn't check".

The verdict is the *worst* of three independent signals, not an average, so
one strong number can't paper over another weak one:

- **seed agreement** -- how often resampling lands back inside the same
  plateau (the tolerance-band region of tied-best layers), not just on the
  exact same layer. Two seeds picking neighbouring, statistically tied
  layers count as agreement, not disagreement.
- **peak margin** -- how many (paired) standard deviations separate the
  plateau from the best layer outside it, across resamples. This is a
  resampling-stability statistic, not a hypothesis test: seeds are
  overlapping subsamples of one dataset, not independent draws, so no
  p-value is implied.
- **seed spread** -- how far, on average, resampled picks land from the
  chosen layer, as a fraction of the network's depth.

A `low` verdict prints a warning block to stderr (not stdout, so `--json`
and piped output stay clean) with concrete, run-specific hints. **The most
important thing to know before reading them:** raising `--n-seeds` measures
confidence more precisely -- it does not raise it. Only one of the following
actually reduces variance:

1. **Raise `--sample`**, if the hint says there's headroom (the tool tracks
   how many rows were available versus used). Score variance falls roughly
   as `1/N`, so ~3x the data tightens the estimate by roughly `1.7x`. If
   you're already using every row in the file, this can't help -- the hint
   says so explicitly rather than telling you to raise a flag that won't do
   anything.
2. **Try `--probe lr`** (or raise `--k` for the kNN probe). A linear probe has
   much lower variance than kNN at typical sample sizes; two probes agreeing
   on a layer is stronger evidence than either one's confidence number alone.
3. **Check the validation set.** Below roughly 50 val rows, none of these
   statistics are trustworthy -- `--sample` and the train/val split both
   drive this.
4. **Check the data** for heavy class imbalance, a near-constant label
   column, or a `split` column that leaves a tiny validation side.
5. **Accept the plateau.** If the curve is genuinely flat, low confidence is
   the *correct* answer, not a problem to fix -- the tool already returns the
   shallowest layer in the tied region, which is the smallest and fastest
   model at equal performance.

## Development

See [CONTRIBUTING.md](CONTRIBUTING.md) for the dev install and workflow, and
[RELEASING.md](RELEASING.md) for how a release is cut and published.

```bash
pip install -e ".[dev]"
pytest tests -q                      # full suite
pytest tests -q -m "not weights"     # no model downloads needed
```

## Citation

```bibtex
@article{joeres2026taskdataset,
  title  = {Task- and dataset-specific information in protein language models},
  author = {Joeres, Roman and Senatorov, Ilya and Kolchina, Anastasia and
            Klakow, Dietrich and Kalinina, Olga V.},
  journal = {arXiv preprint arXiv:2608.12090},
  year   = {2026}
}
```

See [CITATION.cff](CITATION.cff) for the software-citation form.

## License

MIT — see [LICENSE](LICENSE).
