Upload folder using huggingface_hub

LICENSE

@@ -0,0 +1,50 @@
+GENBIO AI COMMUNITY LICENSE AGREEMENT
+
+This GenBio AI Community License Agreement (the “License”) constitutes an agreement between you or the legal entity you represent (“you” or “your”) and GENBIO.AI, INC. (“GenBio”), governing your use of the GenBio Materials. If you are using the GenBio Materials on behalf of a legal entity, you represent and warrant to GenBio that you have full legal authority to act on behalf of that legal entity as applicable under the License. If you do not have the authority to accept this License or if you disagree with any or all of the License, you shall not use the GenBio Materials in any manner. By using or distributing any portion or element of the GenBio Materials, you imply your agreement to be bound by the License.
+
+“GenBio Materials” means any datasets, code, model weights or any other materials provided by GenBio at the following GitHub Page https://github.com/genbio-ai or Hugging Face Page https://huggingface.co/genbio-ai, including any updates or modifications made from time to time, whether in Source or Object form, and is made available to you under this License.
+
+
+1. License Grant.
+      1.1 License Scope. Subject to the terms of this License, GenBio grants you a non-exclusive, worldwide, non-transferable, non-sublicensable, revocable and royalty-free limited license under GenBio’s intellectual property or other rights owned by GenBio embodied in the GenBio Materials to use, reproduce, distribute, and create Derivative Works of, and make modifications to, the GenBio Materials for any Non-Commercial Purposes.
+      1.2 Use Restrictions. Restricted activities in relation to the License or use of GenBio Materials include:
+            1.2.1 You shall use the GenBio Materials, Contributions, Derivative Works, Outputs and Output Derivatives (as defined below) solely for Non-Commercial Purposes;
+            1.2.2 You shall not, directly or indirectly: (a) use or provide access to any Outputs or Output Derivatives to train, optimize, improve, or otherwise enhance the functionality or performance of any machine learning models or related technologies that are similar to the GenBio Materials; (b) engage in any form of model distillation or other methods that would achieve the purposes described in subsection (a) above. Notwithstanding the foregoing, you may use Outputs and Output Derivatives to train, optimize, improve, or enhance the functionality or performance of: (i) The GenBio Materials itself; and (ii) downstream Derivative Works of the GenBio Materials;
+            1.2.3 Your use of the GenBio Materials shall be subject to any additional terms and conditions that: (a) GenBio provides to you separately; or (b) GenBio otherwise makes available to you.
+
+2. Sharing and Distribution.  
+      2.1 Subject to Section 1, if you distribute or make available the GenBio Materials or a Derivative Work to a third party for your Non-Commercial Purposes, in Source or Object form, you shall: 
+            2.1.1 provide a copy of this License to that third party;
+            2.1.2 retain the following attribution notice within a “Notice” text file distributed as a part of such copies: “This is licensed under the GenBio AI Community License Agreement, Copyright ©  GENBIO.AI, INC. All Rights Reserved”; and
+            2.1.3 prominently display “Powered by GenBio AI” on a related website, user interface, blogpost, about page, or product documentation. 
+      2.2 If You create a Derivative Work, you may add your own attribution notice(s) to the “Notice” text file included with that Derivative Work, provided that you clearly indicate which attributions apply to the GenBio Materials and state in the “Notice” text file that you changed the GenBio Materials and how it was modified.
+
+3. Submission of Contribution.
+      Unless you explicitly state otherwise, any Contribution intentionally submitted for inclusion in the GenBio Materials by you to GenBio shall be under the terms and conditions of this License, without any additional terms or conditions. Notwithstanding the above, nothing herein shall supersede or modify the terms of any separate license agreement you may have executed with GenBio regarding such Contributions.
+
+4. Export Control.
+      You shall comply with the applicable U.S. Foreign Corrupt Practices Act and all applicable export laws, restrictions and regulations of the U.S. Department of Commerce, and any other applicable U.S. and foreign authority.
+
+5. Disclaimer of Warranty.
+      GENBIO MATERIALS PROVIDED BY GENBIO OR ANY OUTPUT YOU RECEIVED ARE PROVIDED “AS IS.” EXCEPT TO THE EXTENT PROHIBITED BY LAW. GENBIO MAKES NO REPRESENTATIONS OR WARRANTIES OF ANY KIND, WHETHER EXPRESS, IMPLIED OR OTHERWISE, REGARDING THE ACCURACY, COMPLETENESS OR PERFORMANCE OF THE SERVICES AND YOUR OUTPUT, OR WITH RESPECT TO SATISFACTORY QUALITY, FITNESS FOR A PARTICULAR PURPOSE OR NON-INFRINGEMENT. 
+
+6. Limitation of Liability.  
+      In no event and under no legal theory, whether in tort (including negligence), contract, or otherwise, unless required by applicable law (such as deliberate and grossly negligent acts) or agreed to in writing, shall any Contributor be liable to You for damages, including any direct, indirect, special, incidental, or consequential damages of any character arising as a result of this License or out of the use or inability to use the GenBio Materials (including but not limited to damages for loss of goodwill, work stoppage, computer failure or malfunction, or any and all other commercial damages or losses), even if such Contributor has been advised of the possibility of such damages.
+
+7. General Terms. 
+      7.1 Relationship of Parties.  You and GenBio are independent contractors, and nothing herein shall be deemed to constitute either party as the agent or representative of the other or both parties as joint venturers or partners for any purpose.
+      7.2 Assignment.  This License and the rights and obligations herein may not be assigned or transferred, in whole or in part, by You without the prior written consent of GenBio. Any assignment in violation of this provision is void. GenBio may freely assign or transfer this License, in whole or in part. This License shall be binding upon, and inure to the benefit of, the successors and permitted assigns of the parties.
+      7.3 Governing Law.  This License shall be governed, construed and interpreted in accordance with the laws of the State of California, without giving effect to principles of conflicts of law. Each of the parties to this License consents to the exclusive jurisdiction and venue of the courts of the state and federal courts of California.
+      7.4 Severability. If any provision of this License is held to be invalid, illegal or unenforceable in any respect, that provision shall be limited or eliminated to the minimum extent necessary so that this License otherwise remains in full force and effect and enforceable.
+
+8. Definitions.
+      8.1 “Commercial Entity” means any entity engaged in any activity intended for or directed toward commercial advantage or monetary compensation, including, without limitation, the development of any product or service intended to be sold or made available for a fee. For the purpose of this License, references to a Commercial Entity expressly exclude any universities, non-profit organizations, not-for-profit entities, research institutes and educational and government bodies.
+      8.2 “Contribution” means any work of authorship, including the original version of the GenBio Materials and any modifications or additions to that GenBio Materials or Derivative Works thereof, that is intentionally submitted to GenBio for inclusion in the GenBio Materials by the copyright owner or by an individual or legal entity authorized to submit on behalf of the copyright owner. For the purposes of this definition, “submitted” means any form of electronic, verbal, or written communication sent to GenBio or its representatives, including but not limited to communication on electronic mailing lists, source code control systems, and issue tracking systems that are managed by, or on behalf of, GenBio for the purpose of discussing and improving the GenBio Materials, but excluding Outputs and all communications that are conspicuously marked or otherwise designated in writing by the copyright owner as “Not a Contribution”.
+      8.3 “Contributor” means GenBio and any individual or legal entity on behalf of whom a Contribution has been received by GenBio and subsequently incorporated within the GenBio Materials.
+      8.4 “Derivative Work” means any work, whether in Source or Object form, that is based on (or derived from) the GenBio Materials and for which the editorial revisions, annotations, elaborations, or other modifications represent, as a whole, an original work of authorship. For the purposes of this License, Derivative Works shall not include works that remain separable from, or merely link (or bind by name) to the interfaces of, the GenBio Materials and Derivative Works thereof.
+      8.5 “Non-Commercial Purposes” means uses not intended for or directed toward commercial advantage or monetary compensation, or the facilitation of development of any product or service to be sold or made available for a fee. For the avoidance of doubt, the provision of Outputs as a service is not a Non-Commercial Purpose.
+      8.6 “Object” means any form resulting from mechanical transformation or translation of a Source form, including but not limited to compiled object code, generated documentation, and conversions to other media types.
+      8.7 “Output” means any output, including any protein sequence, structure prediction, functional annotation, molecule, descriptions of a molecule, model, sequence, text, and/or image that is elicited directly or indirectly by, or otherwise made available to, you in connection with your use of the GenBio Materials, including, but not limited to, the use of AI-Powered Technology. For the avoidance of doubt, it includes any intermediate results, such as activations across model layers, intermediate outputs from model layers (e.g., attention maps), as well as gradients and embeddings produced by the GenBio Materials.
+      8.8 “Output Derivatives” means any enhancements, modifications and derivative works of Outputs (including, but not limited to, any derivative sequences or molecules).
+      8.9 “Source” means the preferred form for making modifications, including but not limited to GenBio Materials source code, documentation source, and configuration files.
+

