Initial ERNIE-RNA 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,154 @@
+---
+library_name: transformers
+tags:
+- RNA
+- language-model
+license: mit
+---
+
+# ERNIE-RNA
+
+ERNIE-RNA is an RNA-specific large language model that incorporates RNA base-pairing potential as a recurrent 2D structural bias into each attention layer, enabling the model to capture secondary structure information during pretraining.
+
+## Architecture
+
+| Parameter | Value |
+|---|---|
+| Layers | 12 |
+| Attention heads | 12 |
+| Embedding dimension | 768 |
+| FFN hidden dimension | 3072 (GELU) |
+| Vocabulary size | 25 |
+| Positional encoding | Sinusoidal (fairseq-style) |
+| Normalization | LayerNorm (embedding and post-residual) |
+| Architecture | Post-LN Transformer with recurrent 2D RNA pairing bias |
+| Max sequence length | 1024 |
+
+### Vocabulary
+
+| Token | ID | Notes |
+|---|---|---|
+| `<cls>` | 0 | Prepended to every sequence |
+| `<pad>` | 1 | Padding token |
+| `<eos>` | 2 | Appended to every sequence |
+| `<unk>` | 3 | Unknown token |
+| G | 4 | |
+| A | 5 | |
+| U | 6 | T is silently mapped to U during tokenization |
+| C | 7 | |
+| N | 8 | Ambiguous nucleotide |
+| Y, R, S, K, W, M, D, H, V, B | 9-18 | IUPAC ambiguity codes |
+| X | 19 | Unknown/ambiguous nucleotide |
+| I | 20 | Inosine |
+| madeupword0000-madeupword0002 | 21-23 | Placeholder tokens from the original fairseq vocabulary |
+| `<mask>` | 24 | MLM mask token |
+
+### 2D RNA Pairing Bias
+
+ERNIE-RNA computes a pairwise RNA base-pairing potential matrix from the input sequence at the start of each forward pass. This matrix (shape `[B, T, T, 1]`) is projected to `[B, H, T, T]` via a 2-layer MLP (1 -> 6 -> H, with GELU) and added to the attention logits in the first layer. The pre-softmax attention scores then become the updated 2D bias for the next layer, creating a recurrent structural information pathway across all 12 transformer layers.
+
+Base-pairing scores: A-U = 2.0, G-C = 3.0, G-U wobble = 0.8.
+
+## Pretraining
+
+- **Objective:** Masked language modeling (MLM) on RNA sequences
+- **Data:** RNAcentral (non-redundant RNA sequences)
+- **Source checkpoint:** `ERNIE-RNA_pretrain.pt`
+
+### Checkpoint selection
+
+Single pretrained checkpoint from the original repository. Used as-is; no fine-tuned variants are included in this release.
+
+## Parity Verification
+
+Hidden-state representations verified identical (max abs diff = 1.82e-06) to the original
+implementation at all 13 representation levels (embedding + 12 transformer layers).
+Verified on GPU with PyTorch 2.7 / CUDA 12.
+
+Only `attn_implementation="eager"` is supported (see Implementation Notes).
+
+## Related Models
+
+See the full [ERNIE-RNA collection](https://huggingface.co/collections/Taykhoom/ernie-rna-6a20c1a8ea56c00a74e2dd93).
+
+| Model | Notes |
+|---|---|
+| **[Taykhoom/ERNIE-RNA](https://huggingface.co/Taykhoom/ERNIE-RNA)** | **Pretrained model (this model)** |
+| [Taykhoom/ERNIE-RNA-SS](https://huggingface.co/Taykhoom/ERNIE-RNA-SS) | SS fine-tuned (RNA3DB), backbone only |
+| [Taykhoom/ERNIE-RNA-MRL](https://huggingface.co/Taykhoom/ERNIE-RNA-MRL) | UTR MRL fine-tuned, backbone only |
+
+## Usage
+
+### Embedding generation
+
+```python
+import torch
+from transformers import AutoTokenizer, AutoModel
+
+tokenizer = AutoTokenizer.from_pretrained("Taykhoom/ERNIE-RNA", trust_remote_code=True)
+model = AutoModel.from_pretrained("Taykhoom/ERNIE-RNA", trust_remote_code=True)
+model.eval()
+
+sequences = ["AUGCAUGCAUGC", "GGGGCCCCGGGG"]
+enc = tokenizer(sequences, return_tensors="pt", padding=True)
+
+with torch.no_grad():
+    out = model(**enc)
+
+cls_emb   = out.last_hidden_state[:, 0, :]   # (batch, 768) -- CLS token
+token_emb = out.last_hidden_state             # (batch, seq_len, 768)
+
+# Intermediate layers
+out_all = model(**enc, output_hidden_states=True)
+layer6_emb = out_all.hidden_states[6]         # (batch, seq_len, 768)
+```
+
+### MLM logits
+
+```python
+import torch
+from transformers import AutoTokenizer, AutoModelForMaskedLM
+
+tokenizer = AutoTokenizer.from_pretrained("Taykhoom/ERNIE-RNA", trust_remote_code=True)
+model = AutoModelForMaskedLM.from_pretrained("Taykhoom/ERNIE-RNA", trust_remote_code=True)
+model.eval()
+
+enc = tokenizer(["AUG<mask>AUG"], return_tensors="pt")
+with torch.no_grad():
+    logits = model(**enc).logits   # (1, seq_len, 25)
+```
+
+### Fine-tuning
+
+Use the CLS token embedding (`last_hidden_state[:, 0, :]`) as input to a prediction head for sequence-level tasks. For token-level tasks, use `last_hidden_state` directly.
+
+## Implementation Notes
+
+ERNIE-RNA's recurrent 2D bias is updated from the pre-softmax attention scores at every layer (the raw QK logits become the bias input for the next layer). Fused attention kernels (SDPA, FlashAttention) do not expose pre-softmax scores, so they cannot maintain this recurrent pathway. Only `attn_implementation="eager"` is supported; requesting `sdpa` or `flash_attention_2` raises a `ValueError`.
+
+The `twod_proj` MLP is always run in float32 (matching the original) regardless of the model's compute dtype.
+
+## Citation
+
+```bibtex
+@article{yin2025_ernierna,
+  title   = {{ERNIE-RNA}: an {RNA} language model with structure-enhanced representations},
+  author  = {Yin, Weijie and Zhang, Zhaoyu and Zhang, Shuo and He, Liang and Zhang, Ruiyang and Jiang, Rui and Liu, Gan and Wang, Jingyi and Zhang, Xuegong and Qin, Tao and Xie, Zhen},
+  journal = {Nature Communications},
+  volume  = {16},
+  number  = {1},
+  pages   = {8407},
+  year    = {2025},
+  doi     = {10.1038/s41467-025-64972-0}
+}
+```
+
+## Credits
+
+Original model and code by Yin et al. Source: [GitHub](https://github.com/Bruce-ywj/ERNIE-RNA).
+The HF conversion code was authored primarily by [Claude Code](https://claude.ai/code)
+and reviewed manually by Taykhoom Dalal.
+
+## License
+
+MIT, following the original repository.

config.json

@@ -0,0 +1,26 @@
+{
+  "activation_dropout": 0.0,
+  "activation_fn": "gelu",
+  "architectures": [
+    "ErnieRNAForMaskedLM"
+  ],
+  "attention_dropout": 0.1,
+  "attention_heads": 12,
+  "auto_map": {
+    "AutoConfig": "configuration_ernierna.ErnieRNAConfig",
+    "AutoModel": "modeling_ernierna.ErnieRNAModel",
+    "AutoModelForMaskedLM": "modeling_ernierna.ErnieRNAForMaskedLM"
+  },
+  "dropout": 0.1,
+  "embed_dim": 768,
+  "ffn_embed_dim": 3072,
+  "mask_idx": 24,
+  "max_positions": 1024,
+  "model_max_length": 1024,
+  "model_type": "ernie_rna",
+  "num_layers": 12,
+  "num_segments": 2,
+  "padding_idx": 1,
+  "transformers_version": "4.57.6",
+  "vocab_size": 25
+}

configuration_ernierna.py

@@ -0,0 +1,45 @@
+from transformers import PretrainedConfig
+
+
+class ErnieRNAConfig(PretrainedConfig):
+    model_type = "ernie_rna"
+
+    auto_map = {
+        "AutoConfig": "configuration_ernierna.ErnieRNAConfig",
+        "AutoModel": "modeling_ernierna.ErnieRNAModel",
+        "AutoModelForMaskedLM": "modeling_ernierna.ErnieRNAForMaskedLM",
+    }
+
+    def __init__(
+        self,
+        vocab_size=25,
+        num_layers=12,
+        embed_dim=768,
+        ffn_embed_dim=3072,
+        attention_heads=12,
+        dropout=0.1,
+        attention_dropout=0.1,
+        activation_dropout=0.0,
+        activation_fn="gelu",
+        max_positions=1024,
+        padding_idx=1,
+        mask_idx=24,
+        num_segments=2,
+        model_max_length=1024,
+        **kwargs,
+    ):
+        super().__init__(padding_idx=padding_idx, **kwargs)
+        self.auto_map = dict(type(self).auto_map)
+        self.vocab_size = vocab_size
+        self.num_layers = num_layers
+        self.embed_dim = embed_dim
+        self.ffn_embed_dim = ffn_embed_dim
+        self.attention_heads = attention_heads
+        self.dropout = dropout
+        self.attention_dropout = attention_dropout
+        self.activation_dropout = activation_dropout
+        self.activation_fn = activation_fn
+        self.max_positions = max_positions
+        self.mask_idx = mask_idx
+        self.num_segments = num_segments
+        self.model_max_length = model_max_length

modeling_ernierna.py

@@ -0,0 +1,322 @@
+import math
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+
+from transformers import PreTrainedModel
+from transformers.activations import ACT2FN
+from transformers.modeling_outputs import BaseModelOutput, MaskedLMOutput
+
+try:
+    from .configuration_ernierna import ErnieRNAConfig
+except ImportError:
+    from configuration_ernierna import ErnieRNAConfig
+
+
+class ErnieRNASinusoidalPositionalEmbedding(nn.Module):
+    def __init__(self, num_positions, embed_dim, padding_idx):
+        super().__init__()
+        self.embedding_dim = embed_dim
+        self.padding_idx = padding_idx
+        # Table size: need indices up to padding_idx + 1 + num_positions
+        table_size = padding_idx + 1 + num_positions
+        self.register_buffer("weights", self._get_embedding(table_size, embed_dim, padding_idx))
+
+    @staticmethod
+    def _get_embedding(num_embeddings, embedding_dim, padding_idx):
+        half_dim = embedding_dim // 2
+        emb = math.log(10000) / (half_dim - 1)
+        emb = torch.exp(torch.arange(half_dim, dtype=torch.float) * -emb)
+        emb = torch.arange(num_embeddings, dtype=torch.float).unsqueeze(1) * emb.unsqueeze(0)
+        emb = torch.cat([torch.sin(emb), torch.cos(emb)], dim=1).view(num_embeddings, -1)
+        if embedding_dim % 2 == 1:
+            emb = torch.cat([emb, torch.zeros(num_embeddings, 1)], dim=1)
+        if padding_idx is not None:
+            emb[padding_idx, :] = 0
+        return emb
+
+    def forward(self, input_ids):
+        mask = input_ids.ne(self.padding_idx).int()
+        positions = (torch.cumsum(mask, dim=1) * mask).long() + self.padding_idx
+        return self.weights.index_select(0, positions.view(-1)).view(
+            input_ids.shape[0], input_ids.shape[1], -1
+        ).detach()
+
+
+class ErnieRNATwodProj(nn.Module):
+    def __init__(self, config):
+        super().__init__()
+        self.linear1 = nn.Linear(1, 6)
+        self.linear2 = nn.Linear(6, config.attention_heads)
+        self.activation_fn = ACT2FN[config.activation_fn]
+
+    def forward(self, x):
+        x = self.linear1(x)
+        x = self.activation_fn(x)
+        x = self.linear2(x)
+        return x
+
+
+def _compute_pairing_bias(input_ids):
+    B, T = input_ids.shape
+    xi = input_ids.unsqueeze(2).expand(B, T, T)
+    xj = input_ids.unsqueeze(1).expand(B, T, T)
+
+    score = torch.zeros(B, T, T, dtype=torch.float32, device=input_ids.device)
+    score[(xi == 5) & (xj == 6)] = 2.0
+    score[(xi == 6) & (xj == 5)] = 2.0
+    score[(xi == 4) & (xj == 7)] = 3.0
+    score[(xi == 7) & (xj == 4)] = 3.0
+    score[(xi == 4) & (xj == 6)] = 0.8
+    score[(xi == 6) & (xj == 4)] = 0.8
+    return score.unsqueeze(-1)  # [B, T, T, 1]
+
+
+class ErnieRNAAttention(nn.Module):
+    def __init__(self, config):
+        super().__init__()
+        self.embed_dim = config.embed_dim
+        self.num_heads = config.attention_heads
+        self.head_dim = self.embed_dim // self.num_heads
+        assert self.head_dim * self.num_heads == self.embed_dim
+
+        self.q_proj = nn.Linear(self.embed_dim, self.embed_dim)
+        self.k_proj = nn.Linear(self.embed_dim, self.embed_dim)
+        self.v_proj = nn.Linear(self.embed_dim, self.embed_dim)
+        self.out_proj = nn.Linear(self.embed_dim, self.embed_dim)
+        self.dropout = nn.Dropout(config.attention_dropout)
+
+    def _to_bh_t_hd(self, tensor, tgt_len, bsz):
+        return tensor.contiguous().view(tgt_len, bsz * self.num_heads, self.head_dim).transpose(0, 1)
+
+    def forward(self, x, key_padding_mask=None, twod_bias=None, output_attentions=False):
+        tgt_len, bsz, _ = x.size()
+
+        q = self._to_bh_t_hd(self.q_proj(x), tgt_len, bsz)
+        k = self._to_bh_t_hd(self.k_proj(x), tgt_len, bsz)
+        v = self._to_bh_t_hd(self.v_proj(x), tgt_len, bsz)
+
+        scale = self.head_dim ** -0.5
+        q = q * scale
+
+        attn_weights = torch.bmm(q, k.transpose(-2, -1))  # [B*H, T, T]
+
+        if key_padding_mask is not None:
+            attn_weights = attn_weights.view(bsz, self.num_heads, tgt_len, tgt_len)
+            attn_weights = attn_weights.masked_fill(
+                key_padding_mask.unsqueeze(1).unsqueeze(2), float("-inf")
+            )
+            attn_weights = attn_weights.view(bsz * self.num_heads, tgt_len, tgt_len)
+
+        if twod_bias is not None:
+            attn_weights = attn_weights + twod_bias.reshape(bsz * self.num_heads, tgt_len, tgt_len)
+
+        # Pre-softmax attention becomes the 2D bias for the next layer
+        twod_bias_new = attn_weights.view(bsz, self.num_heads, tgt_len, tgt_len)
+
+        attn_probs = F.softmax(attn_weights, dim=-1)
+        attn_weights_out = None
+        if output_attentions:
+            attn_weights_out = attn_probs.view(bsz, self.num_heads, tgt_len, tgt_len)
+        attn_probs = self.dropout(attn_probs)
+
+        out = torch.bmm(attn_probs, v)
+        out = out.transpose(0, 1).contiguous().view(tgt_len, bsz, self.embed_dim)
+        out = self.out_proj(out)
+
+        return out, attn_weights_out, twod_bias_new
+
+
+class ErnieRNALayer(nn.Module):
+    def __init__(self, config):
+        super().__init__()
+        self.self_attn = ErnieRNAAttention(config)
+        self.self_attn_layer_norm = nn.LayerNorm(config.embed_dim)
+        self.fc1 = nn.Linear(config.embed_dim, config.ffn_embed_dim)
+        self.fc2 = nn.Linear(config.ffn_embed_dim, config.embed_dim)
+        self.final_layer_norm = nn.LayerNorm(config.embed_dim)
+        self.dropout = nn.Dropout(config.dropout)
+        self.activation_dropout = nn.Dropout(config.activation_dropout)
+        self.activation_fn = ACT2FN[config.activation_fn]
+
+    def forward(self, x, key_padding_mask=None, twod_bias=None, output_attentions=False):
+        residual = x
+        x, attn_weights, twod_bias_new = self.self_attn(
+            x,
+            key_padding_mask=key_padding_mask,
+            twod_bias=twod_bias,
+            output_attentions=output_attentions,
+        )
+        x = self.dropout(x)
+        x = self.self_attn_layer_norm(residual + x)
+
+        residual = x
+        x = self.activation_fn(self.fc1(x))
+        x = self.activation_dropout(x)
+        x = self.fc2(x)
+        x = self.dropout(x)
+        x = self.final_layer_norm(residual + x)
+
+        return x, attn_weights, twod_bias_new
+
+
+class ErnieRNAModel(PreTrainedModel):
+    config_class = ErnieRNAConfig
+    base_model_prefix = "model"
+    _supports_sdpa = False
+    _supports_flash_attn_2 = False
+
+    def __init__(self, config):
+        super().__init__(config)
+        self.padding_idx = config.padding_idx
+
+        self.embed_tokens = nn.Embedding(config.vocab_size, config.embed_dim, padding_idx=config.padding_idx)
+        self.embed_positions = ErnieRNASinusoidalPositionalEmbedding(
+            config.max_positions, config.embed_dim, config.padding_idx
+        )
+        self.segment_embeddings = nn.Embedding(config.num_segments, config.embed_dim)
+        self.emb_layer_norm = nn.LayerNorm(config.embed_dim)
+        self.dropout = nn.Dropout(config.dropout)
+        self.layers = nn.ModuleList([ErnieRNALayer(config) for _ in range(config.num_layers)])
+        self.twod_proj = ErnieRNATwodProj(config)
+
+        self.post_init()
+
+    def forward(
+        self,
+        input_ids=None,
+        attention_mask=None,
+        token_type_ids=None,
+        output_attentions=None,
+        output_hidden_states=None,
+        return_dict=None,
+    ):
+        output_attentions = output_attentions if output_attentions is not None else self.config.output_attentions
+        output_hidden_states = output_hidden_states if output_hidden_states is not None else self.config.output_hidden_states
+        return_dict = return_dict if return_dict is not None else self.config.use_return_dict
+
+        # HF: 1=attend, 0=pad -> True=padding
+        if attention_mask is not None:
+            padding_mask = attention_mask.eq(0)
+        else:
+            padding_mask = input_ids.eq(self.padding_idx)
+
+        # Zero out padding positions after masking (matches fairseq behavior)
+        x = self.embed_tokens(input_ids)
+        # Sinusoidal PE is a float32 buffer; cast to activation dtype for bfloat16 compat.
+        x = x + self.embed_positions(input_ids).to(x.dtype)
+        if token_type_ids is not None:
+            x = x + self.segment_embeddings(token_type_ids)
+        x = self.emb_layer_norm(x)
+        if padding_mask.any():
+            x = x * (~padding_mask).unsqueeze(-1).to(x.dtype)
+        x = self.dropout(x)
+
+        # Compute initial 2D bias from sequence (always float32 as in original)
+        pairing = _compute_pairing_bias(input_ids)  # [B, T, T, 1]
+        twod_proj_f32 = self.twod_proj.float()
+        twod_bias = twod_proj_f32(pairing.float())  # [B, T, T, H]
+        twod_bias = twod_bias.permute(0, 3, 1, 2).contiguous().to(x.dtype)  # [B, H, T, T]
+
+        # Transpose to [T, B, C] for attention
+        x = x.transpose(0, 1)
+
+        all_hidden_states = []
+        all_attentions = []
+
+        if output_hidden_states:
+            all_hidden_states.append(x.transpose(0, 1))
+
+        key_padding_mask = padding_mask if padding_mask.any() else None
+
+        for layer in self.layers:
+            x, attn_weights, twod_bias = layer(
+                x,
+                key_padding_mask=key_padding_mask,
+                twod_bias=twod_bias,
+                output_attentions=output_attentions,
+            )
+            if output_hidden_states:
+                all_hidden_states.append(x.transpose(0, 1))
+            if output_attentions:
+                all_attentions.append(attn_weights)
+
+        x = x.transpose(0, 1)  # [B, T, C]
+
+        if not return_dict:
+            return tuple(v for v in [x, tuple(all_hidden_states) or None, tuple(all_attentions) or None] if v is not None)
+
+        return BaseModelOutput(
+            last_hidden_state=x,
+            hidden_states=tuple(all_hidden_states) if output_hidden_states else None,
+            attentions=tuple(all_attentions) if output_attentions else None,
+        )
+
+
+class ErnieRNALMHead(nn.Module):
+    def __init__(self, config):
+        super().__init__()
+        self.dense = nn.Linear(config.embed_dim, config.embed_dim)
+        self.layer_norm = nn.LayerNorm(config.embed_dim)
+        self.activation_fn = ACT2FN[config.activation_fn]
+        self.decoder = nn.Linear(config.embed_dim, config.vocab_size)
+
+    def forward(self, x):
+        x = self.layer_norm(self.activation_fn(self.dense(x)))
+        x = self.decoder(x)
+        return x
+
+
+class ErnieRNAForMaskedLM(PreTrainedModel):
+    config_class = ErnieRNAConfig
+    base_model_prefix = "model"
+    _supports_sdpa = False
+    _supports_flash_attn_2 = False
+
+    def __init__(self, config):
+        super().__init__(config)
+        self.model = ErnieRNAModel(config)
+        self.lm_head = ErnieRNALMHead(config)
+        self.post_init()
+
+    def forward(
+        self,
+        input_ids=None,
+        attention_mask=None,
+        token_type_ids=None,
+        labels=None,
+        output_attentions=None,
+        output_hidden_states=None,
+        return_dict=None,
+    ):
+        return_dict = return_dict if return_dict is not None else self.config.use_return_dict
+
+        out = self.model(
+            input_ids,
+            attention_mask=attention_mask,
+            token_type_ids=token_type_ids,
+            output_attentions=output_attentions,
+            output_hidden_states=output_hidden_states,
+            return_dict=return_dict,
+        )
+
+        logits = self.lm_head(out[0])
+
+        loss = None
+        if labels is not None:
+            loss = F.cross_entropy(
+                logits.view(-1, self.config.vocab_size),
+                labels.view(-1),
+                ignore_index=-100,
+            )
+
+        if not return_dict:
+            output = (logits,) + out[1:]
+            return ((loss,) + output) if loss is not None else output
+
+        return MaskedLMOutput(
+            loss=loss,
+            logits=logits,
+            hidden_states=out.hidden_states,
+            attentions=out.attentions,
+        )

pytorch_model.bin

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:10de02fd6d0034cd2e6987abffb3abafad75ac0c7420697d0572c45c4fac1bfb
+size 345970527

special_tokens_map.json

@@ -0,0 +1,7 @@
+{
+  "cls_token": "<cls>",
+  "eos_token": "<eos>",
+  "mask_token": "<mask>",
+  "pad_token": "<pad>",
+  "unk_token": "<unk>"
+}

tokenization_ernierna.py

@@ -0,0 +1,115 @@
+import json
+import os
+from transformers import PreTrainedTokenizer
+
+
+_VOCAB = {
+    "<cls>": 0,
+    "<pad>": 1,
+    "<eos>": 2,
+    "<unk>": 3,
+    "G": 4,
+    "A": 5,
+    "U": 6,
+    "C": 7,
+    "N": 8,
+    "Y": 9,
+    "R": 10,
+    "S": 11,
+    "K": 12,
+    "W": 13,
+    "M": 14,
+    "D": 15,
+    "H": 16,
+    "V": 17,
+    "B": 18,
+    "X": 19,
+    "I": 20,
+    "madeupword0000": 21,
+    "madeupword0001": 22,
+    "madeupword0002": 23,
+    "<mask>": 24,
+}
+
+
+class ErnieRNATokenizer(PreTrainedTokenizer):
+    vocab_files_names = {"vocab_file": "vocab.json"}
+    model_input_names = ["input_ids", "attention_mask"]
+
+    def __init__(
+        self,
+        vocab_file=None,
+        cls_token="<cls>",
+        pad_token="<pad>",
+        eos_token="<eos>",
+        unk_token="<unk>",
+        mask_token="<mask>",
+        **kwargs,
+    ):
+        if vocab_file is not None and os.path.isfile(vocab_file):
+            with open(vocab_file) as f:
+                self._vocab = json.load(f)
+        else:
+            self._vocab = dict(_VOCAB)
+        self._ids_to_tokens = {v: k for k, v in self._vocab.items()}
+
+        super().__init__(
+            cls_token=cls_token,
+            pad_token=pad_token,
+            eos_token=eos_token,
+            unk_token=unk_token,
+            mask_token=mask_token,
+            **kwargs,
+        )
+
+    @property
+    def vocab_size(self):
+        return len(self._vocab)
+
+    def get_vocab(self):
+        return dict(self._vocab)
+
+    def _tokenize(self, text):
+        tokens = []
+        for ch in text.upper():
+            if ch == "T":
+                tokens.append("U")
+            elif ch in self._vocab:
+                tokens.append(ch)
+            else:
+                tokens.append("<unk>")
+        return tokens
+
+    def _convert_token_to_id(self, token):
+        return self._vocab.get(token, self._vocab["<unk>"])
+
+    def _convert_id_to_token(self, index):
+        return self._ids_to_tokens.get(index, "<unk>")
+
+    def save_vocabulary(self, save_directory, filename_prefix=None):
+        os.makedirs(save_directory, exist_ok=True)
+        fname = (filename_prefix + "-" if filename_prefix else "") + "vocab.json"
+        path = os.path.join(save_directory, fname)
+        with open(path, "w") as f:
+            json.dump(self._vocab, f, indent=2)
+        return (path,)
+
+    def build_inputs_with_special_tokens(self, token_ids_0, token_ids_1=None):
+        cls = [self.cls_token_id]
+        eos = [self.eos_token_id]
+        if token_ids_1 is None:
+            return cls + token_ids_0 + eos
+        return cls + token_ids_0 + eos + cls + token_ids_1 + eos
+
+    def get_special_tokens_mask(self, token_ids_0, token_ids_1=None, already_has_special_tokens=False):
+        if already_has_special_tokens:
+            return super().get_special_tokens_mask(token_ids_0, token_ids_1, already_has_special_tokens=True)
+        mask = [1] + [0] * len(token_ids_0) + [1]
+        if token_ids_1 is not None:
+            mask += [1] + [0] * len(token_ids_1) + [1]
+        return mask
+
+    def create_token_type_ids_from_sequences(self, token_ids_0, token_ids_1=None):
+        if token_ids_1 is None:
+            return [0] + token_ids_0 + [0]
+        return [0] + token_ids_0 + [0, 0] + token_ids_1 + [0]

tokenizer_config.json

@@ -0,0 +1,56 @@
+{
+  "added_tokens_decoder": {
+    "0": {
+      "content": "<cls>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "1": {
+      "content": "<pad>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "2": {
+      "content": "<eos>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "3": {
+      "content": "<unk>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "24": {
+      "content": "<mask>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    }
+  },
+  "clean_up_tokenization_spaces": false,
+  "cls_token": "<cls>",
+  "eos_token": "<eos>",
+  "extra_special_tokens": {},
+  "mask_token": "<mask>",
+  "auto_map": {
+    "AutoTokenizer": ["tokenization_ernierna.ErnieRNATokenizer", null]
+  },
+  "model_max_length": 1024,
+  "pad_token": "<pad>",
+  "tokenizer_class": "ErnieRNATokenizer",
+  "unk_token": "<unk>"
+}

vocab.json

@@ -0,0 +1,27 @@
+{
+  "<cls>": 0,
+  "<pad>": 1,
+  "<eos>": 2,
+  "<unk>": 3,
+  "G": 4,
+  "A": 5,
+  "U": 6,
+  "C": 7,
+  "N": 8,
+  "Y": 9,
+  "R": 10,
+  "S": 11,
+  "K": 12,
+  "W": 13,
+  "M": 14,
+  "D": 15,
+  "H": 16,
+  "V": 17,
+  "B": 18,
+  "X": 19,
+  "I": 20,
+  "madeupword0000": 21,
+  "madeupword0001": 22,
+  "madeupword0002": 23,
+  "<mask>": 24
+}