ablms

A unified Python API for antibody language models.

Overview

Working with antibody language models often means dealing with different architectures, tokenizers, input formats, and output structures. ablms provides a consistent interface across multiple models, so you can focus on your research instead of wrestling with model-specific quirks.

from ablms import AntibodySequence, load_model

# Same API for any model
model = load_model("balm")  # or "antiberty", "ablang2", "igbert", etc.
input = AntibodySequence(
    heavy="EVQLVESGGGLVQPGRSLRLSCAASGFTFS",
    light="DIQMTQSPSSLSASVGDRVTITCRASQSIS",
)
embeddings = model.get_embeddings([input])  # get_embeddings accepts a list of AntibodySequence objects

Supported Models

Installation

pip install ablms

This installs ablms along with all required dependencies including PyTorch, Transformers, and the model-specific packages (antiberty, ablang2, iglm).

From Source

git clone https://github.com/bryanbriney/ablms.git
cd ablms
pip install -e .

Quickstart

Creating Antibody Sequences

The AntibodySequence class provides a unified way to represent antibody sequences. All arguments must be passed as keywords to ensure the chain type is always explicit:

from ablms import AntibodySequence, Species

# Single heavy chain
heavy_seq = AntibodySequence(heavy="EVQLVESGGGLVQPGRSLRLSCAASGFTFS")

# Single light chain
light_seq = AntibodySequence(light="DIQMTQSPSSLSASVGDRVTITCRASQSIS")

# Paired heavy and light chains
paired_seq = AntibodySequence(
    heavy="EVQLVESGGGLVQPGRSLRLSCAASGFTFS",
    light="DIQMTQSPSSLSASVGDRVTITCRASQSIS",
    species=Species.HUMAN
)

# Check sequence properties
print(paired_seq.is_paired)      # True
print(paired_seq.length)         # {'heavy': 30, 'light': 30}
print(paired_seq.total_length)   # 60

Getting Embeddings

Extract token-level or sequence-level embeddings from any encoder model:

from ablms import load_model, AntibodySequence

# Load a model
model = load_model("balm")

# Prepare sequences
sequences = [
    AntibodySequence(heavy="EVQLVESGGGLVQPGRSLRLSCAASGFTFS"),
    AntibodySequence(heavy="QVQLVQSGAEVKKPGASVKVSCKASGYTFT"),
]

# Get token-level embeddings
output = model.get_embeddings(sequences)
print(output.embeddings.shape)  # [2, seq_len, 1024]

# Get sequence-level embeddings with pooling
# Pooling options are "mean", "max", "cls", "first", and "last" 
pooled = model.get_embeddings(sequences, pooling="mean")
print(pooled.embeddings.shape)  # [2, 1024]

# Select several layers, or every layer, by passing a list or "all"
# A layer axis is inserted at dimension 1
multi = model.get_embeddings(sequences, layer=[0, 6, 12], pooling="cls")
print(multi.embeddings.shape)  # [2, 3, 1024]
print(multi.get_layer(6).shape)  # [2, 1024]

# Concatenate every layer into one feature vector per sequence,
# the usual input for a UMAP or t-SNE projection
every = model.get_embeddings(sequences, layer="all", pooling="cls")
print(every.concat_layers().shape)  # [2, 25 * 1024]

Token-level output for many layers is large — layer="all" on BALM's 24-block model is roughly 25x the single-layer payload — so pair it with iter_embeddings() rather than get_embeddings() for anything sizeable. Pooled multi-layer runs stay small: pooling is applied per layer before the layers are stacked.

AbLang exposes only its final layer and raises UnsupportedOperationError for any other selection.

# Stream batches instead of accumulating, for datasets larger than memory
for batch in model.iter_embeddings(sequences, pooling="mean", batch_size=64):
    ...  # batch is an EmbeddingOutput covering just this batch

Working with Paired Sequences

