| """Single-nucleotide tokenizer for SeqLens. |
| |
| Maps individual nucleotides to token IDs. No BPE, no k-mers — each base |
| is one token. This is the simplest tokenization strategy and matches |
| HyenaDNA, Caduceus, and Evo2. |
| """ |
|
|
| from typing import List, Optional |
|
|
| import torch |
|
|
|
|
| |
| VOCAB = { |
| "A": 0, "T": 1, "G": 2, "C": 3, "N": 4, |
| "[CLS]": 5, "[SEP]": 6, "[PAD]": 7, "[MASK]": 8, |
| } |
| ID_TO_TOKEN = {v: k for k, v in VOCAB.items()} |
| COMPLEMENT = {0: 1, 1: 0, 2: 3, 3: 2, 4: 4, 5: 5, 6: 6, 7: 7, 8: 8} |
| NUCLEOTIDE_IDS = {0, 1, 2, 3, 4} |
|
|
|
|
| class NucleotideTokenizer: |
| """Tokenizes raw DNA strings into integer token IDs. |
| |
| Usage: |
| tok = NucleotideTokenizer(max_len=16384) |
| ids = tok.encode("ATGCNATGC") # -> [0, 1, 2, 3, 4, 0, 1, 2, 3] |
| ids = tok.encode("ATGC", add_special=True) # -> [5, 0, 1, 2, 3, 6] |
| seq = tok.decode(ids) # -> "ATGCNATGC" |
| """ |
|
|
| def __init__(self, max_len: int = 16_384, pad_token_id: int = 7): |
| self.max_len = max_len |
| self.pad_token_id = pad_token_id |
| self.vocab_size = len(VOCAB) |
|
|
| |
| self._char_to_id = {} |
| for char in "ATGCNatgcn": |
| self._char_to_id[char] = VOCAB[char.upper()] |
|
|
| def encode( |
| self, |
| sequence: str, |
| add_special: bool = False, |
| max_len: Optional[int] = None, |
| ) -> List[int]: |
| """Encode a DNA string to token IDs. |
| |
| Args: |
| sequence: Raw DNA string (ATGCN characters). |
| add_special: If True, prepend [CLS] and append [SEP]. |
| max_len: Override max sequence length. Truncates if exceeded. |
| |
| Returns: |
| List of integer token IDs. |
| """ |
| max_len = max_len or self.max_len |
| ids = [] |
|
|
| if add_special: |
| ids.append(VOCAB["[CLS]"]) |
| max_len -= 2 |
|
|
| for char in sequence[:max_len]: |
| token_id = self._char_to_id.get(char) |
| if token_id is not None: |
| ids.append(token_id) |
| else: |
| ids.append(VOCAB["N"]) |
|
|
| if add_special: |
| ids.append(VOCAB["[SEP]"]) |
|
|
| return ids |
|
|
| def decode(self, token_ids: List[int]) -> str: |
| """Decode token IDs back to a DNA string.""" |
| chars = [] |
| for tid in token_ids: |
| token = ID_TO_TOKEN.get(tid, "N") |
| if token in ("A", "T", "G", "C", "N"): |
| chars.append(token) |
| |
| return "".join(chars) |
|
|
| def batch_encode( |
| self, |
| sequences: List[str], |
| add_special: bool = False, |
| pad: bool = True, |
| ) -> torch.Tensor: |
| """Encode and pad a batch of sequences. |
| |
| Args: |
| sequences: List of DNA strings. |
| add_special: Whether to add [CLS]/[SEP]. |
| pad: Whether to pad to max length in batch. |
| |
| Returns: |
| LongTensor of shape (B, L). |
| """ |
| encoded = [self.encode(seq, add_special=add_special) for seq in sequences] |
|
|
| if pad: |
| max_len = max(len(e) for e in encoded) |
| for i in range(len(encoded)): |
| pad_len = max_len - len(encoded[i]) |
| encoded[i] = encoded[i] + [self.pad_token_id] * pad_len |
|
|
| return torch.tensor(encoded, dtype=torch.long) |
|
|
| @staticmethod |
| def reverse_complement_ids(token_ids: torch.Tensor) -> torch.Tensor: |
| """Reverse complement a tensor of token IDs. |
| |
| Args: |
| token_ids: LongTensor of shape (..., L). |
| |
| Returns: |
| LongTensor of same shape with RC transformation applied. |
| """ |
| |
| comp_map = torch.tensor( |
| [1, 0, 3, 2, 4, 5, 6, 7, 8], |
| dtype=torch.long, |
| device=token_ids.device, |
| ) |
| complemented = comp_map[token_ids] |
| |
| reversed_comp = complemented.flip(-1) |
| return reversed_comp |
