Keep original model card; update model list (add ESM-C V0.3 line) + license note

README.md

@@ -1,57 +1,188 @@
 ---
 license: mit
-library_name: pytorch
-pipeline_tag: other
-tags:
-- protein
-- tcr
-- tcr-pmhc
-- immunology
-- esm-c
-- esm2
-- masked-language-model
 ---
 
-# DecoderTCR V0.3
+# DecoderTCR
+DecoderTCR is a protein language model for T-cell receptor (TCR) & peptide-MHC complexes. The models are based on the ESM-2 and ESM-C model families.
 
-TCR–pMHC binding scoring by masked-language-model pseudo-log-likelihood (PLL), on ESM-2 and
-ESM-C backbones. Given a TCR (V/J genes + CDR3) paired with an HLA allele and a peptide, the
-models score how well the TCR is predicted to bind the peptide–MHC.
+For Model Code and additional information on installation/usage please see [the associated GitHub repository](https://github.com/czbiohub-chi/DecoderTCR)
 
-This repository hosts **model weights only**. The code, loader, and prediction CLIs live in the
-DecoderTCR package; weights are fetched into the expected paths by its `download_weights.py`
-script — you do not load these checkpoints with `transformers`.
+## Model Architecture
 
-## Models
+DecoderTCR is built on a Transformer-based protein language model (ESM-2 and ESM-C families).
 
-| File | Registry name | Backbone | Params | Version | Notes |
-|------|---------------|----------|--------|---------|-------|
-| `DecoderTCR-C-V0.3/600M.ckpt` | `DecoderTCR-C_600M` | ESM-C | 600M | V0.3 | **default**, runs on ≤24 GB GPUs |
-| `DecoderTCR-C-V0.3/300M.ckpt` | `DecoderTCR-C_300M` | ESM-C | 300M | V0.3 | lightest ESM-C |
-| `DecoderTCR-C-V0.3/6B.ckpt` | `DecoderTCR-C_6B` | ESM-C | 6B | V0.3 | larger variant (80 GB GPU) |
-| `DecoderTCR-ESM2-V0.1/650M_DecoderTCR.ckpt` | `DecoderTCR_650M` | ESM-2 | 650M | V0.1 | paper reproduction |
-| `DecoderTCR-ESM2-V0.1/3B_DecoderTCR.ckpt` | `DecoderTCR_3B` | ESM-2 | 3B | V0.1 | paper reproduction |
+### Core Architecture
 
-The V0.3 ESM-C line (`DecoderTCR-C`) is the current default; the V0.1 ESM-2 line is kept for
-paper reproduction and backward compatibility.
+The model follows the **ESM2** architecture, a deep Transformer encoder designed for protein sequences.
 
-## Usage
+#### Embedding Layer
+- Token embedding dimension: *d* (e.g., 1280)
+- Learned positional embeddings
+- Vocabulary includes:
+  - 20 standard amino acids
+  - Special tokens (mask, padding, BOS/EOS, unknown)
 
-Get the DecoderTCR code, then fetch weights (they download into these same nested paths):
+#### Transformer Stack
+- Number of layers: *L* (e.g., 33)
+- Hidden dimension: *d*
+- Number of attention heads: *h* (e.g., 20)
+- Multi-head self-attention:
+  - Full-sequence, bidirectional attention
+- Feed-forward network:
+  - Intermediate dimension ≈ 4× *d*
+  - Activation function: GELU
+- Layer normalization: Pre-LayerNorm
+- Residual connections around attention and feed-forward blocks
 
-```bash
-uv sync
-uv run python scripts/download_weights.py                       # all models
-uv run python scripts/download_weights.py -m DecoderTCR-C_600M   # just the default
-```
+### Continual Training Setup
 
-```python
-from DecoderTCR.utils.model_zoo import load
-model, n_layers = load()                       # default = DecoderTCR-C_600M
-```
+The model is initialized from a pretrained **ESM2 checkpoint** and further trained via continual pretraining with MLM objectives.
+
+
+### Released Models
+
+This repository hosts two model lines. The **V0.3 ESM-C** line (`DecoderTCR-C`) is the current default; the **V0.1 ESM-2** line is retained for paper reproduction.
+
+| Model | File | Backbone | Parameters | Layers | Hidden Dim | Attention Heads |
+| --- | --- | --- | --- | --- | --- | --- |
+| DecoderTCR-C 300M | `DecoderTCR-C-V0.3/300M.ckpt` | ESM-C | ~300M | 30 | 960 | 15 |
+| DecoderTCR-C 600M (default) | `DecoderTCR-C-V0.3/600M.ckpt` | ESM-C | ~600M | 36 | 1152 | 18 |
+| DecoderTCR-C 6B | `DecoderTCR-C-V0.3/6B.ckpt` | ESM-C | ~6B | 80 | 2560 | 40 |
+| DecoderTCR 650M | `DecoderTCR-ESM2-V0.1/650M_DecoderTCR.ckpt` | ESM-2 | ~650M | 33 | 1280 | 20 |
+| DecoderTCR 3B | `DecoderTCR-ESM2-V0.1/3B_DecoderTCR.ckpt` | ESM-2 | ~3B | 36 | 2560 | 40 |
+
+### Model Card Authors
+
+Ben Lai
+
+### Primary Contact Email
+
+Ben Lai ben.lai@czbiohub.org
+To submit feature requests or report issues with the model, please open an issue on [the GitHub repository](https://github.com/czbiohub-chi/DecoderTCR).
+
+### System Requirements
+
+- Compute Requirements: GPU
+
+## Intended Use
+
+### Primary Use Cases
+
+The DecoderTCR models are designed for the following primary use cases:
+
+1. **TCR-pMHC Binding Prediction**: Predict the interaction between T-cell receptors (TCRs) and peptide-MHC complexes
+2. **Interaction Scoring**: Calculate interface energy scores for TCR-pMHC interactions
+3. **Sequence Analysis**: Analyze TCR sequences and their interactions with specific peptides
+4. **Immunology Research**: Support research in adaptive immunity, T-cell recognition, and antigen presentation
+
+The models are particularly useful for:
+- Identifying potential TCR-peptide binding pairs
+- Screening TCR sequences for specific antigen recognition
+- Understanding the molecular basis of T-cell recognition
+- Supporting vaccine design and immunotherapy development
+
+### Out-of-Scope or Unauthorized Use Cases
+
+Do not use the model for the following purposes:
+- Use that violates applicable laws, regulations (including trade compliance laws), or third party rights such as privacy or intellectual property rights
+- Any use that is prohibited by the [MIT license](https://github.com/czbiohub-chi/DecoderTCR/blob/main/LICENSE) and [Acceptable Use Policy](https://virtualcellmodels.cziscience.com/acceptable-use-policy).
+- Clinical diagnosis or treatment decisions without proper validation
+- Direct use in patient care without appropriate clinical validation and regulatory approval
+- Use for purposes that could cause harm to individuals or groups
+
+The models are research tools and should not be used as the sole basis for clinical or diagnostic decisions.
+
+## Training Data
+
+The models are trained with multi-component large-scale protein sequence databases. The training data consists of:
+
+- **TCR sequences**: Observerd T-cell Space(OTS) for paired $\alpha/\beta$ TCR sequences.
+- **Peptide-MHC sequences**: MHC Motif Atlas for peptide-MHC ligandomes and high confidence synthetic interactions via MixMHCpred predictions.
+- **Paired TCR-pMHC Interactions**: VDJdb for paired TCR-pMHC interaction data.
+
+
+## Continual Pre-training Strategy
+
+This model is trained using a **continual pre-training curriculum** that adapts a pretrained ESM2 backbone to new protein domains while preserving previously learned representations.
+
+### Overview
+
+Continual pre-training proceeds in **multiple stages**, each leveraging different data regimes and masking strategies:
+
+- Stage 1 emphasize **abundant marginal sequence data**, encouraging robust component-level representations.
+- Stage 2 incorporate **scarcer, structured, or interaction-rich data**, refining conditional dependencies without overwriting earlier knowledge.
+
+The architecture, tokenizer, and objective remain unchanged throughout training; only the data distribution and masking strategy evolve.
+
+### Stage 1: Component-Level Adaptation
+
+In the first stage, the model is further pretrained on large collections of unpaired or weakly structured protein sequences relevant to the target domain.
+
+- **Objective:** Masked Language Modeling (MLM)
+- **Masking:** Component- or region-aware masking schedules that upweight functionally relevant positions
+- **Purpose:**  
+  - Adapt the pretrained ESM2 representations to the target protein subspace  
+  - Learn domain-specific sequence statistics while retaining general protein knowledge
+
+This stage acts as a regularizer, anchoring learning in large-scale marginal data before introducing more complex dependencies.
+
+### Stage 2: Conditional / Interaction-Aware Refinement
+
+In subsequent stages, the model is continually trained on **structured or paired sequences** that encode higher-order dependencies (e.g., interactions between protein regions or components).
+
+- **Objective:** Masked Language Modeling (MLM)
+- **Masking:** Joint masking across interacting regions to encourage cross-context conditioning
+- **Purpose:**  
+  - Refine conditional relationships learned from limited paired data  
+  - Align representations across components without degrading Stage 1 task performance
+
+## Biases, Risks, and Limitations
+
+### Potential Biases
+
+- The model may reflect biases present in the training data, including:
+  - Overrepresentation of certain HLA alleles or peptide types
+  - Limited diversity in TCR sequences from specific populations
+  - Bias toward well-studied antigen systems
+- Certain TCR clonotypes or peptide types may be underrepresented in training data
+
+### Risks
+
+Areas of risk may include but are not limited to:
+- **Inaccurate predictions**: The model may produce incorrect binding predictions, especially for novel sequences or rare HLA-peptide combinations
+- **Overconfidence**: The model may assign high confidence to predictions that are actually uncertain
+- **Biological misinterpretation**: Users may misinterpret model outputs as definitive biological facts rather than predictions
+- **Clinical misuse**: Use in clinical settings without proper validation could lead to incorrect treatment decisions
+
+### Limitations
+
+- **Sequence length**: The model has limitations on maximum sequence length (typically ~1024 tokens)
+- **Novel sequences**: Performance may degrade on sequences very different from training data
+- **HLA diversity**: Limited training data for rare HLA alleles may affect prediction accuracy
+- **Context dependency**: The model may not capture all biological context (e.g., post-translational modifications, cellular environment)
+- **Computational requirements**: GPU is recommended for optimal performance
+
+### Caveats and Recommendations
+
+- **Review and validate outputs**: Always review and validate model predictions, especially for critical applications
+- **Experimental validation**: Model predictions should be validated experimentally before use in research or clinical contexts
+- **Uncertainty awareness**: Be aware that predictions are probabilistic and may have uncertainty
+- **Domain expertise**: Use the model in conjunction with domain expertise in immunology and T-cell biology
+- **Version tracking**: Keep track of which model version and checkpoint you are using
+
+We are committed to advancing the responsible development and use of artificial intelligence. Please follow our [Acceptable Use Policy](/acceptable-use-policy)  when engaging with our services.
+
+Should you have any security or privacy issues or questions related to the services, please reach out to our team at security@chanzuckerberg.com or privacy@chanzuckerberg.com respectively.
+
+## Acknowledgements
+
+This model builds upon:
+- **ESM-2** by Meta AI (Facebook Research) for the ESM-2 base protein language model
+- **ESM-C** released by Chan Zuckerberg Biohub (https://github.com/Biohub/esm) for the ESM-C base protein language model
+- The broader computational biology and immunology research communities
+
+Special thanks to the developers and contributors of the ESM models and the open-source tools that made this work possible.
 
 ## License
 
-MIT (see [`LICENSE`](LICENSE)). The bundled backbones are also MIT: ESM-2 (Meta) and the
-Chan Zuckerberg Biohub ESM-C release (https://github.com/Biohub/esm). The released checkpoints
-contain the full fine-tuned weights and are distributed under the MIT license.
+DecoderTCR code and all released weights are distributed under the [MIT license](https://github.com/czbiohub-chi/DecoderTCR/blob/main/LICENSE). The base backbones are likewise MIT: ESM-2 (Meta AI) and ESM-C (Chan Zuckerberg Biohub, https://github.com/Biohub/esm).