README.md

@@ -0,0 +1,142 @@
+---
+language:
+- rna
+library_name: transformers
+tags:
+- RNA
+- language-model
+license: other
+---
+
+# AIDO.RNA-650M-CDS
+
+650M-parameter model fine-tuned on coding sequences (CDS). This is a standalone HuggingFace port that loads without the
+[ModelGenerator](https://github.com/genbio-ai/ModelGenerator) package.
+
+## Architecture
+
+| Parameter | Value |
+|---|---|
+| Layers | 33 |
+| Attention heads | 20 |
+| Embedding dimension | 1280 |
+| Intermediate (MLP) size | 3392 |
+| Vocabulary size | 16 |
+| Positional encoding | RoPE (rotary_percent=1.0) |
+| Normalization | LayerNorm |
+| MLP activation | SwiGLU |
+| Architecture | Pre-LN Transformer |
+| Max sequence length | 1024 (training truncation; RoPE has no hard limit) |
+
+**Vocabulary:** `[PAD]`, `[MASK]`, `[CLS]`, `[SEP]`, `[UNK]`, `A`, `G`, `C`, `T`, `U`, `N`,
+`[BOS]`, `[EOS]`, `[UNUSED1]`, `[UNUSED2]`, `[UNUSED3]`
+
+## Pretraining
+
+- **Objective:** Masked language modeling (MLM) on RNA sequences
+- **Data:** 42M ncRNA sequences (pre-training) + CDS fine-tuning
+- **Source checkpoint:** `genbio-ai/AIDO.RNA-650M-CDS`
+
+### Checkpoint selection
+
+Preferred over the base 650M for tasks involving coding regions or translation.
+
+## Parity Verification
+
+Hidden-state representations compared against the original `genbio-ai/AIDO.RNA-650M-CDS` weights at all
+34 representation levels (embedding + 33 transformer layers).
+Intermediate layer differences are floating-point accumulation noise normalised away by the
+final layer norm; the final output matches the original within 1e-5.
+max diff = 4.77e-06 (final output), 3.81e-05 (intermediate layers). Verified on GPU with PyTorch 2.7 / CUDA 12.
+
+## Related Models
+
+See the full [AIDO.RNA collection](COLLECTION_URL).
+
+| Model | Parameters | Data | Notes |
+|---|---|---|---|
+| [AIDO.RNA-1M-MARS](COLLECTION_URL) | 1M | MARS ncRNA | Smallest MARS variant |
+| [AIDO.RNA-25M-MARS](COLLECTION_URL) | 25M | MARS ncRNA | Mid-size MARS variant |
+| [AIDO.RNA-300M-MARS](COLLECTION_URL) | 300M | MARS ncRNA | Large MARS variant |
+| [AIDO.RNA-650M](COLLECTION_URL) | 650M | 42M ncRNA | Base model |
+| [AIDO.RNA-650M-CDS](COLLECTION_URL) | 650M | 42M ncRNA + CDS | CDS-adapted |
+| [AIDO.RNA-1.6B](COLLECTION_URL) | 1.6B | 42M ncRNA | Largest base model |
+| [AIDO.RNA-1.6B-CDS](COLLECTION_URL) | 1.6B | 42M ncRNA + CDS | Largest CDS-adapted |
+
+## Usage
+
+### Embedding generation
+
+```python
+import torch
+from transformers import AutoTokenizer, AutoModel
+
+tokenizer = AutoTokenizer.from_pretrained("Taykhoom/AIDO.RNA-650M-CDS", trust_remote_code=True)
+model = AutoModel.from_pretrained("Taykhoom/AIDO.RNA-650M-CDS", trust_remote_code=True)
+model.eval()
+
+sequences = ["ACGUGCUAGCUAGCUA", "AUGCUAGCUAGCUAGC"]
+enc = tokenizer(sequences, return_tensors="pt", padding=True)
+
+with torch.no_grad():
+    out = model(**enc)
+
+cls_emb   = out.last_hidden_state[:, 0, :]   # (batch, 1280) -- CLS token
+token_emb = out.last_hidden_state             # (batch, seq_len, 1280)
+
+# Intermediate layers
+out_all = model(**enc, output_hidden_states=True)
+layer3_emb = out_all.hidden_states[3]
+```
+
+### MLM logits
+
+```python
+from transformers import AutoTokenizer, AutoModelForMaskedLM
+
+tokenizer = AutoTokenizer.from_pretrained("Taykhoom/AIDO.RNA-650M-CDS", trust_remote_code=True)
+model = AutoModelForMaskedLM.from_pretrained("Taykhoom/AIDO.RNA-650M-CDS", trust_remote_code=True)
+model.eval()
+
+enc = tokenizer(["ACGU[MASK]GCUA"], return_tensors="pt")
+with torch.no_grad():
+    logits = model(**enc).logits   # (1, seq_len, 16)
+```
+
+### Fine-tuning
+
+Standard HF conventions. Use `cls_emb = out.last_hidden_state[:, 0, :]` (CLS token) as
+input to a task-specific head for sequence-level tasks.
+
+## Implementation Notes
+
+The original `genbio-ai/AIDO.RNA-650M-CDS` checkpoint requires the
+[ModelGenerator](https://github.com/genbio-ai/ModelGenerator) package to load.
+This port is a clean standalone re-implementation:
+
+- All model logic is contained in `modeling_aidorna.py` and `configuration_aidorna.py`.
+- `attn_implementation="sdpa"` and `attn_implementation="flash_attention_2"` are added
+  (not present in the original genbio-ai implementation).
+- Architecture: pre-LN Transformer with SwiGLU MLP and RoPE positional embeddings.
+
+## Citation
+
+```bibtex
+@article{zou2024_aidorna,
+  title   = {A Large-Scale Foundation Model for {RNA} Function and Structure Prediction},
+  author  = {Zou, Shuxian and Tao, Tianhua and Mahbub, Sazan and Ellington, Caleb N. and Algayres, Robin and Li, Dian and Zhuang, Yonghao and Wang, Hongyi and Song, Le and Xing, Eric P.},
+  journal = {bioRxiv},
+  year    = {2024},
+  doi     = {10.1101/2024.11.28.625345}
+}
+```
+
+## Credits
+
+Original model and code by Zou et al. Source: [GitHub](https://github.com/genbio-ai/AIDO).
+The HF conversion code was authored primarily by [Claude Code](https://claude.ai/code)
+and reviewed manually by Taykhoom Dalal.
+
+## License
+
+GenBio AI Community License, following the original repository. See `LICENSE` for details.

config.json

@@ -0,0 +1,26 @@
+{
+  "add_linear_bias": true,
+  "attention_probs_dropout_prob": 0.0,
+  "auto_map": {
+    "AutoConfig": "configuration_aidorna.AIDORNAConfig",
+    "AutoModel": "modeling_aidorna.AIDORNAModel",
+    "AutoModelForMaskedLM": "modeling_aidorna.AIDORNAForMaskedLM"
+  },
+  "hidden_act": "swiglu",
+  "hidden_dropout_prob": 0.0,
+  "hidden_size": 1280,
+  "initializer_range": 0.02,
+  "intermediate_size": 3392,
+  "layer_norm_eps": 1e-05,
+  "max_position_embeddings": 1024,
+  "model_type": "aidorna",
+  "normalization_type": "LayerNorm",
+  "num_attention_heads": 20,
+  "num_hidden_layers": 33,
+  "pad_token_id": 0,
+  "position_embedding_type": "rope",
+  "rotary_percent": 1.0,
+  "seq_len_interpolation_factor": null,
+  "transformers_version": "4.57.6",
+  "vocab_size": 16
+}

configuration_aidorna.py

@@ -0,0 +1,55 @@
+from transformers import PretrainedConfig
+
+
+class AIDORNAConfig(PretrainedConfig):
+    model_type = "aidorna"
+
+    auto_map = {
+        "AutoConfig": "configuration_aidorna.AIDORNAConfig",
+        "AutoModel": "modeling_aidorna.AIDORNAModel",
+        "AutoModelForMaskedLM": "modeling_aidorna.AIDORNAForMaskedLM",
+    }
+
+    def __init__(
+        self,
+        vocab_size: int = 16,
+        hidden_size: int = 1280,
+        num_hidden_layers: int = 33,
+        num_attention_heads: int = 20,
+        intermediate_size: int = 3392,
+        hidden_act: str = "swiglu",
+        hidden_dropout_prob: float = 0.0,
+        attention_probs_dropout_prob: float = 0.0,
+        max_position_embeddings: int = 1024,
+        initializer_range: float = 0.02,
+        layer_norm_eps: float = 1e-5,
+        pad_token_id: int = 0,
+        position_embedding_type: str = "rope",
+        rotary_percent: float = 1.0,
+        seq_len_interpolation_factor=None,
+        normalization_type: str = "LayerNorm",
+        add_linear_bias: bool = True,
+        **kwargs,
+    ):
+        super().__init__(pad_token_id=pad_token_id, **kwargs)
+        self.auto_map = {
+            "AutoConfig": "configuration_aidorna.AIDORNAConfig",
+            "AutoModel": "modeling_aidorna.AIDORNAModel",
+            "AutoModelForMaskedLM": "modeling_aidorna.AIDORNAForMaskedLM",
+        }
+        self.vocab_size = vocab_size
+        self.hidden_size = hidden_size
+        self.num_hidden_layers = num_hidden_layers
+        self.num_attention_heads = num_attention_heads
+        self.intermediate_size = intermediate_size
+        self.hidden_act = hidden_act
+        self.hidden_dropout_prob = hidden_dropout_prob
+        self.attention_probs_dropout_prob = attention_probs_dropout_prob
+        self.max_position_embeddings = max_position_embeddings
+        self.initializer_range = initializer_range
+        self.layer_norm_eps = layer_norm_eps
+        self.position_embedding_type = position_embedding_type
+        self.rotary_percent = rotary_percent
+        self.seq_len_interpolation_factor = seq_len_interpolation_factor
+        self.normalization_type = normalization_type
+        self.add_linear_bias = add_linear_bias

model.safetensors

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

modeling_aidorna.py

@@ -0,0 +1,470 @@
+import math
+from typing import Optional, Union
+
+import torch
+import torch.nn.functional as F
+import torch.utils.checkpoint
+from torch import nn
+
+from transformers import PreTrainedModel
+from transformers.activations import ACT2FN
+from transformers.modeling_outputs import BaseModelOutput, MaskedLMOutput
+
+from .configuration_aidorna import AIDORNAConfig
+
+
+class AIDORNARMSNorm(nn.Module):
+    def __init__(self, hidden_size, eps=1e-6):
+        super().__init__()
+        self.weight = nn.Parameter(torch.ones(hidden_size))
+        self.variance_epsilon = eps
+
+    def forward(self, hidden_states):
+        input_dtype = hidden_states.dtype
+        hidden_states = hidden_states.to(torch.float32)
+        variance = hidden_states.pow(2).mean(-1, keepdim=True)
+        hidden_states = hidden_states * torch.rsqrt(variance + self.variance_epsilon)
+        return self.weight * hidden_states.to(input_dtype)
+
+
+def _get_norm_cls(config):
+    if config.normalization_type == "RMSNorm":
+        return AIDORNARMSNorm
+    return nn.LayerNorm
+
+
+class RotaryEmbedding(nn.Module):
+    def __init__(self, kv_channels, rotary_percent=1.0, seq_len_interpolation_factor=None, rotary_base=10000):
+        super().__init__()
+        dim = kv_channels
+        if rotary_percent < 1.0:
+            dim = int(dim * rotary_percent)
+        self.seq_len_interpolation_factor = seq_len_interpolation_factor
+        inv_freq = 1.0 / (rotary_base ** (torch.arange(0, dim, 2, dtype=torch.float32) / dim))
+        self.register_buffer("inv_freq", inv_freq, persistent=False)
+
+    def forward(self, max_seq_len, offset=0):
+        seq = torch.arange(max_seq_len, device=self.inv_freq.device, dtype=self.inv_freq.dtype) + offset
+        if self.seq_len_interpolation_factor is not None:
+            seq = seq / self.seq_len_interpolation_factor
+        freqs = torch.outer(seq, self.inv_freq)
+        emb = torch.cat((freqs, freqs), dim=-1)
+        return emb[:, None, None, :]  # (T, 1, 1, rot_dim*2)
+
+
+def _rotate_half(x):
+    x1, x2 = torch.chunk(x, 2, dim=-1)
+    return torch.cat((-x2, x1), dim=-1)
+
+
+def apply_rotary_pos_emb(t, freqs):
+    rot_dim = freqs.shape[-1]
+    t_rot, t_pass = t[..., :rot_dim], t[..., rot_dim:]
+    cos_ = torch.cos(freqs).to(t.dtype)
+    sin_ = torch.sin(freqs).to(t.dtype)
+    t_rot = (t_rot * cos_) + (_rotate_half(t_rot) * sin_)
+    return torch.cat((t_rot, t_pass), dim=-1)
+
+
+def _qkv_rope(q, k, rotary_pos_emb):
+    q = q.permute(2, 0, 1, 3).contiguous()  # (B,H,T,d) -> (T,B,H,d)
+    k = k.permute(2, 0, 1, 3).contiguous()
+    q_pos, k_pos = (rotary_pos_emb, rotary_pos_emb) if not isinstance(rotary_pos_emb, tuple) else rotary_pos_emb
+    q = apply_rotary_pos_emb(q, q_pos)
+    k = apply_rotary_pos_emb(k, k_pos)
+    q = q.permute(1, 2, 0, 3).contiguous()  # (T,B,H,d) -> (B,H,T,d)
+    k = k.permute(1, 2, 0, 3).contiguous()
+    return q, k
+
+
+class AIDORNASelfAttention(nn.Module):
+    def __init__(self, config):
+        super().__init__()
+        assert config.hidden_size % config.num_attention_heads == 0
+        self.num_heads = config.num_attention_heads
+        self.head_dim = config.hidden_size // config.num_attention_heads
+        self.all_head_size = self.num_heads * self.head_dim
+
+        self.query = nn.Linear(config.hidden_size, self.all_head_size, bias=config.add_linear_bias)
+        self.key = nn.Linear(config.hidden_size, self.all_head_size, bias=config.add_linear_bias)
+        self.value = nn.Linear(config.hidden_size, self.all_head_size, bias=config.add_linear_bias)
+        self.dropout = nn.Dropout(config.attention_probs_dropout_prob)
+
+    def _project(self, x):
+        B, T, _ = x.size()
+        def t(z):
+            return z.view(B, T, self.num_heads, self.head_dim).permute(0, 2, 1, 3)
+        return t(self.query(x)), t(self.key(x)), t(self.value(x))
+
+    def forward(self, hidden_states, attention_mask, output_attentions=False, rotary_pos_emb=None):
+        B, T, _ = hidden_states.size()
+        q, k, v = self._project(hidden_states)
+
+        if rotary_pos_emb is not None:
+            q, k = _qkv_rope(q, k, rotary_pos_emb)
+
+        scores = torch.matmul(q, k.transpose(-1, -2)) / math.sqrt(self.head_dim)
+        if attention_mask is not None:
+            scores = scores + attention_mask.to(scores.dtype)
+
+        probs = F.softmax(scores, dim=-1)
+        if attention_mask is not None:
+            probs = probs.masked_fill(attention_mask < -1e5, 0.0)
+        probs = self.dropout(probs)
+
+        context = torch.matmul(probs, v).permute(0, 2, 1, 3).contiguous().view(B, T, self.all_head_size)
+        return context, probs if output_attentions else None
+
+
+class AIDORNASdpaAttention(AIDORNASelfAttention):
+    def forward(self, hidden_states, attention_mask, output_attentions=False, rotary_pos_emb=None):
+        if output_attentions:
+            return super().forward(hidden_states, attention_mask, output_attentions=True, rotary_pos_emb=rotary_pos_emb)
+
+        B, T, _ = hidden_states.size()
+        q, k, v = self._project(hidden_states)
+
+        if rotary_pos_emb is not None:
+            q, k = _qkv_rope(q, k, rotary_pos_emb)
+
+        attn_mask = attention_mask.to(q.dtype) if attention_mask is not None else None
+        context = F.scaled_dot_product_attention(q, k, v, attn_mask=attn_mask)
+        context = context.permute(0, 2, 1, 3).contiguous().view(B, T, self.all_head_size)
+        return context, None
+
+
+class AIDORNAFlashAttention2(AIDORNASelfAttention):
+    def forward(self, hidden_states, attention_mask, output_attentions=False, rotary_pos_emb=None):
+        if output_attentions:
+            return super().forward(hidden_states, attention_mask, output_attentions=True, rotary_pos_emb=rotary_pos_emb)
+
+        try:
+            from flash_attn import flash_attn_func
+            from flash_attn.bert_padding import pad_input, unpad_input
+        except ImportError as e:
+            raise ImportError(
+                "flash_attn is required for attn_implementation='flash_attention_2'. "
+                "Install with: pip install flash-attn --no-build-isolation"
+            ) from e
+
+        B, T, _ = hidden_states.size()
+        q, k, v = self._project(hidden_states)  # each (B, H, T, d)
+
+        if rotary_pos_emb is not None:
+            q, k = _qkv_rope(q, k, rotary_pos_emb)
+
+        # flash_attn layout: (B, T, H, d)
+        q, k, v = q.permute(0, 2, 1, 3), k.permute(0, 2, 1, 3), v.permute(0, 2, 1, 3)
+
+        orig_dtype = q.dtype
+        if q.dtype not in (torch.float16, torch.bfloat16):
+            q = q.to(torch.bfloat16)
+            k = k.to(torch.bfloat16)
+            v = v.to(torch.bfloat16)
+
+        # Derive key padding mask from extended additive attention mask.
+        # Position 0 (CLS) is always non-padding, so row 0 reliably reflects per-token pad status.
+        if attention_mask is not None:
+            key_padding_mask_bool = (attention_mask[:, 0, 0, :] > -1e10)  # True = attend
+            has_padding = not key_padding_mask_bool.all()
+        else:
+            key_padding_mask_bool = None
+            has_padding = False
+
+        if has_padding and key_padding_mask_bool is not None:
+            from flash_attn import flash_attn_varlen_func
+            q_u, indices, cu_seqlens, max_seqlen, _ = unpad_input(q, key_padding_mask_bool)
+            k_u, *_ = unpad_input(k, key_padding_mask_bool)
+            v_u, *_ = unpad_input(v, key_padding_mask_bool)
+            out_u = flash_attn_varlen_func(
+                q_u, k_u, v_u,
+                cu_seqlens_q=cu_seqlens, cu_seqlens_k=cu_seqlens,
+                max_seqlen_q=max_seqlen, max_seqlen_k=max_seqlen,
+                causal=False,
+            )
+            out = pad_input(out_u, indices, B, T)  # (B, T, H, d)
+        else:
+            out = flash_attn_func(q, k, v, causal=False)  # (B, T, H, d)
+
+        context = out.to(orig_dtype).contiguous().view(B, T, self.all_head_size)
+        return context, None
+
+
+AIDORNA_ATTENTION_CLASSES = {
+    "eager": AIDORNASelfAttention,
+    "sdpa": AIDORNASdpaAttention,
+    "flash_attention_2": AIDORNAFlashAttention2,
+}
+
+
+class AIDORNASelfOutput(nn.Module):
+    def __init__(self, config):
+        super().__init__()
+        self.dense = nn.Linear(config.hidden_size, config.hidden_size, bias=config.add_linear_bias)
+        self.dropout = nn.Dropout(config.hidden_dropout_prob)
+
+    def forward(self, hidden_states, residual):
+        return residual + self.dropout(self.dense(hidden_states))
+
+
+class AIDORNAAttention(nn.Module):
+    def __init__(self, config):
+        super().__init__()
+        norm_cls = _get_norm_cls(config)
+        self.ln = norm_cls(config.hidden_size, eps=config.layer_norm_eps)
+        attn_cls = AIDORNA_ATTENTION_CLASSES[getattr(config, "_attn_implementation", "eager")]
+        self.self = attn_cls(config)
+        self.output = AIDORNASelfOutput(config)
+
+    def forward(self, hidden_states, attention_mask, output_attentions=False, rotary_pos_emb=None):
+        ln_out = self.ln(hidden_states)
+        attn_out, attn_weights = self.self(ln_out, attention_mask, output_attentions, rotary_pos_emb)
+        out = self.output(attn_out, hidden_states)
+        return out, attn_weights
+
+
+class AIDORNAMLP(nn.Module):
+    def __init__(self, config):
+        super().__init__()
+        self.up_proj = nn.Linear(config.hidden_size, config.intermediate_size, bias=config.add_linear_bias)
+        self.gate_proj = nn.Linear(config.hidden_size, config.intermediate_size, bias=config.add_linear_bias)
+        self.down_proj = nn.Linear(config.intermediate_size, config.hidden_size, bias=config.add_linear_bias)
+
+    def forward(self, x):
+        return self.down_proj(F.silu(self.gate_proj(x)) * self.up_proj(x))
+
+
+class AIDORNALayer(nn.Module):
+    def __init__(self, config):
+        super().__init__()
+        self.attention = AIDORNAAttention(config)
+        norm_cls = _get_norm_cls(config)
+        self.ln = norm_cls(config.hidden_size, eps=config.layer_norm_eps)
+        self.mlp = AIDORNAMLP(config)
+        self.dropout = nn.Dropout(config.hidden_dropout_prob)
+
+    def forward(self, hidden_states, attention_mask, output_attentions=False, rotary_pos_emb=None):
+        attn_out, attn_weights = self.attention(hidden_states, attention_mask, output_attentions, rotary_pos_emb)
+        mlp_out = self.dropout(self.mlp(self.ln(attn_out)))
+        return attn_out + mlp_out, attn_weights
+
+
+class AIDORNAEncoder(nn.Module):
+    def __init__(self, config):
+        super().__init__()
+        self.layer = nn.ModuleList([AIDORNALayer(config) for _ in range(config.num_hidden_layers)])
+        norm_cls = _get_norm_cls(config)
+        self.ln = norm_cls(config.hidden_size, eps=config.layer_norm_eps)
+        self.gradient_checkpointing = False
+
+    def forward(self, hidden_states, attention_mask, output_attentions, output_hidden_states, rotary_pos_emb=None):
+        all_hidden_states = (hidden_states,) if output_hidden_states else None
+        all_attentions = () if output_attentions else None
+
+        for layer in self.layer:
+            if self.gradient_checkpointing and self.training:
+                hidden_states, attn_w = self._gradient_checkpointing_func(
+                    layer.__call__, hidden_states, attention_mask, output_attentions, rotary_pos_emb
+                )
+            else:
+                hidden_states, attn_w = layer(hidden_states, attention_mask, output_attentions, rotary_pos_emb)
+
+            if output_hidden_states:
+                all_hidden_states = all_hidden_states + (hidden_states,)
+            if output_attentions:
+                all_attentions = all_attentions + (attn_w,)
+
+        hidden_states = self.ln(hidden_states)
+
+        if output_hidden_states:
+            # Replace the last element with the final LN-normalised version
+            all_hidden_states = all_hidden_states[:-1] + (hidden_states,)
+
+        return hidden_states, all_hidden_states, all_attentions
+
+
+class AIDORNAEmbeddings(nn.Module):
+    def __init__(self, config):
+        super().__init__()
+        self.word_embeddings = nn.Embedding(config.vocab_size, config.hidden_size, padding_idx=config.pad_token_id)
+        self.dropout = nn.Dropout(config.hidden_dropout_prob)
+        self.position_embedding_type = config.position_embedding_type
+
+        if config.position_embedding_type == "absolute":
+            self.position_embeddings = nn.Embedding(config.max_position_embeddings, config.hidden_size)
+        self.register_buffer(
+            "position_ids",
+            torch.arange(config.max_position_embeddings).expand((1, -1)),
+            persistent=False,
+        )
+
+    def forward(self, input_ids):
+        embeddings = self.word_embeddings(input_ids)
+        if self.position_embedding_type == "absolute":
+            seq_len = input_ids.size(1)
+            embeddings = embeddings + self.position_embeddings(self.position_ids[:, :seq_len])
+        return self.dropout(embeddings)
+
+
+def _make_extended_attention_mask(attention_mask):
+    b1s = attention_mask.unsqueeze(1)   # (B, 1, T)
+    bs1 = attention_mask.unsqueeze(2)   # (B, T, 1)
+    bss = (b1s * bs1).unsqueeze(1)      # (B, 1, T, T)
+    return (bss < 0.5).float() * torch.finfo(torch.float).min
+
+
+class AIDORNAPreTrainedModel(PreTrainedModel):
+    config_class = AIDORNAConfig
+    base_model_prefix = "model"
+    supports_gradient_checkpointing = True
+    _supports_sdpa = True
+    _supports_flash_attn_2 = True
+
+    def _init_weights(self, module):
+        if isinstance(module, (nn.Linear, nn.Embedding)):
+            module.weight.data.normal_(mean=0.0, std=self.config.initializer_range)
+        if isinstance(module, nn.LayerNorm):
+            module.bias.data.zero_()
+            module.weight.data.fill_(1.0)
+        if isinstance(module, AIDORNARMSNorm):
+            module.weight.data.fill_(1.0)
+        if isinstance(module, nn.Linear) and module.bias is not None:
+            module.bias.data.zero_()
+
+
+class AIDORNAModel(AIDORNAPreTrainedModel):
+    def __init__(self, config):
+        super().__init__(config)
+        self.embeddings = AIDORNAEmbeddings(config)
+        self.encoder = AIDORNAEncoder(config)
+
+        if config.position_embedding_type == "rope":
+            head_dim = config.hidden_size // config.num_attention_heads
+            self.rotary_pos_emb = RotaryEmbedding(
+                kv_channels=head_dim,
+                rotary_percent=config.rotary_percent,
+                seq_len_interpolation_factor=config.seq_len_interpolation_factor,
+            )
+        self.post_init()
+
+    def get_input_embeddings(self):
+        return self.embeddings.word_embeddings
+
+    def set_input_embeddings(self, value):
+        self.embeddings.word_embeddings = value
+
+    def forward(
+        self,
+        input_ids=None,
+        attention_mask=None,
+        output_hidden_states=None,
+        output_attentions=None,
+        return_dict=None,
+    ):
+        output_hidden_states = output_hidden_states if output_hidden_states is not None else self.config.output_hidden_states
+        output_attentions = output_attentions if output_attentions is not None else self.config.output_attentions
+        return_dict = return_dict if return_dict is not None else self.config.use_return_dict
+
+        if input_ids is None:
+            raise ValueError("input_ids must be provided")
+
+        if attention_mask is None:
+            attention_mask = torch.ones_like(input_ids)
+
+        extended_mask = _make_extended_attention_mask(attention_mask)
+
+        rotary_pos_emb = None
+        if self.config.position_embedding_type == "rope":
+            rotary_pos_emb = self.rotary_pos_emb(input_ids.size(1))
+
+        hidden_states = self.embeddings(input_ids)
+        last_hidden_state, all_hidden_states, all_attentions = self.encoder(
+            hidden_states, extended_mask, output_attentions, output_hidden_states, rotary_pos_emb
+        )
+
+        if not return_dict:
+            return tuple(v for v in [last_hidden_state, all_hidden_states, all_attentions] if v is not None)
+
+        return BaseModelOutput(
+            last_hidden_state=last_hidden_state,
+            hidden_states=all_hidden_states,
+            attentions=all_attentions,
+        )
+
+
+class AIDORNAPredictionHeadTransform(nn.Module):
+    def __init__(self, config):
+        super().__init__()
+        self.dense = nn.Linear(config.hidden_size, config.hidden_size)
+        self.transform_act_fn = ACT2FN["gelu"]
+        norm_cls = _get_norm_cls(config)
+        self.LayerNorm = norm_cls(config.hidden_size, eps=config.layer_norm_eps)
+
+    def forward(self, hidden_states):
+        return self.LayerNorm(self.transform_act_fn(self.dense(hidden_states)))
+
+
+class AIDORNALMPredictionHead(nn.Module):
+    def __init__(self, config):
+        super().__init__()
+        self.transform = AIDORNAPredictionHeadTransform(config)
+        self.decoder = nn.Linear(config.hidden_size, config.vocab_size, bias=False)
+        self.bias = nn.Parameter(torch.zeros(config.vocab_size))
+        self.decoder.bias = self.bias
+
+    def forward(self, hidden_states):
+        return self.decoder(self.transform(hidden_states))
+
+
+class AIDORNAOnlyMLMHead(nn.Module):
+    def __init__(self, config):
+        super().__init__()
+        self.predictions = AIDORNALMPredictionHead(config)
+
+    def forward(self, x):
+        return self.predictions(x)
+
+
+class AIDORNAForMaskedLM(AIDORNAPreTrainedModel):
+    def __init__(self, config):
+        super().__init__(config)
+        self.model = AIDORNAModel(config)
+        self.cls = AIDORNAOnlyMLMHead(config)
+        self.post_init()
+
+    def get_output_embeddings(self):
+        return self.cls.predictions.decoder
+
+    def set_output_embeddings(self, new_embeddings):
+        self.cls.predictions.decoder = new_embeddings
+
+    def get_input_embeddings(self):
+        return self.model.get_input_embeddings()
+
+    def set_input_embeddings(self, value):
+        self.model.set_input_embeddings(value)
+
+    def forward(
+        self,
+        input_ids=None,
+        attention_mask=None,
+        labels=None,
+        output_hidden_states=None,
+        output_attentions=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=input_ids,
+            attention_mask=attention_mask,
+            output_hidden_states=output_hidden_states,
+            output_attentions=output_attentions,
+            return_dict=return_dict,
+        )
+        logits = self.cls(out[0] if not return_dict else out.last_hidden_state)
+        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:
+            return ((loss, logits) + out[1:]) if loss is not None else (logits,) + out[1:]
+        return MaskedLMOutput(loss=loss, logits=logits, hidden_states=out.hidden_states, attentions=out.attentions)

special_tokens_map.json

@@ -0,0 +1,9 @@
+{
+  "bos_token": "[BOS]",
+  "cls_token": "[CLS]",
+  "eos_token": "[EOS]",
+  "mask_token": "[MASK]",
+  "pad_token": "[PAD]",
+  "sep_token": "[SEP]",
+  "unk_token": "[UNK]"
+}

tokenization_aidorna.py

@@ -0,0 +1,100 @@
+import os
+from typing import List, Optional
+
+from transformers import PreTrainedTokenizer
+
+_DEFAULT_VOCAB = [
+    "[PAD]", "[MASK]", "[CLS]", "[SEP]", "[UNK]",
+    "A", "G", "C", "T", "U", "N",
+    "[BOS]", "[EOS]", "[UNUSED1]", "[UNUSED2]", "[UNUSED3]",
+]
+
+
+class AIDORNATokenizer(PreTrainedTokenizer):
+    vocab_files_names = {"vocab_file": "vocab.txt"}
+    model_input_names = ["input_ids", "attention_mask"]
+
+    def __init__(
+        self,
+        vocab_file=None,
+        unk_token="[UNK]",
+        cls_token="[CLS]",
+        pad_token="[PAD]",
+        mask_token="[MASK]",
+        sep_token="[SEP]",
+        bos_token="[BOS]",
+        eos_token="[EOS]",
+        **kwargs,
+    ):
+        if vocab_file is not None and os.path.isfile(vocab_file):
+            with open(vocab_file) as f:
+                self.all_tokens = [line.strip() for line in f if line.strip()]
+        else:
+            self.all_tokens = list(_DEFAULT_VOCAB)
+
+        self._id_to_token = dict(enumerate(self.all_tokens))
+        self._token_to_id = {tok: idx for idx, tok in enumerate(self.all_tokens)}
+
+        super().__init__(
+            unk_token=unk_token,
+            cls_token=cls_token,
+            pad_token=pad_token,
+            mask_token=mask_token,
+            sep_token=sep_token,
+            bos_token=bos_token,
+            eos_token=eos_token,
+            **kwargs,
+        )
+
+        # Register all vocab tokens as no-split so the trie-based tokenizer matches
+        # single characters (A, G, C, T, U, N) and special tokens exactly.
+        self.unique_no_split_tokens = self.all_tokens
+        self._update_trie(self.unique_no_split_tokens)
+
+    @property
+    def vocab_size(self):
+        return len(self.all_tokens)
+
+    def get_vocab(self):
+        vocab = dict(self._token_to_id)
+        vocab.update(self.added_tokens_encoder)
+        return vocab
+
+    def _tokenize(self, text, **kwargs):
+        return text.split()
+
+    def _convert_token_to_id(self, token):
+        return self._token_to_id.get(token, self._token_to_id.get("[UNK]", 4))
+
+    def _convert_id_to_token(self, index):
+        return self._id_to_token.get(index, "[UNK]")
+
+    def build_inputs_with_special_tokens(self, token_ids_0, token_ids_1=None):
+        cls = [self.cls_token_id]
+        sep = [self.sep_token_id]
+        if token_ids_1 is None:
+            return cls + token_ids_0 + sep
+        return cls + token_ids_0 + sep + cls + token_ids_1 + sep
+
+    def get_special_tokens_mask(self, token_ids_0, token_ids_1=None, already_has_special_tokens=False):
+        if already_has_special_tokens:
+            return [1 if t in self.all_special_ids else 0 for t in token_ids_0]
+        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):
+        sep = [self.sep_token_id]
+        cls = [self.cls_token_id]
+        if token_ids_1 is None:
+            return [0] * (1 + len(token_ids_0) + 1)
+        return [0] * (1 + len(token_ids_0) + 1) + [1] * (1 + len(token_ids_1) + 1)
+
+    def save_vocabulary(self, save_directory, filename_prefix=None):
+        os.makedirs(save_directory, exist_ok=True)
+        fname = (filename_prefix + "-" if filename_prefix else "") + "vocab.txt"
+        path = os.path.join(save_directory, fname)
+        with open(path, "w") as f:
+            f.write("\n".join(self.all_tokens))
+        return (path,)

tokenizer_config.json

@@ -0,0 +1,77 @@
+{
+  "added_tokens_decoder": {
+    "0": {
+      "content": "[PAD]",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "1": {
+      "content": "[MASK]",
+      "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": "[UNK]",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "11": {
+      "content": "[BOS]",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "12": {
+      "content": "[EOS]",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    }
+  },
+  "bos_token": "[BOS]",
+  "clean_up_tokenization_spaces": false,
+  "cls_token": "[CLS]",
+  "eos_token": "[EOS]",
+  "extra_special_tokens": {},
+  "mask_token": "[MASK]",
+  "model_max_length": 1024,
+  "pad_token": "[PAD]",
+  "sep_token": "[SEP]",
+  "tokenizer_class": "AIDORNATokenizer",
+  "unk_token": "[UNK]",
+  "auto_map": {
+    "AutoTokenizer": [
+      "tokenization_aidorna.AIDORNATokenizer",
+      null
+    ]
+  }
+}

vocab.txt

@@ -0,0 +1,16 @@
+[PAD]
+[MASK]
+[CLS]
+[SEP]
+[UNK]
+A
+G
+C
+T
+U
+N
+[BOS]
+[EOS]
+[UNUSED1]
+[UNUSED2]
+[UNUSED3]