mindspeed-llm是昇腾模型套件代码仓，原来叫"modelLink"。这篇文章带大家阅读一下数据处理脚本preprocess_data.py（基于1.0.0分支），数据处理是模型训练的第一步，经常会用到。
文章中贴的源码加了相关注释，同学们可以把源码和注释结合起来看。
首先来看一下main函数
def main():
# 获取入参，通过后面的代码可以知道有哪些关键参数
args = get_args()
# 参数校验
validate_args(args)
# 合并已经处理好的数据集
if args.merge_group_keys is not None:
merge_datasets(args)
return
# 创建splitter，用来把文章段落分割成句子
splitter = build_splitter(args)
# 创建tokenizer，用来把句子切分成单个的词
tokenizer = build_tokenizer(args)
logger.info("building dataset: %s", args.input)
# 加载数据，把CSV、JSON、TXT等格式的数据加载到内存
raw_data = build_dataset(args)
# 保存到一个文件
if args.n_subs == 1:
# 获取处理后的数据句柄
handler = get_dataset_handler(args, raw_data, tokenizer, splitter)
# 数据落盘
handler.serialize_to_disk()
# 保存到多个文件，使用多进程处理，单文件的处理方式和if条件中是一致的
else:
target_prefix = args.output_prefix
target_prefixname = os.path.basename(target_prefix)
num_samples = len(raw_data)
start_ends = cut_range_to_subs(num_samples, num_samples // args.n_subs)
subsets = [raw_data.select(range(x[0], x[1])) for x in start_ends]
# multiprocessing
params_list = []
for k, subset in enumerate(subsets):
args_ = copy.deepcopy(args)
args_.output_prefix = target_prefix.replace(target_prefixname, f'{str(k).zfill(3)}_of_{str(len(subsets)-1).zfill(3)}_{target_prefixname}')
params = [args_, subset, tokenizer, splitter]
params_list.append(params)
pool = multiprocessing.Pool()
sub_idx_files = pool.map(handle_subset, params_list)
pool.close()
pool.join()
for key in sub_idx_files[0].keys():
idx_files = [x[key] for x in sub_idx_files]
idx_files.sort()
target_idx = idx_files[0].replace(f'000_of_{str(len(subsets)-1).zfill(3)}_{target_prefixname}', target_prefixname)
target_bin = target_idx.replace('.idx', '.bin')
idx = IndexedDatasetBuilder(target_bin)
for idx_file in idx_files:
idx.add_index(idx_file.replace('.idx', ''))
idx.finalize(target_idx)
for idx_file in idx_files:
os.remove(idx_file)
os.remove(idx_file.replace('.idx', '.bin'))可以看到，main函数处理逻辑主要由这几个函数组成：build_splitter、build_tokenizer、build_dataset、get_dataset_handler、serialize_to_disk。
build_splitter
这个函数的功能是把文字段落分割成单个句子，查看源码，主要使用的是三方库nltk的函数：
def build_splitter(args):
if nltk and args.split_sentences:
nltk.download("punkt", quiet=True)
if args.split_sentences:
if not nltk:
logger.error("NLTK is not available to split sentences.")
raise Exception("nltk is not available")
splitter = nltk.load("tokenizers/punkt/english.pickle")
if args.keep_newlines:
# this prevents punkt from eating newlines after sentences
final_splitter = nltk.tokenize.punkt.PunktSentenceTokenizer(
train_text=splitter._params,
lang_vars=CustomLanguageVars())
else:
final_splitter = splitter
else:
# 自定义splitter
final_splitter = IdentitySplitter()
return final_splitterbuild_tokenizer
这个函数的主要功能是把句子切分成单个的词，比如说把 "今天是星期几" 切分成 "今天"、"是"、"星期几"，然后转成对应的整数。
def build_tokenizer(args):
"""Initialize tokenizer."""
# 获取huggingface的tokenizer
if args.tokenizer_type == "PretrainedFromHF":
if args.rank == 0:
print(' > building PretrainFromHF tokenizer. Vocab file is un-used, '
'loading tokenizer from pre-trained model', flush=True)
if args.tokenizer_name_or_path is None:
raise ValueError("Missing tokenizer_name_or_path while building PretrainFromHF tokenizer.")
hf_tokenizer_kwargs = dict()
if hasattr(args, "tokenizer_kwargs") and args.tokenizer_kwargs:
if len(args.tokenizer_kwargs) % 2 != 0:
raise ValueError("The token name and token value must be entered in pairs.")
for i in range(0, len(args.tokenizer_kwargs), 2):
hf_tokenizer_kwargs[args.tokenizer_kwargs[i]] = \
args.tokenizer_kwargs[i + 1]
# 基于MegatronTokenizer构建的类
tokenizer = _AutoTokenizer(
args.tokenizer_name_or_path,
vocab_extra_ids=args.vocab_extra_ids,
model_max_length=args.seq_length,
use_fast=args.tokenizer_not_use_fast,
**hf_tokenizer_kwargs
)
# Add vocab size (if not already set from a checkpoint).
if getattr(args, "padded_vocab_size", None) is None:
args.padded_vocab_size = _vocab_size_with_padding(tokenizer.vocab_size,
args)
else:
#
tokenizer = TokenizerAdaptor(megatron_build_tokenizer(args))
# 根据prompt_type完善tokenizer
if hasattr(args, "prompt_type") and args.prompt_type is not None:
if ("PreTrainedTokenizerBase" not in str(tokenizer.tokenizer._pad.__func__)):
tokenizer.tokenizer._pad = MethodType(PreTrainedTokenizerBase._pad, tokenizer.tokenizer)
tokenizer.tokenizer.padding_side = "right"
fix_model_tokenizer(tokenizer.tokenizer, args.prompt_type.strip())
return tokenizer分成了2类tokenizer，一类是PretrainedFromHF，也就是使用预训练的 HuggingFace 分词器；如果不适用hf的，则则使用 TokenizerAdaptor 类和 megatron_build_tokenizer 函数创建分词器实例 tokenizer。
build_dataset
这个函数的功能是把数据文件加载到内存，返回DatasetDict 或Dataset，也就是一个Python容器。这个函数中调用的load_dataset是huggingface的datasets库的函数。
def build_dataset(args):
"""loading dataset by huggingface"""
raw_datasets = None
if args.handler_name == "LlamaFactoryInstructionHandler":
all_datasets = []
for dataset_attr in get_dataset_list(args):
# 加载单个数据集
all_datasets.append(load_single_dataset(dataset_attr, args))
# 合并数据集
raw_datasets = merge_dataset(all_datasets, args)
else:
if args.handler_name == "MOSSInstructionHandler" or args.handler_name == "MOSSMultiTurnHandler":
# for MOSS, streaming is needed.流式加载数据
args.streaming = True
if args.hf_datasets_params:
with open(args.hf_datasets_params, 'r') as fin:
param_dict = json.load(fin)
return load_dataset(**param_dict)
cache_dir = args.cache_dir
split_flag = "train"
load_from_local = os.path.exists(args.input)
# 从本地加载
if load_from_local:
# args.input 是一个有效的 Python 脚本路径
if _has_py_script(args.input):
logger.info("loading data from a local python script")
raw_datasets = load_dataset(
args.input,
data_dir='./' if not args.script_data_dir else args.script_data_dir,
split=split_flag,
num_proc=None if args.streaming else args.workers,
cache_dir=cache_dir,
streaming=args.streaming,
trust_remote_code=False
)
else:
# args.input 是一个文件或目录路径
data_files = [args.input] if os.path.isfile(args.input) else \
glob.glob(os.path.join(args.input, '*'))
# 获取文件格式
ext, data_format = _get_data_format(data_files)
# 筛选合法的文件格式
filtered_data_files = list(filter(lambda x: x.split('.')[-1] == ext, data_files))
if filtered_data_files:
logger.info("loading data from local file, format: %s,"
" file num: %s", data_format, len(data_files))
raw_datasets = load_dataset(
data_format,
split=split_flag,
data_files=filtered_data_files,
num_proc=None if args.streaming else args.workers,
cache_dir=cache_dir,
streaming=args.streaming,
trust_remote_code=False
)
else:
raise Exception("unknown local data!")
else:
logger.info("loading data from remote huggingface") # 从远程 Hugging Face 数据集加载
raw_datasets = load_dataset(
args.input,
split=split_flag,
num_proc=None if args.streaming else args.workers,
cache_dir=cache_dir,
streaming=args.streaming,
trust_remote_code=False
)
if raw_datasets is None:
raise Exception("unknown data!")
if args.handler_name in [
"AlpacaStyleInstructionHandler",
"SharegptStyleInstructionHandler",
"AlpacaStylePairwiseHandler",
"SharegptStylePairwiseHandler"
]:
handler_dataset_attr = get_handler_dataset_attr(args, raw_datasets)
return align_dataset(raw_datasets, handler_dataset_attr, args)
return raw_datasetsget_dataset_handler
这个函数的功能是创建数据集处理实例，_get_handler_cls会根据args.handler_name选择对应的handler。handler的基类和子类都在mindspeed_llm/tasks/preprocess/data_handler.py里面定义了，查看BaseDatasetHandler可以知道，这个类的对外函数有这几个：get_tokenized_data、serialize_to_disk，功能分别是对数据进行令牌化]、数据序列化。
serialize_to_disk
接着上面讲，这个函数是handler的类函数，用于将分词后的数据集保存到磁盘。具体来说，它将数据集的每个样本（或句子）序列化为二进制文件，并生成相应的索引文件。代码如下：
def _serialize_to_disk(self, iteration_batch_size=50):
startup_start = time.time()
if not self.tokenized_dataset:
self.tokenized_dataset = self.get_tokenized_data()
output_bin_files = {} # 保存数据的文件路径
output_idx_files = {} # 保存数据的文件路径
builders = {} # 用于构建索引数据集的字典
level = "document"
if self.args.split_sentences:
level = "sentence"
logger.info("Vocab size: %s", self.tokenizer.vocab_size)
logger.info("Output prefix: %s", self.args.output_prefix)
# 字典的key就是文件名，json_keys就是类似"input_ids", "attention_mask", "labels"的string
for key in self.args.json_keys:
output_bin_files[key] = f"{self.args.output_prefix}_{key}_{level}.bin"
output_idx_files[key] = f"{self.args.output_prefix}_{key}_{level}.idx"
# vocab_size=None : use int32 dtype for -100 will be used in labels
# 为每个文件创建一个数据字典
builders[key] = indexed_dataset.IndexedDatasetBuilder(output_bin_files[key])
self.output_idx_files = output_idx_files
startup_end = time.time()
proc_start = time.time()
total_bytes_processed = 0
logger.info("Time to startup:%s", startup_end - startup_start)
skip_num = 0
# 遍历每个文件的内容
for i, doc in enumerate(self.tokenized_dataset.iter(batch_size=iteration_batch_size), start=1):
# In post-training stage, we need to drop the data exceeded set sequence-length
skip_indices = set()
# 进行一次筛选
for key in self.args.json_keys:
batch = [sentences for sentences in doc[key] if len(sentences) > 0]
if len(batch) == 0:
continue
for j, sentences in enumerate(batch):
for k, sentence in enumerate(sentences):
if self.args.seq_length is not None and len(sentence) >= self.args.seq_length:
skip_indices.add((j, k))
# 正式开始处理每个句子
for key in self.args.json_keys:
batch = [sentences for sentences in doc[key] if len(sentences) > 0]
if len(batch) == 0:
continue
for j, sentences in enumerate(batch):
for k, sentence in enumerate(sentences):
if (j, k) in skip_indices:
skip_num = skip_num + 1
continue
# 记录处理的字节数
total_bytes_processed += len(sentence) * np.int32().itemsize
# 把合法的句子加到builders里面
builders[key].add_item(sentence)
builders[key].end_document()
batch_id = i * iteration_batch_size
if batch_id % self.args.log_interval == 0:
current = time.time()
elapsed = current - proc_start
mbs = total_bytes_processed / elapsed / 1024 / 1024
logger.info("Processed %s documents (%s docs/s, %s MB/s).", batch_id, batch_id / elapsed, mbs)
logger.info("Skip %s sample exceeded seq-length(%s)", skip_num / len(self.args.json_keys), self.args.seq_length)
for key in self.args.json_keys:
builders[key].finalize(output_idx_files[key])以上就是mindspeed-llm处理数据的主要函数了，大家还有什么想了解的呢？欢迎评论区提问！
本文由博客一文多发平台 OpenWrite 发布！