Function tokenize

training/preprocess.py:17–61

Tokenize the input text and split into chunks of specified context length. Args: examples: Dictionary containing the input text. tokenizer: Initialized tokenizer. seq_len: Total sequence length for each training sample. Default: 2

(
    examples: Dict[str, List[Any]],
    tokenizer: AutoTokenizer,
    seq_len: int = 2048,
    ctx_len: int = None,
    return_offsets: bool = False
)

def tokenize(
    examples: Dict[str, List[Any]],
    tokenizer: AutoTokenizer,
    seq_len: int = 2048,
    ctx_len: int = None,
    return_offsets: bool = False
) -> Dict[str, List[List[int]]]:
    """
    Tokenize the input text and split into chunks of specified context length.

    Args:
        examples:
            Dictionary containing the input text.
        tokenizer:
            Initialized tokenizer.
        seq_len:
            Total sequence length for each training sample. Default: 2048.
        ctx_len:
            Max contiguous length to preserve (will not be split). Default: `None`.
        return_offsets:
            Return cumulative offsets for concatenated inputs. Default: `False`.

    Returns:
        Dictionary containing tokenized and chunked input ids, and optionally offsets.
    """
    text = examples['text']
    input_ids = tokenizer(text)['input_ids']
    # further split each input into chunks of length `ctx_len` if provided
    if ctx_len is not None:
        input_ids = [seq[i:i+ctx_len] for seq in input_ids for i in range(0, len(seq), ctx_len)]
    lens = torch.tensor([len(seq) for seq in input_ids]).cumsum(0)
    total_len = lens[-1] // seq_len * seq_len

    input_ids = list(chain(*input_ids))
    # each yielded sample is of length `seq_len`
    input_ids = [input_ids[i:i+seq_len] for i in range(0, total_len, seq_len)]

    if not return_offsets:
        return {'input_ids': input_ids}

    # insert boundaries into cumulative offsets
    offsets = torch.cat((lens, torch.arange(0, total_len, seq_len))).unique().sort()[0] % seq_len
    # split offsets according the start positions
    offsets = [i.tolist() + [seq_len] for i in offsets.tensor_split(torch.where(offsets.eq(0))[0][1:])][:len(input_ids)]
    return {'input_ids': input_ids, 'offsets': offsets}


def preprocess(