PTM-LLaMA

LoRA fine-tune of GreatCaptainNemo/ProLLaMA_Stage_1 instruction-tuned to predict post-translational modification (PTM) sites for three PTM types in a single adapter: methylation, phosphorylation, and ubiquitination.

Task

Given a 21-residue peptide window and a PTM-type instruction, the model generates the list of modified positions in the format Sites=<R5,D12,...> (residue letter + 1-indexed position within the window). The PTM type to predict is selected by the instruction prompt; the output format is shared across PTM types.

Inference architecture

Full-protein inference for a single PTM type proceeds in three steps:

1. Sliding windows. A 21-residue window is slid across the input sequence with stride 5; a tail window is appended so the final residues are covered.
2. Per-window generation. The model generates Sites=<...> for each window, prompted with the PTM-type instruction. Each predicted residue letter is validated against the window sequence at the indicated position; mismatches are discarded. Validated predictions are mapped to full-protein coordinates and accumulated into a per-residue consensus score, defined as (# windows predicting the residue as a site) / (# windows covering the residue).
3. Thresholding. A per-PTM-type F1-optimal threshold (derived from the held-out calibration set) is applied to the consensus scores.

The per-PTM-type thresholds, windowing parameters, and prompt template are persisted in inference_config.json.

Reference implementation

import re, json, torch
from huggingface_hub import hf_hub_download
from transformers import AutoTokenizer, AutoModelForCausalLM

REPO = "jbenbudd/ptm-llama"

cfg = json.load(open(hf_hub_download(REPO, "inference_config.json")))
tok = AutoTokenizer.from_pretrained(REPO)
if tok.pad_token is None:
    tok.pad_token = tok.unk_token
mdl = AutoModelForCausalLM.from_pretrained(
    REPO, torch_dtype=torch.float16, device_map="auto"
).eval()

SITE_RE  = re.compile(r"^([A-Z])(\d+)$")
SITES_RE = re.compile(r"Sites=<([^>]*)>")


@torch.no_grad()
def predict_sites(seq: str, ptm_type: str):
    """Predict PTM-site positions in a full protein for a given PTM type.

    Returns sites like ['K161', 'S203'] in 1-indexed full-protein coords.
    """
    if ptm_type not in cfg["instructions"]:
        raise ValueError(f"Unknown PTM type {ptm_type}; supported: {list(cfg['instructions'])}")
    w, s = cfg["window_size"], cfg["stride"]
    t = cfg["consensus_thresholds"][ptm_type]
    instruction = cfg["instructions"][ptm_type]
    L = len(seq)

    if L <= w:
        starts = [0]
    else:
        starts = list(range(0, L - w + 1, s))
        if starts[-1] + w < L:
            starts.append(L - w)

    covered, predicted = [0] * L, [0] * L
    for st in starts:
        win = seq[st:st + w]
        prompt = cfg["prompt_template"].format(instruction=instruction, input=f"Seq=<{win}>")
        enc = tok(prompt, return_tensors="pt").to(mdl.device)
        out = mdl.generate(
            **enc, max_new_tokens=cfg["max_new_tokens"], do_sample=False,
            pad_token_id=tok.pad_token_id,
        )
        text = tok.decode(out[0][enc.input_ids.shape[1]:], skip_special_tokens=True)

        for i in range(st, min(st + w, L)):
            covered[i] += 1
        m = SITES_RE.search(text)
        if not m:
            continue
        for part in m.group(1).split(","):
            mm = SITE_RE.match(part.strip())
            if not mm:
                continue
            letter, pos_local = mm.group(1), int(mm.group(2))
            pos_full = st + pos_local
            if 1 <= pos_full <= L and seq[pos_full - 1] == letter:
                predicted[pos_full - 1] += 1

    return [f"{seq[i]}{i + 1}" for i in range(L)
            if covered[i] > 0 and predicted[i] / covered[i] >= t]


# Example usage.
sites = predict_sites(
    "MASDEGKLFVGGLSFDTNEQALEQVFSKYGQISEVVVVKDRETQRSRGFGFVTFENIDDAKDAMMAMNGK",
    ptm_type="Phosphorylation",
)
print(sites)

Training

• Base model: GreatCaptainNemo/ProLLaMA_Stage_1
• Method: LoRA SFT via trl.SFTTrainer + peft.LoraConfig
• LoRA config: r=64, alpha=128, dropout=0.05, target modules = q,k,v,o,gate,down,up_proj
• Optimizer: AdamW, lr=3e-4, cosine schedule, warmup=40 steps, max_grad_norm=1.0
• Batching: per-device batch 16 × grad_accum 8 (effective 128) at bf16, max_seq_length 2048
• Epochs: up to 8, with EarlyStoppingCallback(patience=3) on eval_loss and load_best_model_at_end=True
• Source data: datasets/all_ptm_sites_site_level.csv (long format: one row per annotated PTM site across methylation, phosphorylation, and ubiquitination)
• Split: protein-level partition with seed 42 and ratios 0.80 / 0.10 / 0.10 (train / calibration / test). All annotations of a given uniprot_id are assigned to a single split.
• Training distribution: natural — the relative training-set abundance across PTM types reflects the natural distribution of annotated sites in the source data (phosphorylation ≫ ubiquitination ≫ methylation). Each (protein, PTM_type) record contributes all sliding windows over its sequence, including windows with no in-window site of that PTM type as in-context negatives.

Evaluation methodology

Two disjoint protein-level partitions are used to separate per-PTM-type threshold selection from final metric reporting:

• Calibration set (5,943 proteins): the F1-optimal threshold over the per-residue consensus ROC is selected per PTM type.
• Test set (5,944 proteins): the calibration-derived per-PTM-type thresholds are applied without modification; the metrics below are unbiased point estimates of generalization.

Per-PTM-type results

PTM type	Test residues	Positive prevalence	Calibration AUC	Locked t	Test AUC	Accuracy	Precision	Recall	Specificity	F1	TN / FP / FN / TP
Methylation	399,838	0.37%	0.619	0.333	0.613	99.44%	20.18%	16.70%	99.75%	18.28%	397,362 / 985 / 1,242 / 249
Phosphorylation	1,732,712	2.52%	0.652	0.200	0.653	96.33%	29.45%	32.61%	97.98%	30.95%	1,654,945 / 34,107 / 29,424 / 14,236
Ubiquitination	2,321,360	0.77%	0.762	0.200	0.765	97.85%	18.91%	54.75%	98.19%	28.11%	2,261,794 / 41,772 / 8,051 / 9,743

Per-residue-type breakdown on test

PTM type	Residue	Total	Positive	AUC	Accuracy	Precision	Recall	Specificity	F1
Methylation	K	24,891	699	0.516	96.99%	14.29%	1.43%	99.75%	2.60%
Methylation	R	24,176	746	0.674	94.08%	20.53%	32.04%	96.05%	25.03%
Phosphorylation	S	149,662	26,557	0.611	73.21%	30.83%	40.99%	80.16%	35.19%
Phosphorylation	T	95,276	11,688	0.574	82.52%	26.37%	23.71%	90.75%	24.97%
Phosphorylation	Y	44,179	5,211	0.531	85.11%	22.90%	11.09%	95.01%	14.95%
Ubiquitination	K	146,021	17,792	0.623	65.88%	18.91%	54.76%	67.42%	28.12%

Cross-instruction ablation

For a sub-sample of test proteins, sliding-window inference was run three times — once per PTM-type instruction — on the same protein sequences. AUC is reported against each PTM type's ground-truth labels.

A working instruction-tuned model should have higher AUC on the diagonal (matched instruction) than off-diagonal (mismatched instruction).

	Methylation	Phosphorylation	Ubiquitination
Methylation	0.657	0.485	0.580
Phosphorylation	0.498	0.656	0.491
Ubiquitination	0.514	0.490	0.743

Instruction-following rate (fraction of windows whose output differs across the three prompts on a 50-protein sub-sample): 15.58%.

Limitations

• Window-local outputs. The model emits positions inside a 21-residue window. Full-protein predictions are produced by the sliding-window aggregation described in Inference architecture; very short proteins (length < 21) are scored as a single window.
• No structural context. The model sees only primary sequence; structurally-disfavored false positives cannot be filtered without external 3D information.
• Three PTM types only. Predictions are well-defined for methylation, phosphorylation, and ubiquitination. Generalization to other PTM types is not evaluated.

Reproduction

Execute training/train_ptm_llama.ipynb followed by evaluation/evaluate_ptm_llama.ipynb from the source repository. Both notebooks are self-contained and intended for execution in Google Colab. The protein-level split is deterministic given SPLIT_SEED = 42.