Metadata-Version: 2.5
Name: ablms
Version: 0.1.0
Summary: Unified API for antibody language models
Project-URL: Homepage, https://github.com/briney/ablms
Project-URL: Documentation, https://github.com/briney/ablms#readme
Project-URL: Repository, https://github.com/briney/ablms
Project-URL: Issues, https://github.com/briney/ablms/issues
Author-email: Bryan Briney <bryan.briney@gmail.com>
License-Expression: MIT
License-File: LICENSE
Keywords: antibody,bioinformatics,deep-learning,immunology,language-model
Classifier: Development Status :: 3 - Alpha
Classifier: Intended Audience :: Science/Research
Classifier: License :: OSI Approved :: MIT License
Classifier: Programming Language :: Python :: 3
Classifier: Programming Language :: Python :: 3.10
Classifier: Programming Language :: Python :: 3.11
Classifier: Programming Language :: Python :: 3.12
Classifier: Topic :: Scientific/Engineering :: Artificial Intelligence
Classifier: Topic :: Scientific/Engineering :: Bio-Informatics
Requires-Python: >=3.10
Requires-Dist: ablang2>=0.1.0
Requires-Dist: antiberty>=0.0.4
Requires-Dist: iglm>=0.1.0
Requires-Dist: numpy>=1.24.0
Requires-Dist: rjieba
Requires-Dist: torch>=2.0.0
Requires-Dist: transformers>=4.30.0
Provides-Extra: dev
Requires-Dist: black==26.5.1; extra == 'dev'
Requires-Dist: pytest-cov>=4.0.0; extra == 'dev'
Requires-Dist: pytest>=7.0.0; extra == 'dev'
Requires-Dist: ruff==0.16.2; extra == 'dev'
Requires-Dist: ty==0.0.40; extra == 'dev'
Description-Content-Type: text/markdown

# 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.

```python
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

| Model | Type | Paired Sequences | Source |
|-------|------|------------------|--------|
| **IgBERT** | Encoder | Yes | HuggingFace |
| **IgT5** | Encoder | Yes | HuggingFace |
| **AntiBERTa2** | Encoder | Yes | HuggingFace |
| **BALM** | Encoder | Yes | HuggingFace |
| **ft-ESM** | Encoder | Yes | HuggingFace |
| **ESM-2** | Encoder | No | HuggingFace |
| **AntiBERTy** | Encoder | No | antiberty package |
| **AbLang** | Encoder | No | ablang package |
| **AbLang2** | Encoder | Yes | ablang2 package |
| **IgLM** | Generative | No | iglm package |

## Installation

```bash
pip install ablms
```

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

### From Source

```bash
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:

```python
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:

```python
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.

```python
# 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:

```python
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:

```python
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:

```python
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:

```python
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):

```python
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:

```python
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:

```python
# 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:

```python
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

```python
# 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:

```python
# 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:

```python
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:

```python
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:

```python
# 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:

```python
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:

| Variable | Default | Effect |
| --- | --- | --- |
| `ABLMS_SUBMISSION_WINDOW` | `2` | Batches in flight per GPU. Lower to reduce shared memory use, raise to hide scheduling latency. |
| `ABLMS_WORKER_TIMEOUT` | `300` | Seconds to wait for a batch before failing. Raises `SharedMemoryError` if every worker is still alive, or `MultiGPUError` if a worker has died. |

Disable the progress bar for cleaner output in scripts:

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

## Available Models

List all registered models:

```python
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`:

```python
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:

```python
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`:

```python
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](LICENSE) for details.
