---
base_model: GreatCaptainNemo/ProLLaMA_Stage_1
tags:
  - protein
  - ptm
  - methylation
  - phosphorylation
  - ubiquitination
  - lora
  - peft
library_name: transformers
pipeline_tag: text-generation
---

# PTM-LLaMA

LoRA fine-tune of [`GreatCaptainNemo/ProLLaMA_Stage_1`](https://huggingface.co/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`](./inference_config.json).

### Reference implementation

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

![training_loss](./training_loss.png)

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

![test_metrics](./test_metrics.png)

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

![cross_instruction_ablation](./cross_instruction_ablation.png)

![calibration_threshold_sweep](./calibration_threshold_sweep.png)

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