Check Protein Embeddings Before Training with Python

Protein embeddings can have the right number of rows and still be paired with the wrong labels. Sorting a FASTA file, merging extraction batches, or filtering failed sequences can change row order. A classifier may then train on plausible numerical inputs attached to the wrong proteins.

This tutorial checks a saved matrix before training: one unique sample ID per vector, a documented feature dimension, finite floating-point values, and exact agreement with an independent ID list in label-table order. It follows the ESM-2 embedding tutorial and the protein label join workflow. Basic Python and NumPy knowledge is enough; the demonstration needs no GPU or model download.

1. Define the embedding file contract

Use an NPZ archive with exactly two named arrays. This is a small interchange format for this tutorial; it is not a native ESM export format.

IDs must be nonempty, unique, and contain no whitespace. They are compared exactly: 001 and 1 remain different, as do upper- and lowercase letters. This tutorial requires one row per sample; repeated measurements need distinct sample IDs and a separate mapping to their parent sequence.

Create ids.txt from the final label table, independently of the embedding archive. Copying the archive IDs into the manifest would make the ordering check circular. An ID in the archive must come from the sequence associated with that vector during extraction.

Per-residue representations usually need a separate pooling step before they can fit this contract. The official ESM extraction script offers both per-token and mean representations. Its mean calculation selects residue positions before averaging; a tensor with a residue axis should not be flattened merely to pass the shape check. Record the model checkpoint, revision, representation layer, residue mask, pooling method, and truncation policy alongside the data.

2. Download the example and install NumPy

Download audit_embeddings.py and make_demo.py into a new working directory. Use Python 3.12 or newer for the pinned NumPy version below. The scripts were tested with Python 3.14.7 and NumPy 2.5.3.

python3 -m venv .venv
source .venv/bin/activate
python -m pip install "numpy==2.5.3"
python audit_embeddings.py --help

The fixture generator below writes a new demo-embeddings directory and refuses to overwrite an existing one. All vectors are invented numbers: three samples, four features each. They are not ESM outputs or biological measurements.

"""Create small synthetic fixtures without extracting real protein embeddings."""
from pathlib import Path

import numpy as np


def main():
    root = Path("demo-embeddings")
    root.mkdir()  # Refuse to overwrite a previous demo directory.
    ids = np.array(["protein_A", "protein_B", "protein_C"], dtype=np.str_)
    vectors = np.array([[0.1, 0.2, 0.3, 0.4],
                        [0.4, 0.3, 0.2, 0.1],
                        [-0.1, 0.0, 0.2, 0.5]], dtype=np.float32)
    (root / "ids.txt").write_text("\n".join(ids.tolist()) + "\n", encoding="utf-8")
    cases = {"good": (ids, vectors),
             "reordered": (ids[::-1], vectors[::-1]),
             "duplicate": (np.array(["protein_A", "protein_A", "protein_C"]), vectors)}
    bad = vectors.copy()
    bad[1, 2] = np.nan
    cases["nan"] = (ids, bad)
    for name, (sample_ids, embedding) in cases.items():
        np.savez(root / f"{name}.npz", sample_id=sample_ids, embedding=embedding,
                 allow_pickle=False)
    print("Created demo-embeddings with four synthetic archives and ids.txt")


if __name__ == "__main__":
    main()

numpy.savez stores arrays under the keyword names you supply. Explicit Unicode IDs and floating-point vectors avoid object arrays in these files.

3. Check a valid archive

python make_demo.py
python audit_embeddings.py demo-embeddings/good.npz demo-embeddings/ids.txt --dim 4

Expected output from the checker:

{"status": "ok", "samples": 3, "dimensions": 4, "dtype": "float32", "zero_rows": 0}

Use the feature count from your real extraction configuration in place of 4. Do not infer the expected dimension from the file being checked. Two different checkpoints or pooling workflows can produce the same feature count, so dimension agreement alone cannot establish a common representation.

Here is the complete checker, identical to the download:

"""Check a pooled protein embedding archive against an ordered ID manifest."""
import argparse
import json
from pathlib import Path
import sys
from zipfile import BadZipFile

import numpy as np


def check_ids(ids, source):
    if not ids:
        raise ValueError(f"{source}: no sample IDs")
    seen = set()
    for row, value in enumerate(ids):
        if not value or any(char.isspace() for char in value):
            raise ValueError(f"{source}: empty ID or whitespace at row {row}")
        if value in seen:
            raise ValueError(f"{source}: duplicate ID {value!r}")
        seen.add(value)


