HCIE-AI：数据工程实验手册

import argparse
import json
import os
import numpy as np

from mindspore.mindrecord import FileWriter

from baichuan2_tokenizer import Baichuan2Tokenizer


IGNORE_TOKEN_ID = -100


def preprocess(sources, tokenizer, seq_length, user_tokens=195, assistant_tokens=196):
    """conversation preprocess."""
    input_ids = []
    labels = []
    # 原始数据是一列表,列表里面一组字典是一轮对话, 因为目的是学习 gpt,所以需要对用户的话和特殊标识进行 mask 标识
    for example in sources:
        input_id = []
        label = []
        for message in example["conversations"]:
            from_ = message["from"]
            value = message["value"]
            # 文本编码
            value_ids = tokenizer.encode(value, add_special_tokens=False)
            
            # 分别对用户的话和 GPT 回答进行处理
            if from_ == "human":
                input_id += [user_tokens] + value_ids
                # eos_token_id 相当于一轮对话的分隔符. 其余 mask
                label += [tokenizer.eos_token_id] + [IGNORE_TOKEN_ID] * len(value_ids)
            else:
                input_id += [assistant_tokens] + value_ids
                # GPT 回答需要学习,所以不需要mask
                label += [IGNORE_TOKEN_ID] + value_ids

        input_id.append(tokenizer.eos_token_id)
        label.append(tokenizer.eos_token_id)
        # 如果编码长度超出 seq_length,则需要截断
        if len(input_id) > seq_length:
            input_id = input_id[: seq_length]
            label = label[: seq_length]
        else:
            # 没有超出 seq_length,则需要 padding. padding字符是不需要学习的,所以mask
            input_id += [tokenizer.pad_token_id] * (seq_length - len(input_id))
            label += [IGNORE_TOKEN_ID] * (seq_length - len(label))

        input_ids.append(np.array(input_id).astype(np.int32))
        labels.append(np.array(label).astype(np.int32))

    return dict(
        input_ids=input_ids,
        labels=labels
    )


class SupervisedDataset:
    """Dataset for supervised fine-tuning."""

    def __init__(self, raw_data, tokenizer, seq_length):
        super(SupervisedDataset, self).__init__()

        data_dict = preprocess(raw_data, tokenizer, seq_length)

        self.input_ids = data_dict["input_ids"]
        self.labels = data_dict["labels"]

    def __len__(self):
        return len(self.input_ids)

    def __getitem__(self, i):
        return dict(
            input_ids=self.input_ids[i],
            labels=self.labels[i]
        )


def tokenize_qa(tokenizer, file_path, seq_length):
    # 从文件以 json 格式加载数据
    raw_data = json.load(open(file_path, "r"))
    # 数据转化为Dataset
    dataset_cls = SupervisedDataset(raw_data, tokenizer, seq_length)
    # 生成器
    for i in range(len(dataset_cls)):
        yield dataset_cls[i]


if __name__ == '__main__':
    # 程序入参接收
    parser = argparse.ArgumentParser()
    parser.add_argument("--mindrecord_schema", type=str, default="belle_baichuan2")
    parser.add_argument("--input_glob", type=str, default="./belle_chat_ramdon_10k.json")
    parser.add_argument("--output_file", type=str, default="./belle512.mindrecord")
    parser.add_argument("--model_file", type=str, default="./tokenizer.model")
    parser.add_argument("--file_partition", type=int, default=1)
    parser.add_argument("--seq_length", type=int, default=512)
    args = parser.parse_args()

    # 输出文件保存的位置
    out_dir, out_file = os.path.split(os.path.abspath(args.output_file))
    if not os.path.exists(out_dir):
        os.mkdir(out_dir)

    # 输出文件的 schema 信息
    schema = {'input_ids': {"type": "int32", "shape": [-1]},
              'labels': {"type": "int32", "shape": [-1]}}
    # 初始化FileWriter
    writer = FileWriter(file_name=args.output_file,
                        shard_num=args.file_partition)
    # 添加文件的schema信息
    writer.add_schema(schema, args.mindrecord_schema)

    # Start to load tokenizer
    if not os.path.exists(args.model_file):
        raise FileNotFoundError(f"file {args.model_file} do not exists.")

    transforms_count = 0

    # 加载分词器
    word_tokenizer = Baichuan2Tokenizer(vocab_file=args.model_file)
    # 将文件内容使用分词器分词,编码,并且遵循模型微调所需要的数据格式.
    # 将数据逐行写入
    for x in tokenize_qa(word_tokenizer, args.input_glob, args.seq_length + 1):
        transforms_count += 1
        writer.write_raw_data([x])
    print("Transformed {} records.".format(transforms_count))

    # flush 内存,写入文件
    writer.commit()
    out_file = args.output_file
    if args.file_partition > 1:
        out_file += '0'
    print("Transform finished, output files refer: {}".format(out_file))