FreakingPotato
/

NucEL

+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.
+        """
+        # Load the tokenizer configuration
+        config_file = os.path.join(pretrained_model_name_or_path, "tokenizer_config.json")
+        with open(config_file, 'r', encoding='utf-8') as f:
+            config = json.load(f)
+        # Load the vocabulary
+        vocab_file = os.path.join(pretrained_model_name_or_path, "vocab.json")
+        with open(vocab_file, 'r', encoding='utf-8') as f:
+            vocab = json.load(f)
+        # Extract k from config (add it to your tokenizer_config.json if not present)
+        k = config.get('k', 6)
+        # 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