Push evaluation artifacts and model card

.gitattributes

@@ -33,3 +33,5 @@ saved_model/**/* filter=lfs diff=lfs merge=lfs -text
 *.zip filter=lfs diff=lfs merge=lfs -text
 *.zst filter=lfs diff=lfs merge=lfs -text
 *tfevents* filter=lfs diff=lfs merge=lfs -text
+calibration_threshold_sweep.png filter=lfs diff=lfs merge=lfs -text
+test_metrics.png filter=lfs diff=lfs merge=lfs -text

README.md

@@ -9,10 +9,171 @@ tags:
   - lora
   - peft
 library_name: transformers
+pipeline_tag: text-generation
 ---
 
 # PTM-LLaMA
 
-LoRA-fine-tuned `GreatCaptainNemo/ProLLaMA_Stage_1` for predicting post-translational modification (PTM) sites in protein sequences. A single LoRA adapter is instruction-tuned to handle three PTM types — methylation, phosphorylation, and ubiquitination — selected at inference time by the prompt. Output format: `Sites=<R5,D12,...>` regardless of PTM type.
+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**.
 
-This is a **training-only stub card**. Final metrics (per-PTM-type AUC, accuracy, precision, recall, F1, confusion matrices, and the cross-instruction ablation) are filled in by the companion evaluation notebook after running on the held-out test set.
+## 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`.

cross_instruction_auc.csv

@@ -0,0 +1,4 @@
+,Methylation,Phosphorylation,Ubiquitination
+Methylation,0.6565676387753249,0.4845552830017262,0.5799763786681202
+Phosphorylation,0.49783048886317616,0.6563748251514029,0.49054773775289695
+Ubiquitination,0.5141878873945271,0.4897625870439934,0.7429060321209946

inference_config.json

@@ -0,0 +1,17 @@
+{
+  "window_size": 21,
+  "stride": 5,
+  "max_new_tokens": 64,
+  "prompt_template": "Below is an instruction that describes a task. Write a response that appropriately completes the request.\n\n### Instruction:\n{instruction}\n{input}\n\n### Response:\n",
+  "instructions": {
+    "Methylation": "[Predict the methylation sites given the peptide sequence]",
+    "Phosphorylation": "[Predict the phosphorylation sites given the peptide sequence]",
+    "Ubiquitination": "[Predict the ubiquitination sites given the peptide sequence]"
+  },
+  "consensus_thresholds": {
+    "Methylation": 0.3333333432674408,
+    "Phosphorylation": 0.20000000298023224,
+    "Ubiquitination": 0.20000000298023224
+  },
+  "output_format": "Sites=<X1,X2,...> where Xi is one-letter residue + 1-indexed position within the window"
+}

metrics_summary.json