def audit(path, ids_path, expected_dim):
    if expected_dim < 1:
        raise ValueError("expected dimension must be positive")
    expected = Path(ids_path).read_text(encoding="utf-8-sig").splitlines()
    check_ids(expected, "manifest")
    archive = np.load(path, allow_pickle=False)
    if not isinstance(archive, np.lib.npyio.NpzFile):
        raise ValueError("expected an NPZ archive")
    with archive:
        if len(archive.files) != 2 or set(archive.files) != {"sample_id", "embedding"}:
            raise ValueError("archive must contain only sample_id and embedding")
        ids = archive["sample_id"]
        vectors = archive["embedding"]
    if ids.ndim != 1 or ids.dtype.kind != "U":
        raise ValueError("sample_id must be a one-dimensional Unicode array")
    names = ids.tolist()
    check_ids(names, "archive")
    if vectors.ndim != 2 or vectors.shape != (len(names), expected_dim):
        raise ValueError(f"expected embedding shape ({len(names)}, {expected_dim})")
    if vectors.dtype.kind != "f":
        raise ValueError("embedding must have a floating-point dtype")
    if names != expected:
        if set(names) == set(expected):
            raise ValueError("ID order differs; join by sample_id before export")
        raise ValueError("archive IDs differ from manifest IDs")
    zero_rows = 0
    # Arrays are loaded into RAM; chunking limits temporary boolean arrays.
    for start in range(0, len(names), 1024):
        block = vectors[start:start + 1024]
        finite = np.isfinite(block)
        if not finite.all():
            row, column = np.argwhere(~finite)[0]
            raise ValueError(f"non-finite value at row {start + int(row)}, column {int(column)}")
        zero_rows += int(np.all(block == 0, axis=1).sum())
    return {"status": "ok", "samples": len(names), "dimensions": expected_dim,
            "dtype": str(vectors.dtype), "zero_rows": zero_rows}


def main():
    parser = argparse.ArgumentParser(description=__doc__)
    parser.add_argument("archive", type=Path)
    parser.add_argument("ids", type=Path, help="one ID per line, in downstream label order")
    parser.add_argument("--dim", type=int, required=True, help="expected pooled feature count")
    args = parser.parse_args()
    try:
        result = audit(args.archive, args.ids, args.dim)
    except (OSError, ValueError, EOFError, BadZipFile) as error:
        print(f"ERROR: {error}", file=sys.stderr)
        return 1
    print(json.dumps(result, ensure_ascii=False))
    return 0


if __name__ == "__main__":
    raise SystemExit(main())

The checker uses numpy.load with allow_pickle=False and closes the NPZ archive after reading. Object arrays are rejected. If an older export needs pickle to load its IDs, regenerate the export from its trusted source with a Unicode ID array.

numpy.isfinite rejects NaN and both infinities. All-zero rows are counted in the report rather than rejected automatically. Investigate them before cosine normalization, which requires a nonzero vector. Whether a zero row is an extraction error depends on the documented representation.

4. Reproduce three failures

Run these commands separately; each should print an error on standard error and exit with status 1:

python audit_embeddings.py demo-embeddings/reordered.npz demo-embeddings/ids.txt --dim 4
python audit_embeddings.py demo-embeddings/duplicate.npz demo-embeddings/ids.txt --dim 4
python audit_embeddings.py demo-embeddings/nan.npz demo-embeddings/ids.txt --dim 4

ERROR: ID order differs; join by sample_id before export
ERROR: archive: duplicate ID 'protein_A'
ERROR: non-finite value at row 1, column 2

The reordered fixture reverses IDs and vectors together. Its protein-to-vector mapping is preserved, but its row order differs from the downstream label order. Matching the sets of IDs would miss this problem. All reported array positions use zero-based indexing.

The checker reads files without modifying them and emits a success report only after all checks pass. In an automated pipeline, require a successful exit before starting the training job:

python audit_embeddings.py demo-embeddings/good.npz demo-embeddings/ids.txt --dim 4 && echo "Embedding checks passed"

5. Fix the cause of each failure

Never sort only the ID array, only the labels, or only the vectors. Apply an explicit join or the same permutation to the associated records. The label join tutorial explains how to check key membership and preserve sample identity.

6. Preserve provenance and plan for memory

This checker loads both arrays into RAM. It checks numerical values in blocks of 1,024 rows to limit temporary boolean arrays, but it does not stream the NPZ matrix. A float32 matrix with 100,000 samples and 1,280 features occupies about 488 MiB before IDs, Python objects, and temporary arrays. Compressed archive size is not its loaded memory requirement.

For larger datasets, consider a separate .npy matrix with memory mapping and a versioned ID manifest. That layout needs its own checker; this one deliberately accepts NPZ only. NumPy documents memory mapping in numpy.load.

A passing result cannot detect a vector exported under the wrong ID upstream, incorrect residue pooling, mixed model revisions, or train/test leakage. Keep sequence checksums and extraction metadata, preserve the dataset manifest, and use the same grouped split when comparing embeddings with a composition baseline. Fitted scaling, PCA, and predictor selection still belong inside the appropriate training and validation partitions.

Summary

Validate sample identity, row order, dimensions, dtype, and finite values before fitting a protein predictor. Preserve the extraction metadata and sequence mapping so the saved numbers can be traced back to their source. Once those checks pass, evaluate the representation on a split that matches the biological generalization question.

The cover is conceptual artwork, not an experimental protein structure.