diff --git a/src/configuration_dolly.py b/src/configuration_dolly.py index 7976365..d5dab6d 100755 --- a/src/configuration_dolly.py +++ b/src/configuration_dolly.py @@ -28,9 +28,9 @@ class DollyConfig(PretrainedConfig): def __init__( self, - vocab_size=32000, - hidden_size=2048, - intermediate_size=6144, + vocab_size=151936, + hidden_size=1024, + intermediate_size=3072, num_hidden_layers=28, num_attention_heads=16, num_key_value_heads=8, @@ -57,10 +57,10 @@ class DollyConfig(PretrainedConfig): self.num_hidden_layers = num_hidden_layers self.num_attention_heads = num_attention_heads self.use_sliding_window = use_sliding_window - self.sliding_window = sliding_window # we check `use_sliding_window` in the modeling code + self.sliding_window = sliding_window self.max_window_layers = max_window_layers - # for backward compatibility + if num_key_value_heads is None: num_key_value_heads = num_attention_heads diff --git a/src/step0_train_tokenizer_jsonl.py b/src/step0_train_tokenizer_jsonl.py index eba17ad..1fb86d4 100755 --- a/src/step0_train_tokenizer_jsonl.py +++ b/src/step0_train_tokenizer_jsonl.py @@ -1,6 +1,6 @@ import json import glob -from tokenizers import Tokenizer,Tokenizer, normalizers, pre_tokenizers, processors, decoders +from tokenizers import Tokenizer,Tokenizer, normalizers, pre_tokenizers, processors, decoders,Regex from tokenizers.normalizers import NFC from tokenizers.models import BPE from tokenizers.trainers import BpeTrainer @@ -76,7 +76,7 @@ tokenizer.normalizer = normalizers.Sequence([NFC()]) regex_pattern = r"""(?i:'s|'t|'re|'ve|'m|'ll|'d)|[^\r\n\p{L}\p{N}]?[^\s\p{L}\p{N}]+[\r\n]*|\s*[\r\n]+|\s+""" tokenizer.pre_tokenizer = pre_tokenizers.Sequence([ pre_tokenizers.Split( - pattern=pre_tokenizers.Regex(regex_pattern), + pattern=Regex(regex_pattern), behavior="isolated", invert=False ), @@ -100,7 +100,7 @@ tokenizer.decoder = decoders.ByteLevel( ) trainer = BpeTrainer( - vocab_size=32000, + vocab_size=151936, min_frequency=1, special_tokens=SPECIAL_TOKENS, show_progress=True, diff --git a/src/step0_train_tokenizer_parquet.py b/src/step0_train_tokenizer_parquet.py index 439baf5..4c84351 100755 --- a/src/step0_train_tokenizer_parquet.py +++ b/src/step0_train_tokenizer_parquet.py @@ -1,5 +1,5 @@ import glob -from tokenizers import Tokenizer, normalizers, pre_tokenizers, processors, decoders +from tokenizers import Tokenizer, normalizers, pre_tokenizers, processors, decoders,Regex from tokenizers.normalizers import NFC from tokenizers.models import BPE from tokenizers.trainers import BpeTrainer @@ -68,7 +68,7 @@ tokenizer.normalizer = normalizers.Sequence([NFC()]) regex_pattern = r"""(?i:'s|'t|'re|'ve|'m|'ll|'d)|[^\r\n\p{L}\p{N}]?[^\s\p{L}\p{N}]+[\r\n]*|\s*[\r\n]+|\s+""" tokenizer.pre_tokenizer = pre_tokenizers.Sequence([ pre_tokenizers.Split( - pattern=pre_tokenizers.Regex(regex_pattern), + pattern=Regex(regex_pattern), behavior="isolated", invert=False ), @@ -92,7 +92,7 @@ tokenizer.decoder = decoders.ByteLevel( ) trainer = BpeTrainer( - vocab_size=32000, + vocab_size=151936, min_frequency=1, special_tokens=SPECIAL_TOKENS, show_progress=True, diff --git a/src/step1.1_count_tokens_length.py b/src/step1.1_count_tokens_length.py index be6f4b1..dc3bb54 100755 --- a/src/step1.1_count_tokens_length.py +++ b/src/step1.1_count_tokens_length.py @@ -7,9 +7,8 @@ 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 +NUM_PROC = min(32, os.cpu_count()) +SAMPLE_SIZE = 10000 #None tokenizer = AutoTokenizer.from_pretrained("../weights/weights_tokenizer") if tokenizer.pad_token is None: @@ -36,7 +35,7 @@ def process_file(file_path): str(p): np.percentile(lengths, p) for p in [50, 75, 90, 95, 99] }, - "all_lengths": lengths + "all_lengths": lengths } except Exception as e: print(f"处理 {file_path} 出错: {e}") @@ -51,20 +50,20 @@ if __name__ == "__main__": 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") - + 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) + "recommended_max_length": int(np.percentile(all_lengths, 95) * 1.05) } - + print(f"\n{'='*50}") print("全局统计结果:") print(f"总样本数: {global_stats['total_samples']:,}") diff --git a/src/step2_get_tokens_from_parquet.py b/src/step2_get_tokens_from_parquet.py index 1784234..c804d4c 100755 --- a/src/step2_get_tokens_from_parquet.py +++ b/src/step2_get_tokens_from_parquet.py @@ -1,108 +1,69 @@ -import os -import multiprocessing as mp -from glob import glob -import numpy as np -from transformers import AutoTokenizer -from datasets import load_dataset -import pyarrow as pa -import pyarrow.parquet as pq -from tqdm import tqdm -import gc -import psutil -import time -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("../weights/weights_tokenizer") +from transformers import AutoTokenizer, set_seed +from datasets import load_dataset, DatasetDict +from accelerate import Accelerator +import os +import logging +import glob + +logging.basicConfig(level=logging.INFO) +logger = logging.getLogger(__name__) + +accelerator = Accelerator() +set_seed(42) + +TOKENS_DATA_DIR = "../corpus" +os.makedirs(TOKENS_DATA_DIR, exist_ok=True) + +model_name = "../weights/weights_tokenizer" +resume_option = None +tokenizer = AutoTokenizer.from_pretrained(model_name) if tokenizer.pad_token is None: tokenizer.pad_token = tokenizer.eos_token -file_paths = glob(os.path.join(DATA_DIR, "*.parquet")) -print(f"找到 {len(file_paths)} 个parquet文件") +data_files = [] +supported_formats = ['.jsonl', '.parquet', '.txt', '.json', '.csv'] +for format in supported_formats: + data_files.extend(glob.glob(f"../train/*{format}")) -def memory_safe(): - mem = psutil.virtual_memory() - return mem.percent < MAX_MEMORY_PERCENT - -def process_chunk(texts, tokenizer): - tokenized = tokenizer( - texts, - truncation=True, - max_length=MAX_SEQ_LENGTH, - padding="max_length", - return_tensors="np" - ) - return tokenized["input_ids"], tokenized["attention_mask"] - -def process_file(file_path): - 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() - dataset = load_dataset("parquet", data_files=file_path, split="train") - total_rows = len(dataset) - print(f"文件包含 {total_rows:,} 行") - all_input_ids = [] - all_attention_mask = [] - num_chunks = (total_rows + CHUNK_SIZE - 1) // CHUNK_SIZE - for chunk_idx in range(num_chunks): - if not memory_safe(): - print("内存使用过高，等待释放...") - time.sleep(30) - gc.collect() - 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"] - input_ids, attention_mask = process_chunk(texts, tokenizer) - all_input_ids.append(input_ids) - all_attention_mask.append(attention_mask) - 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}秒") - all_input_ids = np.vstack(all_input_ids) - all_attention_mask = np.vstack(all_attention_mask) - 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) - 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}行/秒") +if not data_files: + logger.error("未找到任何支持的语料文件！支持的格式: .jsonl, .parquet, .txt, .json, .csv") + exit(1) - except Exception as e: - print(f"处理文件 {file_path} 时出错: {str(e)}") - if 'dataset' in locals(): - del dataset - gc.collect() +logger.info(f"找到 {len(data_files)} 个训练语料文件: {data_files}") -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) - gc.collect() +file_ext = os.path.splitext(data_files[0])[1].lower() -if __name__ == "__main__": - print(f"开始预处理 {len(file_paths)} 个文件，使用 {NUM_PROC} 个进程") - 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("所有文件处理完成!") \ No newline at end of file +if file_ext in ['.txt', '.text']: + logger.info("检测到文本文件格式，使用'text'加载器") + dataset = load_dataset('text', data_files=data_files, split='train') + dataset = dataset.map(lambda examples: {'text': examples['text']}, batched=True) +else: + logger.info(f"检测到结构化文件格式 ({file_ext})，使用自动推断加载器") + dataset = load_dataset(file_ext.lstrip('.'), data_files=data_files, split='train') + +if 'text' not in dataset.column_names: + logger.error(f"数据集不包含必需的'text'列！实际列名: {dataset.column_names}") + logger.info("请确保语料文件包含文本内容在'text'字段中") + exit(1) + +def preprocess_function(examples): + return tokenizer( + examples['text'], + truncation=True, + max_length=1024, + padding='max_length', + return_tensors='pt' + ) + +tokenized_datasets = dataset.map( + preprocess_function, + batched=True, + remove_columns=dataset.column_names, + num_proc=16, + desc="Tokenizing dataset" +) + +logger.info(f"保存tokenized数据集到: {TOKENS_DATA_DIR}") +DatasetDict({"train": tokenized_datasets}).save_to_disk(TOKENS_DATA_DIR) +logger.info(f"Tokenized数据已保存，大小: {len(tokenized_datasets)} 样本") \ No newline at end of file diff --git a/src/step3_train_llm.py b/src/step3_train_llm.py index 2bb66c9..5781fb3 100755 --- a/src/step3_train_llm.py +++ b/src/step3_train_llm.py @@ -12,56 +12,38 @@ from accelerate import Accelerator import os import logging import glob - +from datasets import load_dataset, DatasetDict logging.basicConfig(level=logging.INFO) logger = logging.getLogger(__name__) accelerator = Accelerator() set_seed(42) - +TOKENS_DATA_DIR = "../corpus" +tokenizer_path="../weights/weights_tokenizer" model_name = "../weights/weights_llm" -tokenizer_path="" resume_option = None first_time=True # resume_option = True # resume_option = "../checkpoints_ds/checkpoint-500" -if first_time: - model = DollyForCausalLM(DollyConfig()) -tokenizer = AutoTokenizer.from_pretrained(model_name) + +tokenizer = AutoTokenizer.from_pretrained(tokenizer_path) if tokenizer.pad_token is None: tokenizer.pad_token = tokenizer.eos_token - -if isinstance(resume_option, str) and os.path.exists(resume_option): - logger.info(f"从指定检查点 {resume_option} 加载模型...") - model = DollyForCausalLM.from_pretrained(resume_option) + +if first_time: + model = DollyForCausalLM(DollyConfig()) else: - logger.info("未找到检查点，从预训练模型加载...") - logger.info(f"从预训练模型 {model_name} 加载...") - model = DollyForCausalLM.from_pretrained(model_name) + if isinstance(resume_option, str) and os.path.exists(resume_option): + logger.info(f"从指定检查点 {resume_option} 加载模型...") + model = DollyForCausalLM.from_pretrained(resume_option) + else: + logger.info("未找到检查点，从预训练模型加载...") + logger.info(f"从预训练模型 {model_name} 加载...") + model = DollyForCausalLM.from_pretrained(model_name) + model.gradient_checkpointing_enable() -data_files = glob.glob("../corpus/*.jsonl") -logger.info(f"找到训练语料: {data_files}") - -dataset = load_dataset('json', data_files=data_files, split='train') -# dataset = load_dataset('text', data_files={'train': '../corpus/wikipedia.txt'}) - -def preprocess_function(examples): - return tokenizer( - examples['text'], - truncation=True, - max_length=512, - padding='max_length', - return_tensors='pt' - ) - -tokenized_datasets = dataset.map( - preprocess_function, - batched=True, - # remove_columns=['text'], - remove_columns=dataset.column_names, - num_proc=16 -) +tokenized_datasets = DatasetDict.load_from_disk(TOKENS_DATA_DIR)["train"] data_collator = DataCollatorForLanguageModeling( tokenizer=tokenizer, @@ -75,14 +57,14 @@ training_args = TrainingArguments( save_strategy="steps", save_steps=500, learning_rate=1e-4, - per_device_train_batch_size=36, + per_device_train_batch_size=16, num_train_epochs=2, logging_dir="../logs", logging_steps=20, save_total_limit=2, fp16=True, - dataloader_num_workers=3, - gradient_accumulation_steps=2, + dataloader_num_workers=4, + gradient_accumulation_steps=4, report_to="tensorboard", optim="adamw_torch", lr_scheduler_type="cosine",