pretrain_data_process.py

from dataclasses import dataclass, field
from typing import Optional
import math
import glob
import torch
import os
from transformers import HfArgumentParser, LlamaTokenizer

## 禁止HuggingFace联网，加快加载本地数据集的速度
os.environ['HF_DATASETS_OFFLINE'] = '1'
import datasets

@dataclass
class DataArgs:
    model_name_or_path: str = field(default='')   ### tokenizer所在目录
    data_path: str = field(default=None)   ### 待预处理的数据所在目录
    save_dir: str = field(default=None)   ### 预处理后的数据的存放目录
    max_length: Optional[int] = field(default=2048)   ### 每个样本的最大长度
    cache_dir: str = field(default='')   ### hf数据集缓存目录
    num_group: Optional[str] = field(default=10000)   ###concat时，每隔batch包含的样本数量
    num_proc: Optional[int] = field(default=32)

parser = HfArgumentParser(DataArgs)
data_args = parser.parse_args_into_dataclasses()[0]


def tokenizer_fn(tokenizer):
    def tokenize(line):
        ## 使用tokenizer对text进行词令化
        input_ids = tokenizer(
            line['text'],
            return_tensors='pt',
            return_attention_mask=False  
        )['input_ids'][0]

        return {
            "input_ids": input_ids,
        }
    return tokenize

def concat_multiple_sample_fn(max_length, pad_token_id):
    def concat_multiple_sample(batch):
        ## cat需要接收一个List[torch.tensor], 不接收List[list]
        concat_input_ids = torch.cat(batch['input_ids'], dim=0)

        all_length = concat_input_ids.size(0)
        chunks = math.ceil(all_length / max_length)
        ## 拼接的样本长度不足max_length的部分, 使用pad_token_id进行填充
        pad_length = chunks * max_length - all_length
        pad_seq = torch.ones(pad_length, dtype=concat_input_ids.dtype) * pad_token_id
        concat_input_ids = torch.cat([concat_input_ids, pad_seq], dim=0)

        ## chunk返回一个tuple[torch.tensor]，需要转成List[torch.tensor]
        input_ids = torch.chunk(concat_input_ids, chunks)

        return {
            "input_ids": list(input_ids),
        }
    return concat_multiple_sample

def tokenize_and_group_chunk(data_path, save_dir, tokenizer):
    ## 获取目录下的所有json文件名
    filenames = glob.glob(f"{data_path}/*.json") + glob.glob(f"{data_path}/*.jsonl")

    print("\nStep1: load json dataset")
    ## 第一步：加载数据集
    data = datasets.load_dataset(
        "json",
        data_files=filenames,  ### 待加载的文件列表
        num_proc=data_args.num_proc,  ### 并行加载的进程数
        cache_dir=data_args.cache_dir  ### 数据集缓存的目录
    )['train']

    print("\nStep2: Tokenizing")
    ## 第二步：对每个sample进行词令化
    data = data.map(
        tokenizer_fn(tokenizer),
        num_proc=data_args.num_proc,
        desc='tokenize'
    )
    data = data.select_columns("input_ids")
    data.set_format(type="torch")

    print("\nStep3: concat and group")
    ## 第三步：对多个sample进行concat
    concat_data = data.map(
        concat_multiple_sample_fn(data_args.max_length, tokenizer.pad_token_id),
        batched=True,   ### 是否对多个sample进行concat
        batch_size=data_args.num_group,  ### 每个batch包含的样本数
        num_proc=data_args.num_proc,   ### 并行处理的进程数
        drop_last_batch=False,   ### 是否丢弃最后一个batch，其可能不足num_group
        desc="concat_group"
    )

    print("\nStep4: save to disk")
    ## 第四步：将预先concat好的数据保存到本地磁盘， 待后续预训练时高效加载
    concat_data.save_to_disk(save_dir, max_shard_size="500MB")  ## 每个分片最大500MB


if __name__=="__main__":
    ## 加载LLaMA模型专属的tokenizer
    tokenizer = LlamaTokenizer.from_pretrained(data_args.model_name_or_path)
    tokenizer.add_eos_token = True   ### 增加<eos>的token标记
    tokenizer.pad_token = tokenizer.eos_token


    ## 执行预处理函数
    tokenize_and_group_chunk(data_args.data_path, data_args.save_dir, tokenizer)