Initial CodonBERT Hugging Face port

.gitattributes

@@ -0,0 +1,35 @@
+*.7z filter=lfs diff=lfs merge=lfs -text
+*.arrow filter=lfs diff=lfs merge=lfs -text
+*.bin filter=lfs diff=lfs merge=lfs -text
+*.bz2 filter=lfs diff=lfs merge=lfs -text
+*.ckpt filter=lfs diff=lfs merge=lfs -text
+*.ftz filter=lfs diff=lfs merge=lfs -text
+*.gz filter=lfs diff=lfs merge=lfs -text
+*.h5 filter=lfs diff=lfs merge=lfs -text
+*.joblib filter=lfs diff=lfs merge=lfs -text
+*.lfs.* filter=lfs diff=lfs merge=lfs -text
+*.mlmodel filter=lfs diff=lfs merge=lfs -text
+*.model filter=lfs diff=lfs merge=lfs -text
+*.msgpack filter=lfs diff=lfs merge=lfs -text
+*.npy filter=lfs diff=lfs merge=lfs -text
+*.npz filter=lfs diff=lfs merge=lfs -text
+*.onnx filter=lfs diff=lfs merge=lfs -text
+*.ot filter=lfs diff=lfs merge=lfs -text
+*.parquet filter=lfs diff=lfs merge=lfs -text
+*.pb filter=lfs diff=lfs merge=lfs -text
+*.pickle filter=lfs diff=lfs merge=lfs -text
+*.pkl filter=lfs diff=lfs merge=lfs -text
+*.pt filter=lfs diff=lfs merge=lfs -text
+*.pth filter=lfs diff=lfs merge=lfs -text
+*.rar filter=lfs diff=lfs merge=lfs -text
+*.safetensors filter=lfs diff=lfs merge=lfs -text
+saved_model/**/* filter=lfs diff=lfs merge=lfs -text
+*.tar.* filter=lfs diff=lfs merge=lfs -text
+*.tar filter=lfs diff=lfs merge=lfs -text
+*.tflite filter=lfs diff=lfs merge=lfs -text
+*.tgz filter=lfs diff=lfs merge=lfs -text
+*.wasm filter=lfs diff=lfs merge=lfs -text
+*.xz 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

README.md

