Update FastPLMs files

README.md

@@ -14,15 +14,16 @@ This checkpoint contains the FastPLMs `ESMC` implementation.
 
 Accepted inputs are amino-acid sequences tokenized to residue IDs.
 Supported Transformers entry points are `AutoConfig`, `AutoModel`,
-`AutoModelForMaskedLM`.
+`AutoModelForMaskedLM`, `AutoModelForSequenceClassification`,
+`AutoModelForTokenClassification`.
 
 ## Capabilities
 
 | Feature | Status |
 | --- | --- |
-| Sequence classification | Unavailable: no advertised AutoClass |
-| Token classification | Unavailable: no advertised AutoClass |
-| PEFT fine-tuning | Supported pattern: attach LoRA to the pretrained model |
+| Sequence classification | Supported: base weights with an untrained task head |
+| Token classification | Supported: base weights with an untrained task head |
+| PEFT fine-tuning | Supported pattern: preserve the separately trained `classifier` |
 | Embeddings | Supported: shared ordered embedding API |
 | Test-time training | Supported: low-rank masked-residue adaptation |
 | Attention variants | Special: SDPA fidelity path; alternate backends have explicit bands |
@@ -122,6 +123,45 @@ Set `output` and `format="safetensors"` or `"sqlite"` for transactional,
 bounded-memory storage. Resume checks input order, model state, tokenizer
 policy, backend, dtype, and pooling configuration before it appends data.
 
+## Downstream prediction
+
+The sequence and token prediction AutoClasses use the checkpoint backbone and
+create a new, untrained `classifier`. Sequence labels have shape `(b,)`.
+Residue labels have shape `(b, l)` and use `-100` outside biological positions.
+
+```python
+import torch
+from transformers import AutoTokenizer
+from transformers import (
+    AutoModelForSequenceClassification,
+    AutoModelForTokenClassification,
+)
+
+model_id = "Synthyra/ESMplusplus_large"
+sequence_model = AutoModelForSequenceClassification.from_pretrained(
+    model_id, num_labels=2, trust_remote_code=True
+).eval()
+token_model = AutoModelForTokenClassification.from_pretrained(
+    model_id, num_labels=3, trust_remote_code=True
+).eval()
+tokenizer = AutoTokenizer.from_pretrained(model_id, trust_remote_code=True)
+sequences = ["MSTNPKPQRKTKRNT", "MKTIIALSYIFCLVFA"]
+batch = tokenizer(sequences, padding=True, return_tensors="pt")
+biological = batch["attention_mask"].bool()
+for special_id in tokenizer.all_special_ids:
+    biological &= batch["input_ids"].ne(special_id)
+
+sequence_labels = torch.zeros(len(sequences), dtype=torch.long)
+token_labels = torch.full_like(batch["input_ids"], -100)
+token_labels[biological] = 0
+
+with torch.inference_mode():
+    sequence_output = sequence_model(**batch, labels=sequence_labels)
+    token_output = token_model(**batch, labels=token_labels)
+print(sequence_output.logits.shape)  # (b, 2)
+print(token_output.logits.shape)     # (b, l, 3)
+```
+
 ## PEFT fine-tuning
 
 Install the training dependencies. Then attach LoRA to the loaded checkpoint:
@@ -131,20 +171,22 @@ python -m pip install "datasets>=4.8,<5" "peft>=0.19,<0.20"
 ```
 
 ```python