Models that support paired sequences (IgBERT, IgT5, BALM, AbLang2) can process heavy and light chains together:

from ablms import load_model, AntibodySequence

model = load_model("balm")  # Supports paired sequences

paired = AntibodySequence(
    heavy="EVQLVESGGGLVQPGRSLRLSCAASGFTFS",
    light="DIQMTQSPSSLSASVGDRVTITCRASQSIS"
)

output = model.get_embeddings([paired])

# Extract chain-specific embeddings
heavy_emb = output.get_chain_embeddings(0, "heavy")
light_emb = output.get_chain_embeddings(0, "light")

Attention Weights

Visualize or analyze attention patterns:

from ablms import load_model

model = load_model("igbert")
sequences = ["EVQLVESGGGLVQPGRSLRLSCAASGFTFS"]

attention = model.get_attention(sequences)
print(attention.num_layers)  # 30
print(attention.num_heads)   # 16

# Get attention from a specific layer and head
layer_5_head_0 = attention.get_head(layer=5, head=0)

# Get mean attention across all layers and heads
mean_attention = attention.get_mean_attention()

Mask Filling

Predict amino acids at masked positions:

from ablms import load_model, AntibodySequence

model = load_model("igbert")

# Create a sequence with masks
masked_seq = AntibodySequence(heavy="EVQL<MASK>ESGGGLVQPGRSLRL")

# Fill the mask with top predictions
predictions = model.fill_mask([masked_seq], top_k=5)

for pred in predictions[0]:
    print(pred.heavy_chain)

Generating New Sequences

Use generative models like IgLM to create new antibody sequences:

from ablms import load_model, ChainType, Species

model = load_model("iglm")

# Generate new heavy chain sequences
output = model.generate(
    num_sequences=5,
    chain_type=ChainType.HEAVY,
    species=Species.HUMAN,
    temperature=1.0
)

for seq in output.sequences:
    print(seq.heavy_chain)

# Get the best sequences by score
top_sequences = output.get_top_k(k=3)

Computing Sequence Likelihoods

Score sequences using pseudo log-likelihood (encoder models) or log-likelihood (generative models):

from ablms import load_model, AntibodySequence, ChainType, Species

# Encoder model: pseudo log-likelihood
encoder = load_model("igbert")
sequences = [
    AntibodySequence(heavy="EVQLVESGGGLVQPGRSLRL"),
    AntibodySequence(heavy="QVQLVQSGAEVKKPGASVKV"),
]
pll_scores = encoder.pseudo_log_likelihood(sequences)

# Generative model: log-likelihood
generator = load_model("iglm")
ll_scores = generator.log_likelihood(
    sequences,
    chain_type=ChainType.HEAVY,
    species=Species.HUMAN
)

Mask Scanning

Analyze model predictions at every position by masking each residue one at a time:

from ablms import load_model, AntibodySequence

model = load_model("igbert")
seq = AntibodySequence(
    heavy="EVQLVESGGGLVQPGRSLRL",
    light="DIQMTQSPSSLSASVGDRVT"
)

# Scan all positions
output = model.mask_scan(seq)

# Basic metrics
print(output.accuracy(agg="mean"))     # Mean prediction accuracy
print(output.perplexity(agg="mean"))   # Mean perplexity
print(output.entropy(agg="mean"))      # Mean entropy

# Per-position values (no aggregation)
accuracy_per_pos = output.accuracy()   # Tensor of shape [seq_len]

Chain-Specific Metrics

Extract metrics for individual chains:

# Get accuracy for each chain
heavy_acc = output.get_chain_accuracy("heavy", agg="mean")
light_acc = output.get_chain_accuracy("light", agg="mean")

# Same for perplexity and entropy
heavy_ppl = output.get_chain_perplexity("heavy", agg="mean")
light_ent = output.get_chain_entropy("light", agg="mean")

Custom Position Masking

Focus metrics on specific positions (e.g., CDR regions) using boolean masks:

