Function gen_baike

utils/data_preprocess.py:49–111

(origin_file)

    return ""

def gen_baike(origin_file):
    baike_items = []
    eos_token = "[EOS]" 
    max_len = 512
    batch_size, batch_cnt = 2000000, 0
    with open(origin_file, "r", encoding="utf-8") as f:
        while True:
            line = f.readline()
            if not line:
                break

            item = ujson.loads(line)
            cur_txt, cur_len = [], 0

            if not item["title"]:
                continue

            temp_txt = f"{item['title']}：{process_none(item['summary'])}"

            cur_len += len(temp_txt)
            cur_txt.append(temp_txt)

            for section in item["sections"]:

                # 太长的截断不要了
                if cur_len > max_len:
                    break

                title = f"{section['title']}：" if section["title"] else ""
                temp_txt = f"{title}{process_none(section['content'])}"

                cur_len += len(temp_txt)
                cur_txt.append(temp_txt)
            temp_txt = normalize("NFKC", "".join(cur_txt))

            if len(temp_txt) > max_len:
                # 从 max_len 开始找第一个句号，叹号
                n, i = len(temp_txt), max_len
                while i < n and temp_txt[i] not in ("。", "！"):
                    i += 1
                temp_txt = "".join(temp_txt[0 : i + 1])

                # 添加 eos token
            temp_txt = f"{temp_txt}{eos_token}"

            baike_items.append(temp_txt)

            if len(baike_items) % batch_size == 0:

                chunk_data = split_txt_cropus_to_chunk_data(baike_items)
                tb = pa.Table.from_arrays([chunk_data], names=["text"])

                file_name = f"/root/autodl-tmp/data/baike/baike_chunk_512_5.6M_{batch_cnt}.parquet"
                pq.write_table(
                    table=tb,
                    where=file_name,
                    row_group_size=50000,
                )