tokenizer.py

"""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


# Token vocabulary
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}  # Tokens that can be masked


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)

        # Build fast lookup table for encoding (ord -> token_id)
        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  # Reserve space for [CLS] and [SEP]

        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"])  # Unknown bases → 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)
            # Skip special tokens in decode
        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.
        """
        # Complement mapping as a tensor for gather
        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]
        # Reverse along the last dimension
        reversed_comp = complemented.flip(-1)
        return reversed_comp