PhyloGPN demo

Load PhyloGPN through explicit AutoClass registration, inspect its F81 rate parameters and embeddings, convert the rates to nucleotide probabilities, and compute the existing C-to-T substitution score.

Setup

from importlib.metadata import version

from gpn import register_auto_classes
import torch
from transformers import AutoModel, AutoTokenizer

register_auto_classes("phylo")
MODEL_ID = "songlab/PhyloGPN"
MODEL_REVISION = "3556db4c469e67d25f0f7a0a6653b48be3eebf51"
DEVICE = torch.device("cuda" if torch.cuda.is_available() else "cpu")

tokenizer = AutoTokenizer.from_pretrained(MODEL_ID, revision=MODEL_REVISION)
model = AutoModel.from_pretrained(MODEL_ID, revision=MODEL_REVISION).to(DEVICE).eval()
print(
    f"gpn={version('gpn')}, transformers={version('transformers')}, "
    f"torch={version('torch')}, device={DEVICE}, "
    f"model dtype={next(model.parameters()).dtype}; "
    f"{MODEL_ID} @ {MODEL_REVISION}"
)

gpn=0.9.0a1, transformers=5.15.0, torch=2.13.0+cpu, device=cpu, model dtype=torch.float32; songlab/PhyloGPN @ 3556db4c469e67d25f0f7a0a6653b48be3eebf51

Sequence batch

sequences = ['TATAAA', 'GGCCAATCT', 'CACGTG', 'AGGTCACGT', 'GCCAGCC', 'GGGGATTTCC']
receptive_field = 481
pad_size = receptive_field // 2
pad_token = tokenizer.pad_token
padded_sequences = [
    pad_token * pad_size + sequence + pad_token * pad_size
    for sequence in sequences
]
input_ids = tokenizer(
    padded_sequences, return_tensors="pt", padding=True
)["input_ids"].to(DEVICE)

with torch.inference_mode():
    padded_embeddings = model.get_embeddings(input_ids)
    padded_logits = model(input_ids)

embeddings = []
logits = []
for row, sequence in enumerate(sequences):
    length = len(sequence)
    embeddings.append(padded_embeddings[row, :length].cpu())
    logits.append(
        {base: padded_logits[base][row, :length].cpu() for base in "ACGT"}
    )

Position embeddings

embeddings[0]

tensor([[ 0.2391, -0.4661, -0.0159,  ...,  0.4051, -0.7045, -0.5292],
        [-0.3234, -0.5837,  0.5243,  ..., -0.0163, -0.4431, -0.5394],
        [ 0.0789, -0.3639, -0.1462,  ...,  0.3260, -0.6638, -0.3064],
        [-0.7663, -0.1948,  0.2535,  ...,  0.5135, -0.1959, -0.3657],
        [ 0.0413, -0.2288,  0.2194,  ..., -0.1099, -0.0648, -0.6248],
        [-0.1851, -0.1795,  0.3790,  ...,  0.0529,  0.0603,  0.1869]])

F81 log-rate parameters

logits[0]

{'A': tensor([1.6577, 4.0178, 1.0112, 3.9506, 3.7107, 3.6846]),
 'C': tensor([1.1716, 0.4637, 0.9406, 0.4236, 0.8937, 0.7400]),
 'G': tensor([-0.1792,  1.1369, -0.0490,  0.8730,  0.3609,  0.7772]),
 'T': tensor([4.0944, 1.2592, 3.8543, 1.4166, 1.8473, 1.4589])}

Nucleotide probabilities

probability_batches = []
for logit_dictionary in logits:
    logit_tensor = torch.stack(
        [logit_dictionary[base] for base in "ACGT"], dim=-1
    )
    probability_tensor = torch.softmax(logit_tensor, dim=-1)
    probability_batches.append(
        {
            base: probability_tensor[:, nucleotide_index]
            for nucleotide_index, base in enumerate("ACGT")
        }
    )

probability_batches[0]

{'A': tensor([0.0757, 0.8710, 0.0514, 0.8659, 0.8000, 0.8229]),
 'C': tensor([0.0466, 0.0249, 0.0479, 0.0255, 0.0478, 0.0433]),
 'G': tensor([0.0121, 0.0488, 0.0178, 0.0399, 0.0281, 0.0449]),
 'T': tensor([0.8657, 0.0552, 0.8829, 0.0687, 0.1241, 0.0889])}

Zero-shot substitution score

At zero-based position 1 of the first sequence, the model’s hypothetical T-vs-C log-rate difference is shown below. The observed base there is A, so this is a conditional preference comparison, not a reference-allele-checked VEP call.

logits[0]["T"][1] - logits[0]["C"][1]

tensor(0.7955)