Function encode_whitespaces

codegeex/megatron/tokenizer/tokenizer.py:26–34

Encode whitespaces to extra tokens. >>> encode_whitespaces('a\\n b\\n c', 10, 10) 'a\\n<|extratoken_10|>b\\n<|extratoken_11|>c'

(text: str, start_extra_id: int=10, max_len: int=10)

def encode_whitespaces(text: str, start_extra_id: int=10, max_len: int=10):
    """Encode whitespaces to extra tokens.

    >>> encode_whitespaces('a\\n  b\\n   c', 10, 10)
    'a\\n<|extratoken_10|>b\\n<|extratoken_11|>c'
    """
    for i in np.arange(max_len, 1, -1):
        text = text.replace(" " * i, f"<|extratoken_{start_extra_id + i - 2}|>")
    return text


def decode_whitespaces(text: str, start_extra_id: int=10, max_len: int=10):