diff --git a/src/step1.1_count_tokens_length.py b/src/step1.1_count_tokens_length.py new file mode 100755 index 0000000..be6f4b1 --- /dev/null +++ b/src/step1.1_count_tokens_length.py @@ -0,0 +1,74 @@ +import os +import multiprocessing as mp +from glob import glob +import numpy as np +from transformers import AutoTokenizer +from datasets import load_dataset +from tqdm import tqdm + +DATA_DIR = "../corpus" +STATS_FILE = "token_stats.json" +NUM_PROC = min(32, os.cpu_count()) +SAMPLE_SIZE = 10000 #None + +tokenizer = AutoTokenizer.from_pretrained("../weights/weights_tokenizer") +if tokenizer.pad_token is None: + tokenizer.pad_token = tokenizer.eos_token + +def get_token_lengths(texts): + return [len(tokenizer.encode(t, add_special_tokens=False)) for t in texts] + +def process_file(file_path): + try: + dataset = load_dataset("parquet", data_files=file_path, split="train") + if SAMPLE_SIZE and len(dataset) > SAMPLE_SIZE: + dataset = dataset.shuffle().select(range(SAMPLE_SIZE)) + texts = dataset["text"] + lengths = get_token_lengths(texts) + return { + "file": os.path.basename(file_path), + "count": len(lengths), + "mean": np.mean(lengths), + "std": np.std(lengths), + "min": np.min(lengths), + "max": np.max(lengths), + "percentiles": { + str(p): np.percentile(lengths, p) + for p in [50, 75, 90, 95, 99] + }, + "all_lengths": lengths + } + except Exception as e: + print(f"处理 {file_path} 出错: {e}") + return None + +if __name__ == "__main__": + file_paths = glob(os.path.join(DATA_DIR, "*.parquet")) + print(f"找到 {len(file_paths)} 个parquet文件") + all_stats = [] + with mp.Pool(processes=NUM_PROC) as pool: + results = pool.imap(process_file, file_paths) + for stat in tqdm(results, total=len(file_paths)): + if stat is not None: + all_stats.append(stat) + + all_lengths = [] + for stat in all_stats: + all_lengths.extend(stat["all_lengths"]) + stat.pop("all_lengths") + + global_stats = { + "total_samples": len(all_lengths), + "global_mean": np.mean(all_lengths), + "global_max": np.max(all_lengths), + "global_95_percentile": np.percentile(all_lengths, 95), + "recommended_max_length": int(np.percentile(all_lengths, 95) * 1.05) + } + + print(f"\n{'='*50}") + print("全局统计结果:") + print(f"总样本数: {global_stats['total_samples']:,}") + print(f"平均token长度: {global_stats['global_mean']:.1f}") + print(f"最大token长度: {global_stats['global_max']}") + print(f"95百分位长度: {global_stats['global_95_percentile']}") + print(f"建议最大长度: {global_stats['recommended_max_length']}") diff --git a/test2.py b/src/step2_get_tokens_from_parquet.py similarity index 63% rename from test2.py rename to src/step2_get_tokens_from_parquet.py index 0e32fea..1784234 100755 --- a/test2.py +++ b/src/step2_get_tokens_from_parquet.py @@ -3,43 +3,33 @@ import multiprocessing as mp from glob import glob import numpy as np from transformers import AutoTokenizer -from datasets import load_dataset, Dataset +from datasets import load_dataset import pyarrow as pa import pyarrow.parquet as pq -import pandas as pd from tqdm import tqdm import gc import psutil import time -# 配置参数 -MODEL_PATH = "models/qwen2" # 本地模型路径 -DATA_DIR = "datas" # 原始数据目录 -OUTPUT_DIR = "tokenized_data" # tokenized数据输出目录 -MAX_SEQ_LENGTH = 1024 # 最大序列长度 -NUM_PROC = min(32, psutil.cpu_count(logical=False)) # 物理核心数 -CHUNK_SIZE = 50000 # 每个处理块的大小 -MAX_MEMORY_PERCENT = 70 # 最大内存使用百分比 - -# 创建输出目录 +DATA_DIR = "../corpus" +OUTPUT_DIR = "../corpus/tokens" +MAX_SEQ_LENGTH = 1024 +NUM_PROC = min(32, psutil.cpu_count(logical=False)) +CHUNK_SIZE = 50000 +MAX_MEMORY_PERCENT = 70 os.makedirs(OUTPUT_DIR, exist_ok=True) - -# 加载分词器 -tokenizer = AutoTokenizer.from_pretrained("/mnt/d/VMshare/LLM-from-scratch/weights/weights_tokenizer") +tokenizer = AutoTokenizer.from_pretrained("../weights/weights_tokenizer") if tokenizer.pad_token is None: tokenizer.pad_token = tokenizer.eos_token -# 获取所有parquet文件 -file_paths = ["/mnt/d/VMshare/LLM-from-scratch/corpus/parquet/train-00001-of-00399.parquet"] +file_paths = glob(os.path.join(DATA_DIR, "*.parquet")) print(f"找到 {len(file_paths)} 个parquet文件") def memory_safe(): - """检查内存是否安全""" mem = psutil.virtual_memory() return mem.percent < MAX_MEMORY_PERCENT def process_chunk(texts, tokenizer): - """处理文本块并返回tokenized结果""" tokenized = tokenizer( texts, truncation=True, @@ -50,31 +40,20 @@ def process_chunk(texts, tokenizer): return tokenized["input_ids"], tokenized["attention_mask"] def process_file(file_path): - """处理单个大型parquet文件""" try: base_name = os.path.basename(file_path) output_path = os.path.join(OUTPUT_DIR, f"tokenized_{base_name}") - - # 如果文件已处理则跳过 if os.path.exists(output_path): print(f"文件已存在: {output_path}, 跳过") return - print(f"开始处理: {base_name}") start_time = time.time() - - # 1. 获取文件总行数 dataset = load_dataset("parquet", data_files=file_path, split="train") total_rows = len(dataset) print(f"文件包含 {total_rows:,} 行") - - # 2. 创建空DataFrame用于收集结果 all_input_ids = [] all_attention_mask = [] - - # 3. 分块处理 num_chunks = (total_rows + CHUNK_SIZE - 1) // CHUNK_SIZE - for chunk_idx in range(num_chunks): if not memory_safe(): print("内存使用过高，等待释放...") @@ -83,58 +62,39 @@ def process_file(file_path): if not memory_safe(): print("内存仍然不足，跳过当前文件") return - start_row = chunk_idx * CHUNK_SIZE end_row = min((chunk_idx + 1) * CHUNK_SIZE, total_rows) - - # 读取当前块 chunk = dataset.select(range(start_row, end_row)) - texts = chunk["text"] # 根据实际字段名调整 - - # 处理当前块 + texts = chunk["text"] input_ids, attention_mask = process_chunk(texts, tokenizer) - - # 收集结果 all_input_ids.append(input_ids) all_attention_mask.append(attention_mask) - - # 每处理5个块显示进度 if (chunk_idx + 1) % 5 == 0 or (chunk_idx + 1) == num_chunks: elapsed = time.time() - start_time rows_done = min((chunk_idx + 1) * CHUNK_SIZE, total_rows) percent = rows_done / total_rows * 100 print(f"进度: {rows_done:,}/{total_rows:,} 行 ({percent:.1f}%) " f"耗时: {elapsed:.1f}秒") - - # 4. 合并所有块 all_input_ids = np.vstack(all_input_ids) all_attention_mask = np.vstack(all_attention_mask) - - # 5. 创建PyArrow表并保存 table = pa.Table.from_arrays([ pa.array(all_input_ids.reshape(-1, MAX_SEQ_LENGTH).tolist()), pa.array(all_attention_mask.reshape(-1, MAX_SEQ_LENGTH).tolist()) ], names=["input_ids", "attention_mask"]) - pq.write_table(table, output_path) - - # 6. 清理内存 del dataset, all_input_ids, all_attention_mask, table gc.collect() - elapsed = time.time() - start_time print(f"完成处理: {base_name} -> {output_path} " f"总耗时: {elapsed:.1f}秒, 平均速度: {total_rows/elapsed:.1f}行/秒") except Exception as e: print(f"处理文件 {file_path} 时出错: {str(e)}") - # 清理可能的残留资源 if 'dataset' in locals(): del dataset gc.collect() def process_file_wrapper(args): - """包装函数用于处理内存监控""" file_path, idx, total = args print(f"\n{'='*50}\n处理文件 [{idx+1}/{total}]: {os.path.basename(file_path)}\n{'='*50}") process_file(file_path) @@ -142,28 +102,7 @@ def process_file_wrapper(args): if __name__ == "__main__": print(f"开始预处理 {len(file_paths)} 个文件，使用 {NUM_PROC} 个进程") - print(f"每个文件约18万行，总数据量约 {len(file_paths)*180000/1000000:.1f} 百万行") - - # 准备参数 file_args = [(fp, idx, len(file_paths)) for idx, fp in enumerate(file_paths)] - - # 使用进程池并行处理 with mp.Pool(processes=NUM_PROC) as pool: results = list(tqdm(pool.imap(process_file_wrapper, file_args), total=len(file_paths))) - - print("所有文件处理完成!") - - # 生成元数据文件 - metadata = { - "model": MODEL_PATH, - "tokenized_date": time.strftime("%Y-%m-%d %H:%M:%S"), - "max_seq_length": MAX_SEQ_LENGTH, - "num_files": len(file_paths), - "output_dir": OUTPUT_DIR - } - - with open(os.path.join(OUTPUT_DIR, "metadata.json"), "w") as f: - import json - json.dump(metadata, f, indent=2) - - print(f"元数据已保存到 {OUTPUT_DIR}/metadata.json") \ No newline at end of file + print("所有文件处理完成!") \ No newline at end of file diff --git a/test.py b/test.py deleted file mode 100755 index 96a716f..0000000 --- a/test.py +++ /dev/null @@ -1,54 +0,0 @@ -import os -import torch -from transformers import ( - AutoTokenizer, - AutoModelForCausalLM, - TrainingArguments, - Trainer, - DataCollatorForLanguageModeling -) -from datasets import Dataset, load_from_disk -import deepspeed -import numpy as np -from glob import glob -import logging -from torch.utils.data import IterableDataset, DataLoader -import pyarrow.parquet as pq -# 配置参数 -MODEL_PATH = "models/qwen2" # 本地模型路径 -TOKENIZED_DATA_DIR = "corpus" # tokenized数据目录 -OUTPUT_DIR = "./output" # 训练输出目录 -DEEPSPEED_CONFIG = "ds_config.json" # DeepSpeed配置文件 -NUM_GPUS = 4 # GPU数量 -PER_DEVICE_BATCH_SIZE = 4 # 每张GPU的batch size -GRADIENT_ACCUM_STEPS = 8 # 梯度累积步数 -MAX_TRAIN_STEPS = 100000 # 最大训练步数 -SAVE_STEPS = 2000 # 保存间隔步数 -LOGGING_STEPS = 50 # 日志记录间隔 -LEARNING_RATE = 5e-5 # 学习率 -# 设置日志 -logging.basicConfig(level=logging.INFO) -logger = logging.getLogger(__name__) -# ===== 1. 准备模型和分词器 ===== -logger.info("加载模型和分词器...") -from transformers import AutoTokenizer -tokenizer = AutoTokenizer.from_pretrained("/mnt/d/VMshare/LLM-from-scratch/weights/weights_tokenizer") -file_path="tokenized_data/tokenized_train-00001-of-00399.parquet" -table = pq.read_table(file_path) -input_ids = table["input_ids"].to_numpy() -attention_mask = table["attention_mask"].to_numpy() -print(len(input_ids)) -for input_id in input_ids: - print(input_id) - print(len(input_id)) - res=tokenizer.decode(input_id) - # print(res) - break - -file_path2="/mnt/d/VMshare/LLM-from-scratch/corpus/parquet/train-00001-of-00399.parquet" -table2 = pq.read_table(file_path2) -text = table2["text"] -print(len(text)) -for value in text: - print(value) - break \ No newline at end of file