Tokenizers

dnallm.models.tokenizer

Classes

DNAOneHotTokenizer

DNAOneHotTokenizer(
    max_length=196608,
    padding_side="right",
    return_embeds=False,
    embeds_transpose=False,
)

Source code in dnallm/models/tokenizer.py

def __init__(
    self,
    max_length: int | None = 196_608,
    padding_side: str = "right",
    return_embeds: bool = False,
    embeds_transpose: bool = False,
):
    self.max_length = max_length
    self.padding_side = padding_side
    self.return_embeds = return_embeds
    self.embeds_transpose = embeds_transpose

    self.token_to_id = {
        "A": 0,
        "a": 0,
        "C": 1,
        "c": 1,
        "G": 2,
        "g": 2,
        "T": 3,
        "t": 3,
        "U": 3,
        "u": 3,  # RNA
        "N": 4,
        "n": 4,
        "X": 4,
        "x": 4,
        "-": -1,
        ".": -1,  # Padding
    }

    self.id_to_token = {0: "A", 1: "C", 2: "G", 3: "T", 4: "N", -1: "-"}

    self.pad_token_id = -1
    self.pad_token = "-"  # ruff: ignore[hardcoded-password-string]
    self.unk_token_id = 4  # N

    self.vocab_vectors = torch.tensor(
        [
            [1.0, 0.0, 0.0, 0.0],  # 0: A
            [0.0, 1.0, 0.0, 0.0],  # 1: C
            [0.0, 0.0, 1.0, 0.0],  # 2: G
            [0.0, 0.0, 0.0, 1.0],  # 3: T / U
            [0.25, 0.25, 0.25, 0.25],  # 4: N / X
            [0.0, 0.0, 0.0, 0.0],  # 5: Padding, index=-1
        ],
        dtype=torch.float32,
    )

Methods:

from_pretrained classmethod

from_pretrained(pretrained_model_name_or_path, **kwargs)

Load Tokenizer from a specified directory.

Source code in dnallm/models/tokenizer.py

@classmethod
def from_pretrained(cls, pretrained_model_name_or_path: str, **kwargs):
    """
    Load Tokenizer from a specified directory.
    """
    config_file = os.path.join(pretrained_model_name_or_path, "tokenizer_config.json")

    if os.path.isfile(config_file):
        with open(config_file, encoding="utf-8") as f:
            config = json.load(f)

        # Remove some metadata that should not be passed to __init__
        config.pop("tokenizer_class", None)
        config.pop("pad_token_id", None)
        config.pop("unk_token_id", None)

        config.update(kwargs)

        return cls(**config)
    else:
        print(
            "Warning: tokenizer_config.json not found",
            f"in {pretrained_model_name_or_path}. Using default init.",
        )
        return cls(**kwargs)

save_pretrained

save_pretrained(save_directory, **kwargs)

Save tokenizer configuration and vocabulary. Generates: tokenizer_config.json and vocab.json

Source code in dnallm/models/tokenizer.py

def save_pretrained(self, save_directory: str, **kwargs):
    """
    Save tokenizer configuration and vocabulary.
    Generates: tokenizer_config.json and vocab.json
    """
    os.makedirs(save_directory, exist_ok=True)

    config = {
        "max_length": self.max_length,
        "padding_side": self.padding_side,
        "pad_token_id": self.pad_token_id,
        "unk_token_id": self.unk_token_id,
        "tokenizer_class": self.__class__.__name__,
    }

    config_file = os.path.join(save_directory, "tokenizer_config.json")
    with open(config_file, "w", encoding="utf-8") as f:
        json.dump(config, f, indent=2, ensure_ascii=False)

    vocab_file = os.path.join(save_directory, "vocab.json")
    with open(vocab_file, "w", encoding="utf-8") as f:
        json.dump(self.token_to_id, f, indent=2, ensure_ascii=False)

    return [config_file, vocab_file]

Functions:

load_tokenizer_with_fallback

