LangAAI

An antigen-conditioned masked language model for antibody CDR sequences.

This repository holds the weights only. The code that loads and runs them lives on GitHub: https://github.com/fbabd/LangAAI

Most antibody language models read the antibody alone. LangAAI reads the antibody and its target antigen in one joint attention stack, so the residues it predicts for a CDR loop depend on what that antibody is meant to bind. It is sequence-only and never sees a structure.

Usage

pip install git+https://github.com/fbabd/langaai.git
import langaai

model = langaai.load()          # downloads these weights on first use, then reuses them

ab = model.encode_antibody(heavy="QVQLQQPGAEL...", light="DIVMTQSPSSL...")
ag = model.embed_antigen("CPFGEVFNATRFASVY...")

masked = ab.mask_region("cdr3", spans=[(96, 108)])
[predictions] = model.predict_masked([(masked, ag)], top_k=5)
for p in predictions:
    print(p.position, p.top)    # [('W', 0.44), ('L', 0.18), ...]

There is also a command-line interface driven by JSON settings files (langaai predict|design|embed|attention) โ€” see docs/cli.md.

The weights are saved inside the installed package, at langaai/checkpoints/langaai.pt, and reused offline from then on. To fetch them ahead of time:

langaai download                 # prints where they landed

The loader uses a plain HTTPS URL, so nothing ties it to this Hub โ€” set LANGAAI_CHECKPOINT_URL to serve the weights from anywhere. To use a copy you already have:

export LANGAAI_CHECKPOINT=/path/to/langaai.pt

Requirements

transformers>=5.0. The checkpoint stores ESM-2's rotary embeddings in the layout transformers 5.0 introduced. Loading it under 4.x fails with missing keys for every layer's rotary_embeddings.inv_freq.

What's in the file

langaai.pt is a torch.save payload, loadable with weights_only=True.

The state dict holds 48.8M parameters in total: 33.3M across 198 tensors are the frozen ESM-2 tower (ab_tower.*, unchanged from facebook/esm2_t12_35M_UR50D, stored so the file loads standalone), and 15.5M across 65 tensors are the trained parts โ€” the two side adapters, the joint stack, the segment embeddings, the pooled token and the tied LM head. The loading code instantiates the backbone from the Hub first, then restores this state dict over it.

Model details

Architecture Two ESM-2 towers into a shared 4-layer, 8-head joint self-attention stack
Backbone facebook/esm2_t12_35M_UR50D (width 480), frozen
Trained components Per-side adapters, joint stack, tied LM head
Inputs Heavy (+ optional light) variable domain, plus antigen sequence
Max trained joint length 1,024 tokens, covering antibody and antigen
Size ~195 MB

Intended use and limitations

  • score_sequence is not a binding-affinity predictor. It is a masked-reconstruction log-likelihood. For affinity, fit your own regressor on cls_embedding.
  • No structural validation. Everything is sequence-only.
  • CDR spans must be supplied by you. No CDR-numbering tool is bundled, so the model cannot locate CDR-H3 in a raw sequence.
  • Antigens beyond ~780 residues push the joint sequence past the 1,024 tokens seen in training; they run, but out of distribution.
  • Protein sequences only โ€” no nucleotides, no modified residues beyond X, no multi-antigen complexes.

Citation

License

MIT, for both the weights and the code.

Downloads last month

-

Downloads are not tracked for this model. How to track
Inference Providers NEW
This model isn't deployed by any Inference Provider. ๐Ÿ™‹ Ask for provider support

Model tree for faisalashraf/langaai

Finetuned
(67)
this model