tokenizer.py

import math
import numpy as np
from dataclasses import dataclass
from typing import Any, Dict, List, Optional, Union

import torch
from tokenizers import Tokenizer
from tokenizers.models import WordPiece
from tokenizers.pre_tokenizers import Split
from tokenizers.processors import TemplateProcessing
from transformers import BatchEncoding, PreTrainedTokenizerFast


@dataclass
class TokenizerConfig:
    """Configuration used to build a :class:`ProteinTokenizer`.

    Attributes:
        vocab: Ordered list of tokens (index = token ID).
        pad_token: Padding token string.
        mask_token: Masking token string used for MLM.
        bos_token: Beginning-of-sequence token string.
        eos_token: End-of-sequence token string.
        unk_token: Unknown token string.
        other_special_tokens: Additional special tokens to register.
        ambiguous_tokens: Amino-acid tokens considered ambiguous (e.g. B, X, Z).
        remove_ambiguous: If ``True``, strip ambiguous tokens during tokenization.
        vocab_size: Optional explicit vocab size (inferred from ``vocab`` if omitted).
    """

    vocab: List[str]
    pad_token: str
    mask_token: str
    bos_token: str
    eos_token: str
    unk_token: str
    other_special_tokens: Optional[List[str]] = None
    ambiguous_tokens: Optional[List[str]] = None
    remove_ambiguous: bool = False
    vocab_size: Optional[int] = None