load_tokenizer_with_fallback(
    model_name,
    *,
    auto_tokenizer_cls=None,
    trust_remote_code=True,
    add_prefix_space=False,
)

Load a tokenizer compatible with transformers 4.x/5.x.

transformers v5 faithfully rebuilds the slow tokenizer class named in tokenizer_config.json (e.g. the Unigram-based DebertaV2Tokenizer) and crashes when tokenizer.json actually holds a BPE model; v4 mapped the declared slow class straight to its fast counterpart and never hit this. Fall back to PreTrainedTokenizerFast, which loads the same tokenizer.json file v4 loaded, then to DNAOneHotTokenizer.

Capability probing only — no transformers.__version__ checks.

Parameters:

Name	Type	Description	Default
model_name	str	Model name or local checkpoint directory.	required
auto_tokenizer_cls	Any	Primary AutoTokenizer class to try first (pass modules["AutoTokenizer"] to keep ModelScope loaders working); transformers' AutoTokenizer when omitted.	None
trust_remote_code	bool	Forwarded to every from_pretrained attempt.	True
add_prefix_space	bool	Forwarded to every attempt (token-classification tasks need it).	False

Returns:

Type	Description
Any	The loaded tokenizer, or a DNAOneHotTokenizer if all attempts fail.

Source code in dnallm/models/tokenizer.py

def load_tokenizer_with_fallback(
    model_name: str,
    *,
    auto_tokenizer_cls: Any = None,
    trust_remote_code: bool = True,
    add_prefix_space: bool = False,
) -> Any:
    """Load a tokenizer compatible with transformers 4.x/5.x.

    transformers v5 faithfully rebuilds the slow tokenizer class named in
    ``tokenizer_config.json`` (e.g. the Unigram-based ``DebertaV2Tokenizer``)
    and crashes when ``tokenizer.json`` actually holds a BPE model; v4 mapped
    the declared slow class straight to its fast counterpart and never hit
    this. Fall back to ``PreTrainedTokenizerFast``, which loads the same
    ``tokenizer.json`` file v4 loaded, then to ``DNAOneHotTokenizer``.

    Capability probing only — no ``transformers.__version__`` checks.

    Args:
        model_name: Model name or local checkpoint directory.
        auto_tokenizer_cls: Primary ``AutoTokenizer`` class to try first
            (pass ``modules["AutoTokenizer"]`` to keep ModelScope loaders
            working); transformers' ``AutoTokenizer`` when omitted.
        trust_remote_code: Forwarded to every ``from_pretrained`` attempt.
        add_prefix_space: Forwarded to every attempt (token-classification
            tasks need it).

    Returns:
        The loaded tokenizer, or a ``DNAOneHotTokenizer`` if all attempts fail.
    """
    kwargs: dict[str, Any] = {"trust_remote_code": trust_remote_code}
    if add_prefix_space:
        kwargs["add_prefix_space"] = True

    if auto_tokenizer_cls is not None:
        try:
            return auto_tokenizer_cls.from_pretrained(model_name, **kwargs)
        except Exception as e:
            logger.debug(f"AutoTokenizer failed for {model_name}: {e!r}")
    else:
        try:
            from transformers import AutoTokenizer  # ty: ignore[unresolved-import]  # transformers lazy export, resolves live

            return AutoTokenizer.from_pretrained(model_name, **kwargs)
        except Exception as e:
            logger.debug(f"AutoTokenizer failed for {model_name}: {e!r}")

    try:
        from transformers import PreTrainedTokenizerFast  # ty: ignore[unresolved-import]  # transformers lazy export, resolves live

        tokenizer = PreTrainedTokenizerFast.from_pretrained(model_name, **kwargs)
        logger.warning(
            f"AutoTokenizer failed for {model_name}; loaded fast tokenizer from tokenizer.json."
        )
        return tokenizer
    except Exception as e:
        logger.debug(f"PreTrainedTokenizerFast also failed for {model_name}: {e!r}")

    logger.warning(f"All tokenizer loading failed for {model_name}; using DNAOneHotTokenizer.")
    return DNAOneHotTokenizer()