bertose-affinose-training-code / CHECKPOINT_GUIDE.md
supanthadey1's picture
Add checkpoint selection and website integration guide
f21a3d6 verified
|
Raw History Blame Contribute Delete
5.01 kB

Choose and compare BERTose checkpoints

Use the shared code and workbench notebook. The code version is 2026-10-02; model assets remain under releases/2026-09-28/.

Sequence embeddings

Checkpoint Purpose
Variant 3 — default BERTose reference used in the manuscript comparison; start here for glycan embeddings.
Variants 1, 2, 4 and 5 Compare the effect of different pretraining objective weights on the same inputs.

All five variants used pretraining seed 42 and their own selected epoch. They are alternative weighting configurations, not independent seeds or models for different tasks. The reported comparison does not establish universally best weights.

Run from the code directory after following the quickstart. Keep the same WURCS strings, sample order and pooling for each variant:

from bertose_workbench import Encoder, cosine_similarity

wurcs = [
    "WURCS=2.0/1,1,0/[a2122h-1a_1-5]/1/",
    "WURCS=2.0/1,2,1/[a2122h-1a_1-5]/1-1/a4-b1",
]
reference_variant, comparison_variant = "3", "1"  # comparison: "1", "2", "4" or "5"
vectors, hashes = {}, {}
for variant in (reference_variant, comparison_variant):
    encoder = Encoder.from_pretrained(variant=variant, device="cpu")
    vectors[variant] = encoder.encode(wurcs, batch_size=8, pooling="first")
    hashes[variant] = encoder.metadata["sha256"]

reference = cosine_similarity(vectors[reference_variant], vectors[reference_variant])
comparison = cosine_similarity(vectors[comparison_variant], vectors[comparison_variant])
print(reference, comparison, hashes)

Each matrix compares glycans within one variant. Compare those matrices or neighbor rankings; do not mix vectors from different variants in one search index or use cross-variant cosine scores without validating their alignment. Similarity describes representation geometry, not binding performance.

first pooling uses the first sequence token. mean averages attended tokens, including START and END. Keep pooling fixed when comparing vectors. Inputs must be WURCS 2.0 within the trained 256-token limit; use encode_with_diagnostics to inspect tokenizer warnings.

The loader checks the checkpoint against its recorded SHA-256. Save that hash and pooling with your exports. from_pretrained also accepts revision="main" (or a Hub commit) and local_checkpoint="/path/to/sequence_encoder.pt" for a local file.

Binding: choose a regime and seed

The nine fits use the same prediction interface. Regime names describe held-out evaluation splits: glycans, proteins, or both. Choose the intended regime and record the exact cell_id:

Regime Seed 42 Seed 43 Seed 44
Glycan-cold mo_A_glycan_s42 moc_A_glycan_s43 moc_A_glycan_s44
Protein-cold mo_B_protein_s42 mof_B_protein_s43 mof_B_protein_s44
Doubly-cold mo_C_doubly_s42 mof_C_doubly_s43 mof_C_doubly_s44
python download_assets.py --task binding --cell-id mo_A_glycan_s42 --output-dir assets

Use the paths in assets/paths_binding.json with the binding instructions. Prediction requires pairs represented in the supplied glycan/protein cache. The adapter does not generate features for new glycans or protein FASTA. Outputs are regression scores, not calibrated probabilities or affinity constants.

For a backend, keep the selected model loaded and reuse a CachedInputs context across batches. The existing load_model, read_pairs and predict_batch functions provide the same preprocessing and prediction used by the command.

IAR and website reuse

IAR needs the full reference checkpoint and fitted linkage readout. Sequence encoder exports have no trained masked-token output head and cannot substitute for that checkpoint. python download_assets.py --task iar --output-dir assets retrieves it. The supplied command runs only the two recorded single-linkage cases; general uncertain-glycan input still needs an adapter and validation.

The collaborator can reuse Encoder.encode_with_diagnostics, cosine_similarity, nearest_neighbors and pca for embedding, similarity and visualization tools. The notebook shows embedding exports. Use the existing binding API and IAR command for their supported inputs; no replacement inference engine is needed. Follow WEBSITE_INTEGRATION.md for backend reuse, output labels and deployed checkpoint identity.