@@ -0,0 +1,239 @@
+---
+library_name: transformers
+tags:
+- RNA
+- mRNA
+- codon
+- language-model
+license: other
+---
+
+# CodonBERT
+
+Minimal HuggingFace port of the **CodonBERT** checkpoint from
+[CodonBERT](https://github.com/Sanofi-Public/CodonBERT) -- a BERT-based RNA
+language model pretrained on codon-level representations of more than 10 million
+mRNA coding sequences.
+
+## Architecture
+
+| Parameter | Value |
+|---|---|
+| Layers | 12 |
+| Attention heads | 12 |
+| Embedding dimension | 768 |
+| FFN hidden dimension | 3072 (GELU) |
+| Vocabulary size | 69 (5 special + 61 sense codons + 3 stop codons) |
+| Positional encoding | Learned absolute |
+| Normalization | LayerNorm (epsilon=1e-12) |
+| Architecture | Standard post-LN BERT Transformer |
+| Max sequence length | 1024 encoded tokens (up to 1022 codons / 3066 nt for one sequence) |
+
+### Vocabulary
+
+The tokenizer operates at the codon level. Unlike the original tokenizer, this port
+accepts raw nucleotide strings and performs codon splitting automatically.
+The 64 codons cover all combinations of {A, U, G, C}^3 in RNA space, including
+the three stop codons.
+Special tokens follow standard BERT convention: `[PAD]=0`, `[UNK]=1`,
+`[CLS]=2`, `[SEP]=3`, `[MASK]=4`.
+
+## Pretraining
+
+- **Objective:** 15% masked language modeling (MLM) plus paired-sequence
+  taxonomy prediction (STP)
+- **Data:** >10 million mRNA coding sequences from mammals, bacteria, human
+  viruses, and yeast
+- **Source checkpoint:** `pytorch_model.bin` from the official
+  [Sanofi CodonBERT archive](https://cdn.prod.accelerator.sanofi/llm/CodonBERT.zip),
+  mirrored as [`lhallee/CodonBERT`](https://huggingface.co/lhallee/CodonBERT)
+
+### Checkpoint selection
+
+There is a single publicly released checkpoint from the original authors. The backbone
+weights (`bert.*` prefix) and complete MLM prediction head are mapped directly.
+Only the two-tensor STP/sequence-relationship head is discarded.
+
+## Parity Verification
+
+All verified on GPU with PyTorch 2.7.1 / CUDA 12.9:
+
+- **Hidden states (eager):** all 13 levels match the original under
+  `torch.allclose(atol=1e-5, rtol=1e-5)` (observed max abs diff `1.15e-5`
+  on a padded three-sequence batch)
+- **MLM logits and loss:** converted logits match original
+  `BertForPreTraining` logits under the same tolerance (observed max abs diff
+  `1.13e-5`; mixed `0`/`-100` label loss difference `2.87e-6`)
+- **SDPA (evaluation):** final hidden states agree with eager FP32 to
+  `3.58e-6` max abs difference at non-padding positions
+- **Flash attention 2 (evaluation):** verified against eager BF16 at
+  non-padding positions (all-layer differences up to `0.25`, expected BF16
+  accumulation across 12 layers)
+
+## Related Models
+
+See the full [CodonBERT collection](https://huggingface.co/collections/Taykhoom/codonbert).
+
+| Model | Parameters | Notes |
+|---|---:|---|
+| **[CodonBERT](https://huggingface.co/Taykhoom/CodonBERT)** | 87.1M | This model |
+
+## Usage
+
+CodonBERT operates on CDS sequences. The tokenizer handles T->U conversion and codon
+splitting automatically — pass raw nucleotide strings directly.
+
+### Embedding generation
+
+```python
+import torch
+from transformers import AutoTokenizer, AutoModel
+
+tokenizer = AutoTokenizer.from_pretrained("Taykhoom/CodonBERT", trust_remote_code=True)
+model = AutoModel.from_pretrained("Taykhoom/CodonBERT", trust_remote_code=True)
+model.eval()
+
+# Raw CDS nucleotide strings — T or U both accepted
+cds_sequences = ["ATGAAAGGCCCTTAA", "ATGTTTGGG"]
+
+enc = tokenizer(cds_sequences, return_tensors="pt", padding=True)
+
+with torch.no_grad():
+    out = model(**enc)
+
+cls_emb   = out.last_hidden_state[:, 0, :]  # (batch, 768) -- CLS token
+content_mask = enc["attention_mask"].clone()
+content_mask[:, 0] = 0  # exclude CLS
+content_mask[torch.arange(content_mask.size(0)), enc["attention_mask"].sum(1) - 1] = 0
+mean_emb = (out.last_hidden_state * content_mask.unsqueeze(-1)).sum(1) / \
+           content_mask.sum(1, keepdim=True)  # mean over codons only
+
+# Intermediate layers
+out_all = model(**enc, output_hidden_states=True)
+layer6_emb = out_all.hidden_states[6]  # (batch, seq_len, 768)
+```
+
+### CDS-aware encoding (full mRNA input)
+
+For full mRNA sequences where the CDS region must be extracted first:
+
+```python
+import numpy as np
+
+# cds: binary array with 1 at the first nucleotide of each codon
+enc, chunk_counts = tokenizer.batch_encode_with_cds(
+    mrna_sequences,
+    cds_tracks,       # list of numpy arrays
+    return_tensors="pt",
+    padding=True,
+)
+with torch.no_grad():
+    out = model(**enc)
+```
+
+### Faster attention backends
+
+```python
+# Evaluation/inference only; use eager for training (see Implementation Notes).
+model_sdpa = AutoModel.from_pretrained(
+    "Taykhoom/CodonBERT", trust_remote_code=True, attn_implementation="sdpa"
+)
+model_flash = AutoModel.from_pretrained(
+    "Taykhoom/CodonBERT", trust_remote_code=True,
+    attn_implementation="flash_attention_2", dtype=torch.bfloat16
+)
+```
+
+### MLM logits
+
+```python
+from transformers import AutoModelForMaskedLM
+
+model_mlm = AutoModelForMaskedLM.from_pretrained("Taykhoom/CodonBERT", trust_remote_code=True)
+model_mlm.eval()
+
+seq = "AUG [MASK] GGG"
+enc = tokenizer(seq, return_tensors="pt")
+with torch.no_grad():
+    logits = model_mlm(**enc).logits  # (1, seq_len, 69)
+```
+
+The MLM prediction transform (dense + GELU + LayerNorm), decoder weight, and output
+bias are all converted from the original checkpoint. The decoder tensor is initialized
+exactly from the word embedding tensor. In this adapter it is a separate parameter,
+not a runtime-tied alias; see the fine-tuning limitations below.
+
+### Fine-tuning
+
+For sequence-level tasks, use the CLS token embedding as input to a
+classification/regression head. Train this checkpoint with the `eager` backend:
+the current SDPA and Flash Attention 2 paths do not apply the configured `0.1`
+attention-probability dropout during training.
+
+The input embedding and MLM decoder start value-identical but are not storage-tied.
+`tie_weights()` does not tie them, and `resize_token_embeddings()` leaves the decoder
+at 69 outputs; vocabulary resizing is therefore unsupported. The model forward API
+accepts `input_ids`, `attention_mask`, and `token_type_ids`, but not the stock BERT
+`inputs_embeds`, `position_ids`, or `head_mask` arguments.
+
+## Implementation Notes
+
+Two key differences from the original CodonBERT release:
+
+**1. Integrated codon tokenization.** The original repository requires users to
+manually pre-process sequences into space-separated codons before passing them to
+the tokenizer. This port ships `CodonBertTokenizer`, a `BertTokenizer` subclass
+whose `_tokenize` method automatically normalizes sequences (T->U, uppercase) and
+splits them into codon 3-mers. Users can pass raw nucleotide strings directly:
+`tokenizer("AUGAAAGGG")` works without any pre-processing. A
+`batch_encode_with_cds(sequences, cds_tracks)` method handles full mRNA input with
+CDS extraction and codon-boundary-aligned chunking.
+
+**2. SDPA and Flash Attention 2 support.** This port inherits from
+[`Taykhoom/BERT-updated`](https://huggingface.co/Taykhoom/BERT-updated),
+a minimal BERT re-implementation with all three backends (`eager`, `sdpa`,
+`flash_attention_2`). Evaluation parity against the original eager implementation
+is verified at every layer. For training, use eager as noted above. Eager attentions
+requested in training mode are post-dropout tensors, so their rows do not sum to one.
+Although the shared config accepts `hidden_act`, the backend always applies GELU;
+this checkpoint is configured for GELU and is unaffected.
+
+## Citation
+
+```bibtex
+@article{li2024_codonbert,
+  title   = {{CodonBERT} large language model for {mRNA} vaccines},
+  author  = {Li, Sizhen and Moayedpour, Saeed and Li, Ruijiang and Bailey, Michael and Riahi, Saleh and Kogler-Anele, Lorenzo and Miladi, Milad and Miner, Jacob and Pertuy, Fabien and Zheng, Dinghai and Wang, Jun and Balsubramani, Akshay and Tran, Khang and Zacharia, Minnie and Wu, Monica and Gu, Xiaobo and Clinton, Ryan and Asquith, Carla and Skaleski, Joseph and Boeglin, Lianne and Chivukula, Sudha and Dias, Anusha and Strugnell, Tod and Ulloa Montoya, Fernando and Agarwal, Vikram and Bar-Joseph, Ziv and Jager, Sven},
+  journal = {Genome Research},
+  volume  = {34},
+  number  = {7},
+  pages   = {1027--1035},
+  year    = {2024},
+  doi     = {10.1101/gr.278870.123}
+}
+```
+
+## Credits
+
+Original model and code by Li et al. Source: [GitHub](https://github.com/Sanofi-Public/CodonBERT).
+The HF conversion code was authored primarily by [Claude Code](https://claude.ai/code)
+and reviewed manually by Taykhoom Dalal.
+
+## License
+
+Academic/non-commercial use only, following the original
+[artifact license](https://github.com/Sanofi-Public/CodonBERT/blob/master/ARTIFACT_LICENSE.md):
+
+Permission is hereby granted, free of charge, for academic research purposes only
+and for non-commercial use only, to any person from an academic research or non-profit
+organization obtaining a copy of these models, software, datasets and/or algorithms
+(including, but not limited to, machine-learning model code, trained model weights,
+inference-enabling code, training-enabling code, fine-tuning enabling code and other
+elements) and/or associated documentation files (collectively the "Materials") to use,
+copy, modify, or merge the Materials, subject to the following conditions: this IP
+License Notice shall be included in all copies of the Materials or of substantial
+portions of the Materials.
+For purposes of this notice, "non-commercial use" excludes uses foreseeably resulting
+in a commercial benefit or monetary gain. All other rights are reserved. The Materials
+are provided "as is," without warranty of any kind, express or implied, including the
+warranties of noninfringement.

config.json

@@ -0,0 +1,26 @@
+{
+  "architectures": [
+    "BertForMaskedLM"
+  ],
+  "auto_map": {
+    "AutoConfig": "Taykhoom/BERT-updated--configuration_bert_updated.BertUpdatedConfig",
+    "AutoModel": "Taykhoom/BERT-updated--modeling_bert.BertModel",
+    "AutoModelForMaskedLM": "Taykhoom/BERT-updated--modeling_bert.BertForMaskedLM"
+  },
+  "attention_probs_dropout_prob": 0.1,
+  "hidden_act": "gelu",
+  "hidden_dropout_prob": 0.1,
+  "hidden_size": 768,
+  "initializer_range": 0.02,
+  "intermediate_size": 3072,
+  "layer_norm_eps": 1e-12,
+  "max_position_embeddings": 1024,
+  "model_type": "bert_updated",
+  "num_attention_heads": 12,
+  "num_hidden_layers": 12,
+  "pad_token_id": 0,
+  "type_vocab_size": 2,
+  "vocab_size": 69,
+  "transformers_version": "4.57.6",
+  "model_max_length": 1024
+}

model.safetensors

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:534daf559ebc571a3158fed93fe140094f39c62eee74edc7bc319f66d6d89957
+size 348554708

special_tokens_map.json

@@ -0,0 +1,7 @@
+{
+  "cls_token": "[CLS]",
+  "mask_token": "[MASK]",
+  "pad_token": "[PAD]",
+  "sep_token": "[SEP]",
+  "unk_token": "[UNK]"
+}

tokenization_codonbert.py

@@ -0,0 +1,91 @@
+import warnings
+import numpy as np
+from transformers import BertTokenizer
+
+
+class CodonBertTokenizer(BertTokenizer):
+    """BertTokenizer that auto-converts nucleotide sequences to codon-level tokens.
+
+    Raw nucleotide input is normalized (T->U, uppercase, whitespace stripped),
+    then split into non-overlapping 3-mer codons before vocab lookup. Trailing
+    1-2 nucleotides that do not form a complete codon are dropped.
+
+    eos_token is aliased to sep_token ("[SEP]") so that pooling code that
+    excludes both CLS and EOS/SEP positions works correctly.
+
+    Standard usage (raw nucleotides):
+        tokenizer("AUGAAAGGG")
+        tokenizer(["AUGAAAGGG", "AUGUUUCCC"], return_tensors="pt", padding=True)
+
+    CDS-aware usage (full mRNA + CDS track -> extract CDS, chunk, encode):
+        tokenizer.batch_encode_with_cds(
+            ["NNNATGAAAGGGNN"],
+            cds=[np.array([0,0,0,1,0,0,1,0,0,1,0,0,0,0])],
+            return_tensors="pt",
+            padding=True,
+        )
+
+    Works with compare_minimal_vs_mm.py --use_cds out of the box.
+    """
+
+    def __init__(self, *args, **kwargs):
+        kwargs.setdefault("eos_token", "[SEP]")
+        super().__init__(*args, **kwargs)
+
+    def _tokenize(self, text, split_special_tokens=False):
+        seq = "".join(text.split()).upper().replace("T", "U")
+        n = len(seq) - len(seq) % 3
+        return [seq[i:i + 3] for i in range(0, n, 3)]
+
+    @staticmethod
+    def _extract_cds(sequence, cds):
+        if sum(cds) == 0:
+            warnings.warn("No CDS found. Returning truncated sequence.")
+            n = len(sequence) - len(sequence) % 3
+            return sequence[:n]
+        first = int(np.argmax(cds == 1))
+        last = int(len(cds) - 1 - np.argmax(np.flip(cds) == 1)) + 2
+        proposed = sequence[first:last + 1]
+        if len(proposed) % 3 != 0:
+            warnings.warn("Irregular CDS. Returning truncated sequence.")
+            return proposed[:-(len(proposed) % 3)]
+        return proposed
+
+    def batch_encode_with_cds(self, sequences, cds_tracks, max_length=None, **kwargs):
+        """Encode a batch of raw mRNA sequences using CDS-aware preprocessing.
+
+        Args:
+            sequences: List of raw nucleotide strings.
+            cds_tracks: List of numpy arrays (one per sequence). Non-zero values
+                mark the first nucleotide of each codon in the CDS region.
+            max_length: Max content codon-tokens per chunk (special tokens NOT
+                counted). Defaults to model_max_length - 2. This matches the
+                convention in compare_minimal_vs_mm.py where max_length is
+                already adjusted for special tokens.
+            **kwargs: Forwarded to batch_encode_plus (e.g. return_tensors, padding).
+
+        Returns:
+            (BatchEncoding, chunk_counts): chunk_counts[i] is the number of
+            chunks produced from sequence i.
+        """
+        budget_codons = max_length or (self.model_max_length - 2)
+        budget_nt = budget_codons * 3
+
+        all_strings = []
+        chunk_counts = []
+
+        for seq, cds in zip(sequences, cds_tracks):
+            seq = seq.replace("T", "U").replace("t", "u").upper()
+            cds_seq = self._extract_cds(seq, np.asarray(cds))
+            n = len(cds_seq)
+            chunks = []
+            for i in range(0, max(n, 1), budget_nt):
+                chunk = cds_seq[i:i + budget_nt]
+                chunk = chunk[:len(chunk) - len(chunk) % 3]
+                if chunk:
+                    chunks.append(chunk)
+            all_strings.extend(chunks or [""])
+            chunk_counts.append(len(chunks) or 1)
+
+        enc = self.batch_encode_plus(all_strings, **kwargs)
+        return enc, chunk_counts

tokenizer_config.json

@@ -0,0 +1,58 @@
+{
+  "added_tokens_decoder": {
+    "0": {
+      "content": "[PAD]",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "1": {
+      "content": "[UNK]",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "2": {
+      "content": "[CLS]",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "3": {
+      "content": "[SEP]",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "4": {
+      "content": "[MASK]",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    }
+  },
+  "cls_token": "[CLS]",
+  "do_basic_tokenize": true,
+  "do_lower_case": false,
+  "mask_token": "[MASK]",
+  "model_max_length": 1024,
+  "pad_token": "[PAD]",
+  "sep_token": "[SEP]",
+  "strip_accents": null,
+  "tokenize_chinese_chars": false,
+  "tokenizer_class": "CodonBertTokenizer",
+  "auto_map": {
+    "AutoTokenizer": ["tokenization_codonbert.CodonBertTokenizer", null]
+  },
+  "unk_token": "[UNK]"
+}

vocab.txt

@@ -0,0 +1,69 @@
+[PAD]
+[UNK]
+[CLS]
+[SEP]
+[MASK]
+AAA
+AAU
+AAG
+AAC
+AUA
+AUU
+AUG
+AUC
+AGA
+AGU
+AGG
+AGC
+ACA
+ACU
+ACG
+ACC
+UAA
+UAU
+UAG
+UAC
+UUA
+UUU
+UUG
+UUC
+UGA
+UGU
+UGG
+UGC
+UCA
+UCU
+UCG
+UCC
+GAA
+GAU
+GAG
+GAC
+GUA
+GUU
+GUG
+GUC
+GGA
+GGU
+GGG
+GGC
+GCA
+GCU
+GCG
+GCC
+CAA
+CAU
+CAG
+CAC
+CUA
+CUU
+CUG
+CUC
+CGA
+CGU
+CGG
+CGC
+CCA
+CCU
+CCG
+CCC