replit/replit-code-v1_5-3b

1from typing import Any2from transformers import AutoTokenizer, PreTrainedTokenizerBase3NUM_SENTINEL_TOKENS: int = 1004 5def adapt_tokenizer_for_denoising(tokenizer: PreTrainedTokenizerBase) -> None:6    """Adds sentinel tokens and padding token (if missing).7 8    Expands the tokenizer vocabulary to include sentinel tokens9    used in mixture-of-denoiser tasks as well as a padding token.10 11    All added tokens are added as special tokens. No tokens are12    added if sentinel tokens and padding token already exist.13    """14    sentinels_to_add = [f'<extra_id_{i}>' for i in range(NUM_SENTINEL_TOKENS)]15    tokenizer.add_tokens(sentinels_to_add, special_tokens=True)16    if tokenizer.pad_token is None:17        tokenizer.add_tokens('<pad>', special_tokens=True)18        tokenizer.pad_token = '<pad>'19        assert tokenizer.pad_token_id is not None20    sentinels = ''.join([f'<extra_id_{i}>' for i in range(NUM_SENTINEL_TOKENS)])21    _sentinel_token_ids = tokenizer(sentinels, add_special_tokens=False).input_ids22    tokenizer.sentinel_token_ids = _sentinel_token_ids23 24class AutoTokenizerForMOD(AutoTokenizer):25    """AutoTokenizer + Adaptation for MOD.26 27    A simple wrapper around AutoTokenizer to make instantiating28    an MOD-adapted tokenizer a bit easier.29 30    MOD-adapted tokenizers have sentinel tokens (e.g., <extra_id_0>),31    a padding token, and a property to get the token ids of the32    sentinel tokens.33    """34 35    @classmethod36    def from_pretrained(cls, *args: Any, **kwargs: Any) -> PreTrainedTokenizerBase:37        """See `AutoTokenizer.from_pretrained` docstring."""38        tokenizer = super().from_pretrained(*args, **kwargs)39        adapt_tokenizer_for_denoising(tokenizer)40        return tokenizer