import torch

# Build a mask from chain-specific masks
# True = include position, False = exclude position
heavy_cdr_mask = torch.zeros(20, dtype=torch.bool)
heavy_cdr_mask[5:12] = True  # Only include positions 5-11

# Create full-sequence mask from chain masks
mask = output.build_mask(heavy=heavy_cdr_mask)  # light chain defaults to all True

# Compute metrics only for masked positions
cdr_accuracy = output.accuracy(mask=mask, agg="mean")
cdr_perplexity = output.perplexity(mask=mask, agg="mean")

# Or use chain-specific methods directly with chain-length masks
heavy_cdr_acc = output.get_chain_accuracy("heavy", mask=heavy_cdr_mask, agg="mean")

Additional Properties

# Raw predictions
print(output.predictions)        # Predicted token indices
print(output.predicted_tokens)   # Predicted tokens as strings (if vocab available)
print(output.probabilities)      # Softmax probabilities [seq_len, vocab_size]

# Top-k predictions at each position
values, indices = output.top_k_predictions(k=5)

Key Concepts

Unified Mask Token

All models use <MASK> as the mask token internally. ablms automatically converts this to each model's native mask token:

# You always use <MASK>
seq = AntibodySequence(heavy="EVQL<MASK>ESGG")

# ablms converts it to the model's token:
# IgBERT: [MASK]
# AntiBERTy: _
# BALM: <mask>
# AbLang: *
# AbLang2: *
# ft-ESM: <mask>
# ESM-2: <mask>

Output Classes

All methods return structured output objects with helpful properties:

• EmbeddingOutput: Token or sequence embeddings with get_chain_embeddings() for extracting specific chains. Multi-layer results (from layer=[...] or layer="all") carry a layers list of the resolved indices, plus get_layer() and concat_layers() for extracting or flattening the layer axis
• LogitsOutput: MLM logits with probabilities, predictions, and top_k_predictions()
• AttentionOutput: Attention weights with get_layer(), get_head(), and get_mean_attention()
• GenerationOutput: Generated sequences with get_top_k() and filter_by_score()
• MaskScanOutput: Per-position predictions with accuracy(), perplexity(), entropy(), and build_mask() for custom position filtering

Device Management

Models automatically use all available GPUs for parallel inference:

from ablms import load_model

# Auto-detects and uses all available GPUs
model = load_model("igbert")
print(model.num_devices)  # e.g., 4
print(model.devices)      # [device(type='cuda', index=0), ...]

# Or specify specific GPUs
model = load_model("igbert", devices=[0, 2, 3])

# Single GPU (no parallelization overhead)
model = load_model("igbert", devices="cuda:0")

# CPU only
model = load_model("igbert", devices="cpu")

# Move model after loading (resets to single device)
model.to("cuda:1")

Multi-GPU Parallelism

When multiple GPUs are available, inference is automatically parallelized. Work is distributed across GPUs using a worker pool, with each GPU holding a complete model replica:

from ablms import load_model, AntibodySequence

# Load model (auto-detects 4 GPUs)
model = load_model("igbert")

# Process 10,000 sequences - automatically distributed across GPUs
sequences = [AntibodySequence(heavy=seq) for seq in heavy_chains]
embeddings = model.get_embeddings(
    sequences,
    batch_size=64,       # Per-GPU batch size
    show_progress=True,  # tqdm progress bar (default: True)
)

Key features:

• Automatic detection: Uses all available GPUs by default
• Lazy initialization: Worker processes spawn on first inference call
• Single-GPU optimization: No subprocess overhead when using one device
• Bounded in-flight memory: At most a few batches per GPU are in flight at once, so shared memory use does not grow with dataset size. The result get_embeddings() returns is still proportional to the dataset - use iter_embeddings() when that is the constraint
• Progress tracking: Built-in tqdm progress bar for all inference methods

Large datasets

