replit/replit-code-v1-3b

1from typing import Union2from transformers import AutoTokenizer, PreTrainedTokenizer, PreTrainedTokenizerFast3Tokenizer = Union[PreTrainedTokenizer, PreTrainedTokenizerFast]4NUM_SENTINEL_TOKENS: int = 1005 6def adapt_tokenizer_for_denoising(tokenizer: Tokenizer):7    """Adds sentinel tokens and padding token (if missing).8 9    Expands the tokenizer vocabulary to include sentinel tokens10    used in mixture-of-denoiser tasks as well as a padding token.11 12    All added tokens are added as special tokens. No tokens are13    added if sentinel tokens and padding token already exist.14    """15    sentinels_to_add = [f'<extra_id_{i}>' for i in range(NUM_SENTINEL_TOKENS)]16    tokenizer.add_tokens(sentinels_to_add, special_tokens=True)17    if tokenizer.pad_token is None:18        tokenizer.add_tokens('<pad>', special_tokens=True)19        tokenizer.pad_token = '<pad>'20        assert tokenizer.pad_token_id is not None21    sentinels = ''.join([f'<extra_id_{i}>' for i in range(NUM_SENTINEL_TOKENS)])22    _sentinel_token_ids = tokenizer(sentinels, add_special_tokens=False).input_ids23    tokenizer.sentinel_token_ids = _sentinel_token_ids24 25class AutoTokenizerForMOD(AutoTokenizer):26    """AutoTokenizer + Adaptation for MOD.27 28    A simple wrapper around AutoTokenizer to make instantiating29    an MOD-adapted tokenizer a bit easier.30 31    MOD-adapted tokenizers have sentinel tokens (e.g., <extra_id_0>),32    a padding token, and a property to get the token ids of the33    sentinel tokens.34    """35 36    @classmethod37    def from_pretrained(cls, *args, **kwargs):38        """See `AutoTokenizer.from_pretrained` docstring."""39        tokenizer = super().from_pretrained(*args, **kwargs)40        adapt_tokenizer_for_denoising(tokenizer)41        return tokenizer