训练前用 Python 检查蛋白质嵌入文件

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

"""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()

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

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

"""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())

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

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