Function group_texts

utils/preprocess_dataset.py:81–111

(examples)

        return output
    
    def group_texts(examples):
        # Concatenate all texts.
        concatenated_examples = {k: sum(examples[k], []) for k in examples.keys()}
        total_length = len(concatenated_examples[list(examples.keys())[0]])
        if total_length >= block_size:
            total_length = (total_length // block_size) * block_size

        input_ids = []
        attention_mask = []
        labels = []
        # We implement a sliding window, so all tokens have a non-zero context in their prediction.
        # We then mask the duplicate tokens' labels, to not count any token twice in the loss.
        for i in tqdm(range(0, total_length, stride), total=total_length):
            begin_loc = max(i + stride - block_size, 0)
            end_loc = min(i + stride, total_length)
            trg_len = end_loc - i
            cur_input_ids = concatenated_examples['input_ids'][begin_loc:end_loc]
            cur_labels = list(cur_input_ids)
            cur_labels[:-trg_len] = [padding_index] * (len(cur_labels) - trg_len)

            if len(cur_input_ids) < block_size:
                padding_size = block_size - len(cur_input_ids)
                pad_token_id = tokenizer.pad_token_id if tokenizer.pad_token_id is not None else tokenizer.eos_token_id
                cur_input_ids += [pad_token_id] * padding_size
                cur_labels += [padding_index] * padding_size
            input_ids.append(cur_input_ids)
            attention_mask.append([1] * len(cur_labels))
            labels.append(cur_labels)

        result = {'input_ids': input_ids, 'labels': labels, 'attention_mask': attention_mask}
        return result

    tokenized_datasets = raw_datasets.map(
        tokenize_function,