Upload RNAElectra pretrained model weights and tokenizer

README.md

@@ -1,3 +1,35 @@
 ---
 license: apache-2.0
 ---
+# RNAElectra
+
+RNAElectra is a pretrained RNA language model for nucleotide-level sequence representation learning.
+
+## Load model
+
+```python
+import torch
+from transformers import AutoModel
+from tokenizer import NucEL_Tokenizer
+
+device = "cuda" if torch.cuda.is_available() else "cpu"
+
+model = AutoModel.from_pretrained(
+    "FreakingPotato/RNAElectra",
+    trust_remote_code=True
+).to(device)
+
+tokenizer = NucEL_Tokenizer.from_pretrained(
+    "FreakingPotato/RNAElectra",
+    trust_remote_code=True
+)
+
+sequence = "AUGCAUGCAUGCAUGC"
+inputs = tokenizer(sequence, return_tensors="pt")
+inputs = {k: v.to(device) for k, v in inputs.items()}
+
+with torch.no_grad():
+    outputs = model(**inputs)
+
+embeddings = outputs.last_hidden_state
+print(embeddings.shape)

config.json

@@ -0,0 +1,46 @@
+{
+  "architectures": [
+    "ModernBertModel"
+  ],
+  "attention_bias": false,
+  "attention_dropout": 0.0,
+  "bos_token_id": 50281,
+  "classifier_activation": "gelu",
+  "classifier_bias": false,
+  "classifier_dropout": 0.0,
+  "classifier_pooling": "cls",
+  "cls_token_id": 2,
+  "decoder_bias": true,
+  "deterministic_flash_attn": false,
+  "dtype": "float32",
+  "embedding_dropout": 0.0,
+  "eos_token_id": 50282,
+  "global_attn_every_n_layers": 3,
+  "global_rope_theta": 10000,
+  "hidden_activation": "gelu",
+  "hidden_size": 512,
+  "initializer_cutoff_factor": 2.0,
+  "initializer_range": 0.02,
+  "intermediate_size": 2048,
+  "layer_norm_eps": 1e-12,
+  "local_attention": 128,
+  "local_rope_theta": 1000,
+  "mask_token_id": 3,
+  "max_position_embeddings": 8192,
+  "mlp_bias": false,
+  "mlp_dropout": 0.0,
+  "model_type": "modernbert",
+  "norm_bias": false,
+  "norm_eps": 1e-12,
+  "num_attention_heads": 16,
+  "num_hidden_layers": 22,
+  "pad_token_id": 1,
+  "repad_logits_with_grad": false,
+  "sep_token_id": 50282,
+  "sparse_pred_ignore_index": -100,
+  "sparse_prediction": false,
+  "tie_word_embeddings": false,
+  "transformers_version": "4.57.3",
+  "unknown_token_id": 0,
+  "vocab_size": 27
+}

pytorch_model.bin

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

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]"
+}

tokenizer.py

@@ -0,0 +1,236 @@
+from typing import List, Dict, Optional, Union, Any, Tuple
+import os
+from transformers import PreTrainedTokenizer
+from itertools import product
+import json
+
+class NucEL_Tokenizer(PreTrainedTokenizer):
+    """
+    KMER Tokenizer for DNA sequences, inheriting from Hugging Face's PreTrainedTokenizer.
+    Handles k-mer tokenization with support for special tokens, padding, and truncation.
+    """
+    
+    model_input_names = ["input_ids", "attention_mask"]
+    
+    def __init__(
+        self,
+        k: int = 6,
+        model_max_length: int = 2048,
+        pad_token: str = "[PAD]",
+        unk_token: str = "[UNK]",
+        sep_token: str = "[SEP]",
+        cls_token: str = "[CLS]",
+        mask_token: str = "[MASK]",
+        bos_token: str = "[BOS]",
+        eos_token: str = "[EOS]",
+        num_reserved_tokens: int = 16,
+        **kwargs
+    ):
+        """Initialize the KMER tokenizer."""
+        self.k = k
+        self.nucleotides = ['A', 'C', 'G', 'T']
+        self.num_reserved_tokens = num_reserved_tokens
+        
+        # Define special tokens
+        self.special_tokens = {
+            "pad_token": pad_token,
+            "unk_token": unk_token,
+            "sep_token": sep_token,
+            "cls_token": cls_token,
+            "mask_token": mask_token,
+            "bos_token": bos_token,
+            "eos_token": eos_token,
+        }
+        
+        # Build vocabulary (includes special tokens, nucleotides, and k-mers)
+        self._init_vocabulary()
+        
+        # Now initialize the parent class.
+        super().__init__(
+            model_max_length=model_max_length,
+            pad_token=pad_token,
+            unk_token=unk_token,
+            sep_token=sep_token,
+            cls_token=cls_token,
+            mask_token=mask_token,
+            bos_token=bos_token,
+            eos_token=eos_token,
+            **kwargs
+        )
+
+    def _init_vocabulary(self):
+        """Initialize the vocabulary with special tokens, nucleotides, and k-mers."""
+        # Get special tokens in a specific order
+        special_tokens = [
+            self.special_tokens["pad_token"],
+            self.special_tokens["unk_token"],
+            self.special_tokens["cls_token"],
+            self.special_tokens["sep_token"],
+            self.special_tokens["mask_token"],
+            self.special_tokens["bos_token"],
+            self.special_tokens["eos_token"]
+        ]
+        
+        # Add individual nucleotides
+        nucleotides = self.nucleotides
+        
+        # Generate all possible k-mers
+        kmers = [''.join(p) for p in product(self.nucleotides, repeat=self.k)]
+        
+        # Add reserved tokens for future use
+        reserved_tokens = [f"[RESERVED_{i}]" for i in range(self.num_reserved_tokens)]
+
+        # Combine all tokens in a specific order
+        all_tokens = special_tokens + nucleotides + kmers + reserved_tokens
+        
+        # Create vocabulary: token -> index
+        self.vocab = {}
+        for idx, token in enumerate(all_tokens):
+            self.vocab[token] = idx
+            
+        # Create reverse mapping: index -> token
+        self.ids_to_tokens = {idx: token for token, idx in self.vocab.items()}
+
+    @property
+    def vocab_size(self) -> int:
+        """Return the size of vocabulary."""
+        return len(self.vocab)
+
+    def get_vocab(self) -> Dict[str, int]:
+        """Return the vocabulary dictionary."""
+        return self.vocab.copy()
+
+    def _tokenize(self, text: str) -> List[str]:
+        """
+        Tokenize a DNA sequence into k-mers and individual nucleotides.
+        
+        Args:
+            text: DNA sequence to tokenize
+            
+        Returns:
+            List of tokens.
+        """
+        text = text.upper().strip()
+        tokens = [self.cls_token]
+        i = 0
+        
+        while i < len(text):
+            # Try to get a k-mer
+            if i <= len(text) - self.k:
+                kmer = text[i:i+self.k]
+                if kmer in self.vocab:
+                    tokens.append(kmer)
+                    i += self.k
+                    continue
+            
+            # Fallback: tokenize a single nucleotide
+            if i < len(text):
+                nucleotide = text[i]
+                if nucleotide in self.nucleotides:
+                    tokens.append(nucleotide)
+                else:
+                    tokens.append(self.unk_token)
+                i += 1
+                
+        return tokens
+
+    def _convert_token_to_id(self, token: str) -> int:
+        """Convert a token to its ID in the vocabulary."""
+        return self.vocab.get(token, self.vocab[self.unk_token])
+
+    def _convert_id_to_token(self, index: int) -> str:
+        """Convert an ID to its token in the vocabulary."""
+        return self.ids_to_tokens.get(index, self.unk_token)
+
+    def save_vocabulary(self, save_directory: str, filename_prefix: Optional[str] = None) -> Tuple[str]:
+        """Save the tokenizer vocabulary to a directory."""
+        if not filename_prefix:
+            filename_prefix = "vocab"
+
+        vocab_file = os.path.join(save_directory, f"{filename_prefix}.json")
+        
+        with open(vocab_file, 'w', encoding='utf-8') as f:
+            json.dump(self.vocab, f, ensure_ascii=False, indent=2)
+            
+        return (vocab_file,)
+
+    def save_pretrained(self, save_directory: str, legacy_format: bool = True, filename_prefix: Optional[str] = None, **kwargs):
+        """
+        Save the tokenizer configuration and vocabulary.
+        """
+        # Save the vocabulary
+        vocab_files = self.save_vocabulary(save_directory, filename_prefix=filename_prefix)
+        
+        # Save the config
+        config = {
+            'k': self.k,
+            'model_max_length': self.model_max_length,
+            'padding_side': self.padding_side,
+            'truncation_side': self.truncation_side,
+            'special_tokens': {
+                'pad_token': self.pad_token,
+                'unk_token': self.unk_token,
+                'sep_token': self.sep_token,
+                'cls_token': self.cls_token,
+                'mask_token': self.mask_token,
+                'bos_token': self.bos_token,
+                'eos_token': self.eos_token,
+            }
+        }
+        
+        super().save_pretrained(save_directory, config=config, legacy_format=legacy_format, **kwargs)
+        
+        return vocab_files
+
+    @classmethod
+    def from_pretrained(cls, pretrained_model_name_or_path: Union[str, os.PathLike], *init_inputs, **kwargs):
+        """
+        Load a tokenizer from a pretrained model.
+        """
+        from huggingface_hub import hf_hub_download
+        
+        # Check if it's a local path or HuggingFace repo
+        if os.path.isdir(pretrained_model_name_or_path):
+            # Local directory
+            config_file = os.path.join(pretrained_model_name_or_path, "tokenizer_config.json")
+            vocab_file = os.path.join(pretrained_model_name_or_path, "vocab.json")
+        else:
+            # HuggingFace Hub
+            config_file = hf_hub_download(
+                repo_id=pretrained_model_name_or_path,
+                filename="tokenizer_config.json"
+            )
+            vocab_file = hf_hub_download(
+                repo_id=pretrained_model_name_or_path,
+                filename="vocab.json"
+            )
+        
+        # Load config
+        with open(config_file, 'r', encoding='utf-8') as f:
+            config = json.load(f)
+        
+        # Load vocab
+        with open(vocab_file, 'r', encoding='utf-8') as f:
+            vocab = json.load(f)
+        
+        k = config.get('k')
+
+        # Create tokenizer instance - tokens are at top level in tokenizer_config.json
+        tokenizer = cls(
+            k=k,
+            model_max_length=config.get('model_max_length', 2048),
+            pad_token=config.get('pad_token', '[PAD]'),
+            unk_token=config.get('unk_token', '[UNK]'),
+            sep_token=config.get('sep_token', '[SEP]'),
+            cls_token=config.get('cls_token', '[CLS]'),
+            mask_token=config.get('mask_token', '[MASK]'),
+            bos_token=config.get('bos_token', '[BOS]'),
+            eos_token=config.get('eos_token', '[EOS]'),
+            **kwargs
+        )
+        
+        # Override the vocabulary with the saved one
+        tokenizer.vocab = vocab
+        tokenizer.ids_to_tokens = {idx: token for token, idx in vocab.items()}
+        
+        return tokenizer

tokenizer_config.json

@@ -0,0 +1,71 @@
+{
+  "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
+    },
+    "5": {
+      "content": "[BOS]",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "6": {
+      "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": 1025,
+  "pad_token": "[PAD]",
+  "sep_token": "[SEP]",
+  "tokenizer_class": "NucEL_Tokenizer",
+  "unk_token": "[UNK]"
+}

vocab.json

@@ -0,0 +1,29 @@
+{
+  "[PAD]": 0,
+  "[UNK]": 1,
+  "[CLS]": 2,
+  "[SEP]": 3,
+  "[MASK]": 4,
+  "[BOS]": 5,
+  "[EOS]": 6,
+  "A": 11,
+  "C": 12,
+  "G": 13,
+  "T": 14,
+  "[RESERVED_0]": 15,
+  "[RESERVED_1]": 16,
+  "[RESERVED_2]": 17,
+  "[RESERVED_3]": 18,
+  "[RESERVED_4]": 19,
+  "[RESERVED_5]": 20,
+  "[RESERVED_6]": 21,
+  "[RESERVED_7]": 22,
+  "[RESERVED_8]": 23,
+  "[RESERVED_9]": 24,
+  "[RESERVED_10]": 25,
+  "[RESERVED_11]": 26,
+  "[RESERVED_12]": 27,
+  "[RESERVED_13]": 28,
+  "[RESERVED_14]": 29,
+  "[RESERVED_15]": 30
+}