OmniTCR

OmniTCR is a component-aware autoregressive foundation model for learning relationships among peptide epitopes, major histocompatibility complex (MHC) molecules, T-cell receptor alpha chains (TRA) and T-cell receptor beta chains (TRB). The repository provides inference checkpoints for molecular-association prediction, candidate cancer-associated TCR analysis and pMHC-conditioned TRB CDR3 sequence generation.

The checkpoints accompany the manuscript:

OmniTCR: a generative foundation model from molecular recognition to T-cell receptor repertoires
[Paper or preprint URL to be added]

Model overview

OmniTCR uses a decoder-only Transformer pretrained by causal next-token prediction on 328,232,215 human immune-sequence records. Its component-aware format permits training on records containing different combinations of peptide, MHC, TRA and TRB components.

Property Value
Parameters approximately 113 million
Transformer layers 12
Attention heads 12
Hidden dimension 768
Feed-forward dimension 3,072
Vocabulary size 34
Backbone objective causal next-token prediction

The special component tokens are [EPI], [HLA], [TRA] and [TRB]. Amino-acid sequences are represented in their native N-to-C-terminal direction.

Available checkpoints

OmniTCR uses a common pretrained backbone followed by separately adapted task models.

Directory Model Intended inference task
OmniTCR(Base) OmniTCR(Base) Sequence likelihood and perplexity scoring
OmniTCR(FFT)_PM OmniTCR(FFT)-PM Peptide-MHC association prediction
OmniTCR(FFT)_PT OmniTCR(FFT)-PT Peptide-TRB association prediction
OmniTCR(FFT)_PMT OmniTCR(FFT)-PMT Peptide-MHC-TRB association prediction
OmniTCR(FFT)_PMAB OmniTCR(FFT)-PMAB Peptide-MHC-TRA-TRB paired-chain association prediction
OmniTCR(FFT)_C_A OmniTCR(FFT)-C_A Candidate cancer-associated receptor and repertoire analysis
OmniTCR(SFT) OmniTCR(SFT) pMHC-conditioned TRB CDR3 candidate generation

The FFT files contain complete fine-tuned state dictionaries. The public inference code constructs TCRLlamaForBinaryClassification using the shared Base configuration and then loads the selected FFT weights. The FFT directories are therefore not independent AutoModel.from_pretrained() repositories.

Repository structure

OmniTCR/
β”œβ”€β”€ README.md
β”œβ”€β”€ LICENSE
β”œβ”€β”€ OmniTCR(Base)/
β”‚   β”œβ”€β”€ model.safetensors
β”‚   β”œβ”€β”€ config.json
β”‚   β”œβ”€β”€ tokenizer_config.json
β”‚   β”œβ”€β”€ special_tokens_map.json
β”‚   └── vocab.json
β”œβ”€β”€ OmniTCR(FFT)_C_A/
β”‚   └── model.safetensors
β”œβ”€β”€ OmniTCR(FFT)_PM/
β”‚   └── model.safetensors
β”œβ”€β”€ OmniTCR(FFT)_PMAB/
β”‚   └── model.safetensors
β”œβ”€β”€ OmniTCR(FFT)_PMT/
β”‚   └── model.safetensors
β”œβ”€β”€ OmniTCR(FFT)_PT/
β”‚   └── model.safetensors
└── OmniTCR(SFT)/
    β”œβ”€β”€ model.safetensors
    β”œβ”€β”€ config.json
    β”œβ”€β”€ generation_config.json
    β”œβ”€β”€ tokenizer_config.json
    β”œβ”€β”€ special_tokens_map.json
    └── vocab.json

Installation

Clone the inference-code repository and install its dependencies:

git clone [GITHUB_REPOSITORY_URL]
cd OmniTCR
pip install -r requirements.txt

The principal dependencies are PyTorch, Transformers, Safetensors, pandas and NumPy. Exact tested versions are provided in the inference-code repository.

Downloading the checkpoints

Download the complete checkpoint repository:

from huggingface_hub import snapshot_download

checkpoint_root = snapshot_download(
    repo_id="loveCloud/OmniTCR",
    revision="main",
)

print(checkpoint_root)

For reproducible analyses, replace main with the release tag or commit identifier reported in the manuscript.

To download only one checkpoint family:

from huggingface_hub import snapshot_download

checkpoint_root = snapshot_download(
    repo_id="loveCloud/OmniTCR",
    revision="main",
    allow_patterns=[
        "OmniTCR(Base)/*",
        "OmniTCR(FFT)_PMT/*",
    ],
)

Loading an FFT checkpoint

The example below illustrates the shared-Base loading procedure. The model implementation is provided in the inference-code repository.

from pathlib import Path

import torch
from safetensors.torch import load_file

from finetune.src.model.one_layer_full import TCRLlamaForBinaryClassification


checkpoint_root = Path(checkpoint_root)
base_path = checkpoint_root / "OmniTCR(Base)"
fft_path = checkpoint_root / "OmniTCR(FFT)_PMT" / "model.safetensors"

model = TCRLlamaForBinaryClassification(
    pretrained_model_path=str(base_path)
)

state_dict = load_file(str(fft_path), device="cpu")
state_dict = {
    key[7:] if key.startswith("module.") else key: value
    for key, value in state_dict.items()
}

model.load_state_dict(state_dict, strict=True)
model.eval()

device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
model.to(device)

For FFT classification models, the reported positive-class probability is:

probability = torch.softmax(logits, dim=1)[:, 1]

Input construction, tokenization, batching and task-specific prediction commands are documented in the inference-code repository.

Selected evaluation results

Selected external-benchmark results reported in the accompanying manuscript are summarized below. Consult the manuscript and Supplementary Information for dataset construction, overlap removal, comparator settings and uncertainty estimates.

Evaluation AUROC AUPRC
External peptide-MHC prediction 0.9727 0.9592
External peptide-TRB prediction 0.9325 0.8736
External peptide-MHC-TRB prediction 0.9235 0.7949
External paired TRA-TRB prediction 0.9584 0.9083

Intended uses

OmniTCR is intended for research applications including:

  • computational prioritization of peptide-MHC and TCR associations;
  • analysis of candidate cancer-associated receptor and repertoire patterns;
  • representation extraction from immune-sequence records;
  • pMHC-conditioned generation and computational prioritization of TRB CDR3 candidates.

Limitations and out-of-scope uses

  • OmniTCR predictions do not establish biochemical binding, antigen specificity or T-cell function without experimental validation.
  • Repertoire-level cancer associations are not equivalent to experimentally established tumour specificity.
  • Generated TRB CDR3 sequences are computational candidates and require binding, functional and off-target validation.
  • Performance may depend on the representation of peptides, HLA alleles, receptor chains, populations and experimental sources in the training data.
  • The checkpoints are not intended for clinical diagnosis, treatment selection or other direct clinical decision-making.

Reproducibility

To reproduce manuscript inference results, use the checkpoint release tag or commit identifier specified with the paper rather than the mutable main branch. Each published release should report:

  • the checkpoint version and commit identifier;
  • the tokenizer version;
  • the positive-class index and fixed validation-derived threshold;
  • task-specific input formatting and maximum sequence length;
  • software and hardware dependencies;
  • SHA-256 checksums for all weight files.

License

[Add the model-weight license and the inference-code license before public release. Confirm the choice with all authors and your institution.]

Citation

Citation information will be added when the manuscript or preprint becomes publicly available.

Contact

[Corresponding author name and institutional email]

Downloads last month

-

Downloads are not tracked for this model. How to track
Inference Providers NEW
This model isn't deployed by any Inference Provider. πŸ™‹ Ask for provider support