-from peft import LoraConfig, get_peft_model
+from peft import LoraConfig, TaskType, get_peft_model
 
 peft_model = get_peft_model(
-    model,
+    sequence_model,
     LoraConfig(
+        task_type=TaskType.SEQ_CLS,
         r=8,
         lora_alpha=16,
         target_modules="all-linear",
+        modules_to_save=["classifier"],
     ),
 )
 ```
 
-This checkpoint has no advertised classifier. Supply the task objective and
-preserve any new head through `modules_to_save`.
+This checkpoint advertises a classification head. Save the separately trained
+`classifier` with the adapter.
 All FastPLMs checkpoints follow the Transformers `PreTrainedModel` contract and
 can use PEFT. The ESM2-specific shipped CLI is an example, not a
 support boundary. Record the target modules, base revision, data identity, and
@@ -269,8 +311,8 @@ and
 ## Runtime contract
 
 - Public input: Amino-acid sequences tokenized to residue IDs
-- Advertised AutoClasses: `AutoConfig`, `AutoModel`, `AutoModelForMaskedLM`
-- AutoClass weight status: `AutoConfig` = `FastPLMs extension`, `AutoModel` = `pretrained`, `AutoModelForMaskedLM` = `pretrained`
+- Advertised AutoClasses: `AutoConfig`, `AutoModel`, `AutoModelForMaskedLM`, `AutoModelForSequenceClassification`, `AutoModelForTokenClassification`
+- AutoClass weight status: `AutoConfig` = `FastPLMs extension`, `AutoModel` = `pretrained`, `AutoModelForMaskedLM` = `pretrained`, `AutoModelForSequenceClassification` = `base weights + untrained task head`, `AutoModelForTokenClassification` = `base weights + untrained task head`
 - Attention implementations: `eager`, `sdpa`, `flex_attention`, `flash_attention_2`, `flash_attention_3`
 - Precision policies: `default`, `fp8` (experimental)
 - BF16 execution: `static_parameters`
@@ -284,8 +326,8 @@ and
 ## Release record
 
 - FastPLMs weights: `Synthyra/ESMplusplus_large`
-- Runtime revision: recorded in the built artifact and published commit
-- Source-tree and runtime-bundle SHA-256: recorded in the source record
+- Runtime revision: recorded separately in the built artifact and published commit
+- Runtime source identities: recorded in `source-record.json`
 - Official checkpoint: `biohub/ESMC-600M`
 - Artifact source: `fast`
 - State transform: `esmc_to_fastplms_v1`

fastplms/models.toml

@@ -199,7 +199,7 @@ representative = "esmc_small"
 documentation = "docs/models.md#esm-and-esmc"
 test_tiers = ["check", "compliance", "feature", "artifact", "benchmark"]
 runtime_paths = ["__init__.py", "registry.py", "runtime.py", "models.toml", "models/__init__.py", "attention", "embeddings", "models/esm_plusplus", "models/ttt.py"]
-auto_map = { AutoConfig = "fastplms.models.esm_plusplus.modeling_esm_plusplus.ESMplusplusConfig", AutoModel = "fastplms.models.esm_plusplus.modeling_esm_plusplus.ESMplusplusModel", AutoModelForMaskedLM = "fastplms.models.esm_plusplus.modeling_esm_plusplus.ESMplusplusForMaskedLM" }
+auto_map = { AutoConfig = "fastplms.models.esm_plusplus.modeling_esm_plusplus.ESMplusplusConfig", AutoModel = "fastplms.models.esm_plusplus.modeling_esm_plusplus.ESMplusplusModel", AutoModelForMaskedLM = "fastplms.models.esm_plusplus.modeling_esm_plusplus.ESMplusplusForMaskedLM", AutoModelForSequenceClassification = "fastplms.models.esm_plusplus.modeling_esm_plusplus.ESMplusplusForSequenceClassification", AutoModelForTokenClassification = "fastplms.models.esm_plusplus.modeling_esm_plusplus.ESMplusplusForTokenClassification" }
 
 [families.esm3]
 architecture = "ESM3"
@@ -223,7 +223,7 @@ representative = "esm3_small"
 documentation = "docs/models.md#esm3"
 test_tiers = ["check", "compliance", "feature", "artifact", "benchmark"]
 runtime_paths = ["__init__.py", "registry.py", "runtime.py", "models.toml", "models/__init__.py", "attention", "embeddings", "models/esm3", "models/ttt.py"]
-auto_map = { AutoConfig = "fastplms.models.esm3.modeling_esm3.FastESM3Config", AutoModel = "fastplms.models.esm3.modeling_esm3.FastESM3Model" }
+auto_map = { AutoConfig = "fastplms.models.esm3.modeling_esm3.FastESM3Config", AutoModel = "fastplms.models.esm3.modeling_esm3.FastESM3Model", AutoModelForSequenceClassification = "fastplms.models.esm3.modeling_esm3.FastESM3ForSequenceClassification", AutoModelForTokenClassification = "fastplms.models.esm3.modeling_esm3.FastESM3ForTokenClassification" }
 
 [families.e1]
 architecture = "E1"
@@ -370,8 +370,8 @@ conversion_provenance = "Input: the pinned native Meta ESMFold checkpoint plus i
 representative = "esmfold"
 documentation = "docs/models.md#esmfold"
 test_tiers = ["check", "compliance", "structure", "feature", "artifact", "benchmark"]
-runtime_paths = ["__init__.py", "registry.py", "runtime.py", "models.toml", "models/__init__.py", "attention", "embeddings", "models/_esm_rotary.py", "models/esmfold"]
-auto_map = { AutoConfig = "fastplms.models.esmfold.modeling_fast_esmfold.FastEsmFoldConfig", AutoModel = "fastplms.models.esmfold.modeling_fast_esmfold.FastEsmForProteinFolding" }
+runtime_paths = ["__init__.py", "registry.py", "runtime.py", "models.toml", "models/__init__.py", "attention", "embeddings", "models/_esm_rotary.py", "models/classification_probe.py", "models/esmfold"]
+auto_map = { AutoConfig = "fastplms.models.esmfold.modeling_fast_esmfold.FastEsmFoldConfig", AutoModel = "fastplms.models.esmfold.modeling_fast_esmfold.FastEsmForProteinFolding", AutoModelForSequenceClassification = "fastplms.models.esmfold.modeling_fast_esmfold.FastEsmForSequenceClassification", AutoModelForTokenClassification = "fastplms.models.esmfold.modeling_fast_esmfold.FastEsmForTokenClassification" }
 
 [families.esmfold2]
 architecture = "ESMFold2"
@@ -396,8 +396,8 @@ conversion_provenance = "Input: each pinned Biohub ESMFold2 checkpoint and its s
 representative = "esmfold2"
 documentation = "docs/esmfold2.md"
 test_tiers = ["check", "compliance", "structure", "feature", "artifact", "benchmark"]
-runtime_paths = ["__init__.py", "registry.py", "runtime.py", "models.toml", "models/__init__.py", "attention", "embeddings", "models/esmfold2", "models/esm_plusplus", "models/ttt.py"]
-auto_map = { AutoConfig = "fastplms.models.esmfold2.configuration_esmfold2.ESMFold2Config", AutoModel = "fastplms.models.esmfold2.modeling_esmfold2.ESMFold2Model" }
+runtime_paths = ["__init__.py", "registry.py", "runtime.py", "models.toml", "models/__init__.py", "attention", "embeddings", "models/classification_probe.py", "models/_esm_rotary.py", "models/esmfold2", "models/esm_plusplus", "models/ttt.py"]
+auto_map = { AutoConfig = "fastplms.models.esmfold2.configuration_esmfold2.ESMFold2Config", AutoModel = "fastplms.models.esmfold2.modeling_esmfold2.ESMFold2Model", AutoModelForSequenceClassification = "fastplms.models.esmfold2.modeling_esmfold2_classification.ESMFold2ForSequenceClassification", AutoModelForTokenClassification = "fastplms.models.esmfold2.modeling_esmfold2_classification.ESMFold2ForTokenClassification" }
 
 [[models]]
 id = "esm2_8m"
@@ -1215,7 +1215,7 @@ official_files = [
   "config.json=git-sha1:79ed0dc0f867b8f09bfa004d6f77397c2ab9b38d",
   "model.safetensors=sha256:01358c317428d38535e3db513cab177336fc0f7fab0d84002e64b7741d5181b3",
 ]
-auto_map = { AutoConfig = "fastplms.models.esmfold2.configuration_esmfold2.ESMFold2Config", AutoModel = "fastplms.models.esmfold2.modeling_esmfold2_experimental.ESMFold2ExperimentalModel" }
+auto_map = { AutoConfig = "fastplms.models.esmfold2.configuration_esmfold2.ESMFold2Config", AutoModel = "fastplms.models.esmfold2.modeling_esmfold2_experimental.ESMFold2ExperimentalModel", AutoModelForSequenceClassification = "fastplms.models.esmfold2.modeling_esmfold2_classification.ESMFold2ExperimentalForSequenceClassification", AutoModelForTokenClassification = "fastplms.models.esmfold2.modeling_esmfold2_classification.ESMFold2ExperimentalForTokenClassification" }
 
 [[models]]
 id = "esmfold2_experimental_fast_cutoff2025"
@@ -1236,4 +1236,4 @@ official_files = [
   "config.json=git-sha1:0333d68ddb12ed2f066741dcb801142f466c0a2c",
   "model.safetensors=sha256:4e903b740ad6ad704ec60881bfd593e0d6c874a630ffa0f0838276e0b665088f",
 ]
-auto_map = { AutoConfig = "fastplms.models.esmfold2.configuration_esmfold2.ESMFold2Config", AutoModel = "fastplms.models.esmfold2.modeling_esmfold2_experimental.ESMFold2ExperimentalModel" }
+auto_map = { AutoConfig = "fastplms.models.esmfold2.configuration_esmfold2.ESMFold2Config", AutoModel = "fastplms.models.esmfold2.modeling_esmfold2_experimental.ESMFold2ExperimentalModel", AutoModelForSequenceClassification = "fastplms.models.esmfold2.modeling_esmfold2_classification.ESMFold2ExperimentalForSequenceClassification", AutoModelForTokenClassification = "fastplms.models.esmfold2.modeling_esmfold2_classification.ESMFold2ExperimentalForTokenClassification" }

modeling_fastplms.py

@@ -8,11 +8,12 @@ import sys
 import tempfile
 from io import BytesIO
 from pathlib import Path
+from typing import ClassVar
 from zipfile import ZIP_DEFLATED, ZipFile
 
 from .fastplms_bundle import RUNTIME_DATA, RUNTIME_HASH
 
-if RUNTIME_HASH != "23133ece4b4c336a3e782b316afd5bcf7367b522de2827cb6153885c1ffe1e70":
+if RUNTIME_HASH != "65b4cb38129822c36a448e295f77702ee64ca00fcfb9265b8c60f1c3ef169097":
     raise RuntimeError("FastPLMs runtime identity differs from the bridge.")
 
 _RUNTIME_TEMPORARIES = []
@@ -179,10 +180,14 @@ def _install_runtime():
     return package
 
 _install_runtime()
-_module_181 = _import_without_bytecode("fastplms.models.esm_plusplus.modeling_esm_plusplus")
-ESMplusplusConfig = _module_181.ESMplusplusConfig
+_module_182 = _import_without_bytecode("fastplms.models.esm_plusplus.modeling_esm_plusplus")
+ESMplusplusConfig = _module_182.ESMplusplusConfig
 ESMplusplusConfig.__module__ = __name__
-ESMplusplusForMaskedLM = _module_181.ESMplusplusForMaskedLM
+ESMplusplusForMaskedLM = _module_182.ESMplusplusForMaskedLM
 ESMplusplusForMaskedLM.__module__ = __name__
-ESMplusplusModel = _module_181.ESMplusplusModel
+ESMplusplusForSequenceClassification = _module_182.ESMplusplusForSequenceClassification
+ESMplusplusForSequenceClassification.__module__ = __name__
+ESMplusplusForTokenClassification = _module_182.ESMplusplusForTokenClassification
+ESMplusplusForTokenClassification.__module__ = __name__
+ESMplusplusModel = _module_182.ESMplusplusModel
 ESMplusplusModel.__module__ = __name__