@@ -0,0 +1,192 @@
+{
+  "model_repo": "jbenbudd/ptm-llama",
+  "data_csv": "datasets/all_ptm_sites_site_level.csv",
+  "split_seed": 42,
+  "split_ratios": {
+    "train": 0.8,
+    "calibration": 0.1,
+    "test": 0.1
+  },
+  "window_size": 21,
+  "stride": 5,
+  "calibration": {
+    "Methylation": {
+      "n_proteins": 639,
+      "n_residues_evaluated": 423176,
+      "positive_prevalence": 0.0036675047734276048,
+      "roc_auc": 0.6188406072489271,
+      "locked_threshold": 0.3333333432674408,
+      "f1_at_threshold": 0.2036967182195398,
+      "precision_at_threshold": 0.2456778889899909,
+      "recall_at_threshold": 0.17396907216494845
+    },
+    "Phosphorylation": {
+      "n_proteins": 2970,
+      "n_residues_evaluated": 1681943,
+      "positive_prevalence": 0.027348132487248378,
+      "roc_auc": 0.6524042839293995,
+      "locked_threshold": 0.20000000298023224,
+      "f1_at_threshold": 0.32371403661726245,
+      "precision_at_threshold": 0.32454875223984964,
+      "recall_at_threshold": 0.3228836036349407
+    },
+    "Ubiquitination": {
+      "n_proteins": 3914,
+      "n_residues_evaluated": 2340476,
+      "positive_prevalence": 0.007890702575031746,
+      "roc_auc": 0.7622045108091328,
+      "locked_threshold": 0.20000000298023224,
+      "f1_at_threshold": 0.2864640844363501,
+      "precision_at_threshold": 0.19483078123476055,
+      "recall_at_threshold": 0.540827377084687
+    }
+  },
+  "test": {
+    "Methylation": {
+      "n_proteins": 649,
+      "n_residues_evaluated": 399838,
+      "positive_prevalence": 0.003729010249150906,
+      "threshold_applied": 0.3333333432674408,
+      "accuracy": 0.9944302442489208,
+      "precision": 0.20178282009724474,
+      "recall": 0.16700201207243462,
+      "specificity": 0.9975272814907605,
+      "f1": 0.18275229357798167,
+      "roc_auc": 0.6129889531736042,
+      "confusion_matrix": {
+        "TN": 397362,
+        "FP": 985,
+        "FN": 1242,
+        "TP": 249
+      }
+    },
+    "Phosphorylation": {
+      "n_proteins": 3040,
+      "n_residues_evaluated": 1732712,
+      "positive_prevalence": 0.02519749387087987,
+      "threshold_applied": 0.20000000298023224,
+      "accuracy": 0.9633343567771216,
+      "precision": 0.2944790352274373,
+      "recall": 0.32606504809894643,
+      "specificity": 0.979807016006612,
+      "f1": 0.3094681695162114,
+      "roc_auc": 0.6533860232975732,
+      "confusion_matrix": {
+        "TN": 1654945,
+        "FP": 34107,
+        "FN": 29424,
+        "TP": 14236
+      }
+    },
+    "Ubiquitination": {
+      "n_proteins": 3850,
+      "n_residues_evaluated": 2321360,
+      "positive_prevalence": 0.007665334114484612,
+      "threshold_applied": 0.20000000298023224,
+      "accuracy": 0.9785371506358341,
+      "precision": 0.1891293797922935,
+      "recall": 0.5475441159941553,
+      "specificity": 0.9818663758711493,
+      "f1": 0.2811467486185055,
+      "roc_auc": 0.7653999295817027,
+      "confusion_matrix": {
+        "TN": 2261794,
+        "FP": 41772,
+        "FN": 8051,
+        "TP": 9743
+      }
+    }
+  },
+  "cross_instruction_auc": {
+    "Methylation": {
+      "Methylation": 0.6565676387753249,
+      "Phosphorylation": 0.49783048886317616,
+      "Ubiquitination": 0.5141878873945271
+    },
+    "Phosphorylation": {
+      "Methylation": 0.4845552830017262,
+      "Phosphorylation": 0.6563748251514029,
+      "Ubiquitination": 0.4897625870439934
+    },
+    "Ubiquitination": {
+      "Methylation": 0.5799763786681202,
+      "Phosphorylation": 0.49054773775289695,
+      "Ubiquitination": 0.7429060321209946
+    }
+  },
+  "instruction_follow_rate": 0.1557890855457227,
+  "per_residue_breakdown": [
+    {
+      "ptm_type": "Methylation",
+      "residue": "K",
+      "n_residues": 24891,
+      "n_positive": 699,
+      "auc": 0.5158489475706036,
+      "accuracy": 0.9699088023783697,
+      "precision": 0.14285714285714285,
+      "recall": 0.01430615164520744,
+      "specificity": 0.9975198412698413,
+      "f1": 0.02600780234070221
+    },
+    {
+      "ptm_type": "Methylation",
+      "residue": "R",
+      "n_residues": 24176,
+      "n_positive": 746,
+      "auc": 0.6735408592590558,
+      "accuracy": 0.9407677035076109,
+      "precision": 0.20532646048109965,
+      "recall": 0.3203753351206434,
+      "specificity": 0.9605206999573197,
+      "f1": 0.250261780104712
+    },
+    {
+      "ptm_type": "Phosphorylation",
+      "residue": "S",
+      "n_residues": 149662,
+      "n_positive": 26557,
+      "auc": 0.6111038426325143,
+      "accuracy": 0.732096323716107,
+      "precision": 0.30830879021295876,
+      "recall": 0.40994841284783673,
+      "specificity": 0.8015921367937939,
+      "f1": 0.3519371575425496
+    },
+    {
+      "ptm_type": "Phosphorylation",
+      "residue": "T",
+      "n_residues": 95276,
+      "n_positive": 11688,
+      "auc": 0.573697049783009,
+      "accuracy": 0.8252130652000503,
+      "precision": 0.26372894260968877,
+      "recall": 0.2370807665982204,
+      "specificity": 0.9074508302627171,
+      "f1": 0.2496958774498761
+    },
+    {
+      "ptm_type": "Phosphorylation",
+      "residue": "Y",
+      "n_residues": 44179,
+      "n_positive": 5211,
+      "auc": 0.530815164126421,
+      "accuracy": 0.8510830937775866,
+      "precision": 0.22900158478605387,
+      "recall": 0.11091920936480522,
+      "specificity": 0.9500615889960994,
+      "f1": 0.14945054945054945
+    },
+    {
+      "ptm_type": "Ubiquitination",
+      "residue": "K",
+      "n_residues": 146021,
+      "n_positive": 17792,
+      "auc": 0.6234033873578442,
+      "accuracy": 0.6588093493401634,
+      "precision": 0.1891293797922935,
+      "recall": 0.5476056654676259,
+      "specificity": 0.6742390566876447,
+      "f1": 0.28115486170228116
+    }
+  ]
+}

per_residue_breakdown.csv

@@ -0,0 +1,7 @@
+ptm_type,residue,n_residues,n_positive,auc,accuracy,precision,recall,specificity,f1
+Methylation,K,24891,699,0.5158489475706036,0.9699088023783697,0.14285714285714285,0.01430615164520744,0.9975198412698413,0.02600780234070221
+Methylation,R,24176,746,0.6735408592590558,0.9407677035076109,0.20532646048109965,0.3203753351206434,0.9605206999573197,0.250261780104712
+Phosphorylation,S,149662,26557,0.6111038426325143,0.732096323716107,0.30830879021295876,0.40994841284783673,0.8015921367937939,0.3519371575425496
+Phosphorylation,T,95276,11688,0.573697049783009,0.8252130652000503,0.26372894260968877,0.2370807665982204,0.9074508302627171,0.2496958774498761
+Phosphorylation,Y,44179,5211,0.530815164126421,0.8510830937775866,0.22900158478605387,0.11091920936480522,0.9500615889960994,0.14945054945054945
+Ubiquitination,K,146021,17792,0.6234033873578442,0.6588093493401634,0.1891293797922935,0.5476056654676259,0.6742390566876447,0.28115486170228116