class ProteinTokenizer(PreTrainedTokenizerFast):
    """Character-level tokenizer for protein amino acid sequences.

    The tokenizer preserves Hugging Face fast-tokenizer behavior while adding
    AMPLIFY-specific preprocessing:
    - optional random-window truncation
    - optional removal of ambiguous amino-acid tokens
    - ``position_ids`` generation aligned with model inputs
    """

    def __init__(self, config: Optional[TokenizerConfig] = None, max_length: int = 2048, **kwargs):
        """Build from config or load from Hugging Face tokenizer assets.

        Args:
            config: Training-time tokenizer config. If ``None``, this instance is
                initialized from pretrained files via ``kwargs``.
            max_length: Default sequence length limit used by truncation/padding.
            **kwargs: Forwarded to :class:`PreTrainedTokenizerFast`.
        """
        if config is not None:
            # Training-time path: build a character-level WordPiece backend from vocab.
            token_to_id = {token: i for i, token in enumerate(config.vocab)}
            tokenizer_object = Tokenizer(WordPiece(vocab=token_to_id, unk_token=config.unk_token))

            # Split on every character so each amino acid becomes one token.
            tokenizer_object.pre_tokenizer = Split("", behavior="removed")

            # Automatically prepend BOS and append EOS (following ESM convention).
            tokenizer_object.post_processor = TemplateProcessing(
                single=f"{config.bos_token}:0 $A:0 {config.eos_token}:0",
                special_tokens=[
                    (config.bos_token, token_to_id[config.bos_token]),
                    (config.eos_token, token_to_id[config.eos_token]),
                ],
            )

            super().__init__(
                model_max_length=max_length,
                padding_side="right",
                truncation_side="right",
                pad_token=config.pad_token,
                bos_token=config.bos_token,
                eos_token=config.eos_token,
                unk_token=config.unk_token,
                mask_token=config.mask_token,
                model_input_names=["input_ids", "attention_mask", "position_ids"],
                tokenizer_object=tokenizer_object,
            )

            # Register additional special tokens (e.g. domain-specific markers).
            if config.other_special_tokens:
                for token in config.other_special_tokens:
                    if token not in token_to_id:
                        raise ValueError(f"other_special_tokens contains '{token}' which is not in the vocabulary")
                self.add_special_tokens({"additional_special_tokens": config.other_special_tokens})

            # Convert ambiguous amino-acid tokens (e.g. B, X, Z) to their IDs.
            if config.ambiguous_tokens is not None:
                ambiguous_token_ids = [token_to_id[tok] for tok in config.ambiguous_tokens]
            else:
                ambiguous_token_ids = []
            remove_ambiguous = config.remove_ambiguous
        else:
            # Pretrained path: restore from saved tokenizer assets.
            ambiguous_token_ids = kwargs.pop("ambiguous_token_ids", [])
            remove_ambiguous = kwargs.pop("remove_ambiguous", False)
            super().__init__(**kwargs)

        self.ambiguous_token_ids = list(ambiguous_token_ids)
        self.remove_ambiguous = remove_ambiguous

        # Persist AMPLIFY-specific fields in tokenizer_config.json.
        self.init_kwargs["ambiguous_token_ids"] = self.ambiguous_token_ids
        self.init_kwargs["remove_ambiguous"] = self.remove_ambiguous

        # Per-field padding values used by ``pad()``.
        self.key_to_padding = {
            "input_ids": self.pad_token_id,
            "attention_mask": 0,
            "special_tokens_mask": 1,
            "position_ids": 0,
        }

    def _truncate_sequences(
        self,
        encoded_inputs: Dict[str, List[List[int]]],
        max_length: Optional[int] = None,
        random_truncate: bool = True,
    ) -> Dict[str, List[List[int]]]:
        """Truncate sequences longer than ``max_length``, preserving BOS/EOS.

        The first and last tokens (BOS/EOS) are kept intact; only the content
        between them is windowed, following ESM convention.

        Args:
            encoded_inputs: Tokenizer outputs with list-valued fields.
            max_length: Maximum sequence length (including BOS/EOS).
            random_truncate: If ``True``, sample a random content window;
                otherwise keep the left-most window.

        Returns:
            Truncated encoded inputs.
        """
        if max_length is None:
            return encoded_inputs

        for i, sequence in enumerate(encoded_inputs["input_ids"]):
            if len(sequence) > max_length:
                # Preserve BOS (first) and EOS (last); truncate content in between.
                content_max = max_length - 2
                if random_truncate:
                    offset = np.random.randint(0, len(sequence) - 2 - content_max + 1)
                else:
                    offset = 0
                for key in encoded_inputs:
                    vals = encoded_inputs[key][i]
                    encoded_inputs[key][i] = [vals[0]] + vals[1 + offset : 1 + offset + content_max] + [vals[-1]]

        return encoded_inputs

    def _remove_ambiguous_tokens(self, encoded_inputs: Dict[str, List[List[int]]]) -> Dict[str, List[List[int]]]:
        """Drop tokens listed in ``self.ambiguous_token_ids`` from all fields.

        Args:
            encoded_inputs: Tokenizer outputs with list-valued fields.

        Returns:
            Filtered encoded inputs.
        """
        filtered_inputs = {key: [] for key in encoded_inputs}

        for i, sequence in enumerate(encoded_inputs["input_ids"]):
            keep = [j for j, token in enumerate(sequence) if token not in self.ambiguous_token_ids]
            for key in encoded_inputs:
                filtered_inputs[key].append([encoded_inputs[key][i][j] for j in keep])

        return filtered_inputs

    def pad(
        self,
        encoded_inputs: Dict[str, List[List[int]]],
        padding: Union[bool, str] = True,
        max_length: Optional[int] = None,
        pad_to_multiple_of: Optional[int] = 8,
        return_tensors: Optional[str] = None,
        **kwargs,
    ) -> Union[Dict[str, List[List[int]]], BatchEncoding]:
        """Pad batch fields to a shared sequence length and optionally return tensors.

        Args:
            encoded_inputs: Tokenizer outputs or list of sample dictionaries.
            padding: If ``"longest"`` or ``True``, pad to the longest sequence.
            max_length: Maximum length to pad to.
            pad_to_multiple_of: Round target length up to this multiple.
            return_tensors: If set (e.g. ``"pt"``), wrap output in :class:`BatchEncoding`.

        Returns:
            Padded encoded inputs, optionally as :class:`BatchEncoding`.
        """
        # Convert list of sample dicts to dict of lists.
        if isinstance(encoded_inputs, list):
            encoded_inputs = {key: [d[key] for d in encoded_inputs] for key in encoded_inputs[0]}

        if max_length is None:
            max_length = self.model_max_length

        sequence_lengths = [len(sequence) for sequence in encoded_inputs["input_ids"]]
        if not sequence_lengths:
            if return_tensors is not None:
                return BatchEncoding(encoded_inputs, tensor_type=return_tensors)
            return encoded_inputs

        # Cap target length to the longest sequence in the batch.
        if padding == "longest" or padding is True:
            max_length = min(max_length, max(sequence_lengths))

        # Round up to nearest multiple for hardware-friendly tensor shapes.
        if pad_to_multiple_of is not None:
            max_length = math.ceil(max_length / pad_to_multiple_of) * pad_to_multiple_of

        # Pad each sequence to the target length with field-appropriate values.
        for i, seq_len in enumerate(sequence_lengths):
            pad_length = max_length - seq_len
            if pad_length > 0:
                for key in encoded_inputs:
                    encoded_inputs[key][i] += [self.key_to_padding[key]] * pad_length

        if return_tensors is not None:
            return BatchEncoding(encoded_inputs, tensor_type=return_tensors)

        return encoded_inputs

    def prepare_packed_batch(
        self,
        batch: Dict[str, Any],
        device: Optional[Union[str, torch.device]] = None,
    ) -> Dict[str, Any]:
        """Pack a tokenized batch for FlashAttention varlen kernels.

        Accepts padded tensor batches ``(B, S)`` with ``attention_mask``,
        or ragged ``List[List[int]]`` from tokenizer outputs.

        Args:
            batch: Tokenizer output dict containing at least ``input_ids``.
                May also contain ``attention_mask``, ``position_ids``, and ``labels``.
            device: Target device for output tensors.

        Returns:
            Dict with:
                - ``input_ids``: ``(1, total_tokens)`` — all sequences concatenated.
                - ``position_ids``: ``(1, total_tokens)`` — per-sequence positions.
                - ``cu_seqlens``: ``(num_sequences + 1,)`` — cumulative lengths.
                - ``max_seqlen``: ``int`` — longest sequence in the batch.
                - ``labels``: ``(1, total_tokens)`` — only present if input has labels.
        """
        input_ids = batch["input_ids"]

        # Detect input format to normalize all fields uniformly.
        is_2d_tensor = isinstance(input_ids, torch.Tensor) and input_ids.dim() == 2
        is_single = not is_2d_tensor and not isinstance(input_ids[0], (list, tuple, torch.Tensor))

        def normalize(field):
            """Normalize a batch field to a list of per-sequence items."""
            if field is None:
                return None
            if is_2d_tensor and isinstance(field, torch.Tensor):
                return list(field)
            if is_single:
                return [field]
            return field

        input_ids = normalize(input_ids)
        mask_raw = normalize(batch.get("attention_mask"))
        pos_raw = normalize(batch.get("position_ids"))
        labels_raw = normalize(batch.get("labels"))

        target = torch.device(device) if device is not None else torch.as_tensor(input_ids[0]).device

        # Strip padding (via attention_mask) and collect per-sequence tensors.
        packed_ids, packed_pos, packed_labels = [], [], []
        for i, ids in enumerate(input_ids):
            ids = torch.as_tensor(ids, dtype=torch.long)
            if mask_raw is not None:
                m = torch.as_tensor(mask_raw[i], dtype=torch.bool)
                ids = ids[m]
            packed_ids.append(ids.to(target))

            if pos_raw is not None:
                pos_i = torch.as_tensor(pos_raw[i], dtype=torch.long)
                if mask_raw is not None:
                    pos_i = pos_i[m]
                packed_pos.append(pos_i.to(target))
            else:
                packed_pos.append(torch.arange(len(ids), device=target))

            if labels_raw is not None:
                lab_i = torch.as_tensor(labels_raw[i], dtype=torch.long)
                if mask_raw is not None:
                    lab_i = lab_i[m]
                packed_labels.append(lab_i.to(target))

        # Concatenate into a single flat sequence and build cu_seqlens.
        seq_lengths = torch.tensor([len(s) for s in packed_ids], dtype=torch.int32, device=target)
        cu_seqlens = torch.cat([torch.zeros(1, dtype=torch.int32, device=target), seq_lengths])
        cu_seqlens = cu_seqlens.cumsum(dim=0, dtype=torch.int32)

        result = {
            "input_ids": torch.cat(packed_ids).unsqueeze(0),
            "position_ids": torch.cat(packed_pos).unsqueeze(0),
            "cu_seqlens": cu_seqlens,
            "max_seqlen": seq_lengths.max().item(),
        }
        if packed_labels:
            result["labels"] = torch.cat(packed_labels).unsqueeze(0)
        return result

    def __call__(
        self,
        text: Union[str, List[str]],
        max_length: Optional[int] = None,
        padding: Union[bool, str] = False,
        pad_to_multiple_of: Optional[int] = None,
        truncation: bool = False,
        random_truncate: bool = True,
        return_special_tokens_mask: bool = False,
        return_tensors: Optional[str] = None,
        **kwargs,
    ) -> Union[Dict[str, torch.Tensor], BatchEncoding]:
        """Tokenize protein sequences with optional truncation/filtering/padding.

        Args:
            text: Single sequence or list of sequences.
            max_length: Maximum sequence length.
            padding: Padding strategy.
            pad_to_multiple_of: If set, pad sequence length up to a multiple of this value.
            truncation: Whether to truncate sequences.
            random_truncate: If ``True``, truncate at random offset.
            return_special_tokens_mask: Whether to return special tokens mask.
            return_tensors: If ``"pt"``, return PyTorch tensors.

        Returns:
            Dict with ``input_ids``, ``attention_mask``, ``position_ids`` and
            optionally ``special_tokens_mask``. Returns :class:`BatchEncoding`
            when ``return_tensors`` is set.
        """
        # Wrap single string as a list for uniform handling.
        single_input = isinstance(text, str)
        if single_input:
            text = [text]

        # Tokenize each amino-acid sequence into token IDs.
        encoded_inputs = super().__call__(
            text,
            padding=False,
            truncation=False,
            return_special_tokens_mask=return_special_tokens_mask,
            **kwargs,
        )

        if max_length is None:
            max_length = self.model_max_length

        if truncation:
            encoded_inputs = self._truncate_sequences(
                encoded_inputs,
                max_length=max_length,
                random_truncate=random_truncate,
            )

        # Generate position IDs (0, 1, 2, ...) for each sequence.
        encoded_inputs["position_ids"] = [list(range(len(seq))) for seq in encoded_inputs["input_ids"]]

        if self.remove_ambiguous and self.ambiguous_token_ids:
            encoded_inputs = self._remove_ambiguous_tokens(encoded_inputs)

        if padding:
            encoded_inputs = self.pad(
                encoded_inputs,
                padding=padding,
                max_length=max_length,
                pad_to_multiple_of=pad_to_multiple_of,
            )

        # Unwrap single-input results back from list to plain values.
        if single_input and return_tensors is None:
            for key in encoded_inputs:
                encoded_inputs[key] = encoded_inputs[key][0]

        if return_tensors is not None:
            return BatchEncoding(encoded_inputs, tensor_type=return_tensors)

        return encoded_inputs