Results travel from worker processes to the parent through shared memory (/dev/shm), so what matters for very large runs is how much each batch carries. Two things keep that bounded.

Pool inside the batch. When you pass pooling=, the reduction happens on the GPU before the batch is transferred, so the full token-level tensor is never materialized:

# Each batch transfers [batch_size, hidden_dim], not
# [batch_size, seq_len, hidden_dim] - roughly 250x smaller at typical lengths.
embeddings = model.get_embeddings(sequences, pooling="mean", batch_size=64)

Stream token-level output. When you need per-residue embeddings for more sequences than fit in memory, iter_embeddings() yields one batch at a time, in input order, and retains nothing:

import h5py

with h5py.File("embeddings.h5", "w") as f:
    for i, batch in enumerate(model.iter_embeddings(sequences, batch_size=64)):
        for j, tokens in enumerate(batch):  # iterating strips padding
            f.create_dataset(f"seq_{i * 64 + j}", data=tokens.numpy())

If a run still exhausts shared memory, ablms raises SharedMemoryError with the current /dev/shm free space and suggested remedies. Inside a container the usual cause is Docker's 64 MB default; raise it with --shm-size=8g.

Two environment variables tune this:

Disable the progress bar for cleaner output in scripts:

embeddings = model.get_embeddings(sequences, show_progress=False)

Available Models

List all registered models:

from ablms import list_models

print(list_models())
# {'igbert': 'encoder', 'igt5': 'encoder', 'antiberta2': 'encoder',
#  'balm': 'encoder', 'antiberty': 'encoder', 'ablang': 'encoder',
#  'ablang2': 'encoder', 'ftesm': 'encoder', 'esm2-8m': 'encoder',
#  'esm2-35m': 'encoder', 'esm2-150m': 'encoder', 'esm2-650m': 'encoder',
#  'esm2-3b': 'encoder', 'esm2-15b': 'encoder', 'iglm': 'generative'}

Notes on Specific Models

IgT5

IgT5 is an encoder-only T5 model and does not have a masked language modeling head. Methods like get_logits(), pseudo_log_likelihood(), and fill_mask() will raise UnsupportedOperationError:

from ablms import load_model

model = load_model("igt5")

# These work:
embeddings = model.get_embeddings(sequences)
attention = model.get_attention(sequences)

# These raise UnsupportedOperationError:
# model.get_logits(sequences)
# model.fill_mask(sequences)

ft-ESM

ft-ESM is an ESM2-based model (finetuned from facebook/esm2_t33_650M_UR50D) optimized for paired antibody sequences. It uses a unique <cls><cls> separator (two consecutive CLS tokens) between chains:

from ablms import load_model, AntibodySequence

model = load_model("ftesm")

# Paired sequences work well with ft-ESM
paired = AntibodySequence(
    heavy="EVQLVESGGGLVQPGRSLRLSCAASGFTFS",
    light="DIQMTQSPSSLSASVGDRVTITCRASQSIS"
)
embeddings = model.get_embeddings([paired])

# Single chain sequences also work
single = AntibodySequence(heavy="EVQLVESGGGLVQPGRSLRLSCAASGFTFS")
embeddings = model.get_embeddings([single])

Single-Chain Models

AntiBERTy and AbLang only support single chain sequences. Passing paired sequences will raise PairedSequenceError:

from ablms import load_model, AntibodySequence

model = load_model("antiberty")  # Single-chain only
# or
model = load_model("ablang")     # Single-chain only

# This works:
model.get_embeddings([AntibodySequence(heavy="EVQLVESGG...")])

# This raises PairedSequenceError:
# model.get_embeddings([AntibodySequence(heavy="...", light="...")])

Note: AbLang uses separate models for heavy and light chains. The appropriate model is automatically selected based on the input sequence type, and mixed batches (containing both heavy and light chain sequences) are supported.

License

MIT License - see LICENSE for details.