[论文笔记] DCLM 分长度区间进行长文本抽取

import os
import zstandard as zstd
import json
import io
import multiprocessing
from tqdm import tqdm

# 定义根目录路径
root_dir = "dclm-baseline-1.0"
output_base_dir = "dclm"  # 输出的基准路径

# 定义长度区间 (字符数)
length_ranges = {
    "dclm_16k": (8000, 16000),
    "dclm_32k": (16000, 32000),
    "dclm_64k": (32000, 64000),
    "dclm_128k": (64000, 128000)
}

# 处理函数
def process_file(file_path):
    relative_dir = os.path.relpath(file_path, root_dir)
    relative_dir = relative_dir.replace(os.path.basename(file_path), "")
    with open(file_path, 'rb') as f:
        dctx = zstd.ZstdDecompressor()
        with dctx.stream_reader(f) as reader:
            text_stream = io.TextIOWrapper(reader, encoding='utf-8')
            for line in text_stream:
                record = json.loads(line)
                text_length = len(record['text'])