[2025.06.19] update train code, pretoken corpus

2025-06-19 17:00:27 +08:00
parent f28b310bc3
commit bd75a6d3ac
6 changed files with 100 additions and 158 deletions
@@ -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

@@ -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,
@@ -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,
@@ -7,7 +7,6 @@ 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

@@ -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
+if not data_files:
+    logger.error("未找到任何支持的语料文件！支持的格式: .jsonl, .parquet, .txt, .json, .csv")
+    exit(1)
    
-def process_chunk(texts, tokenizer):
-    tokenized = tokenizer(
-        texts,
+logger.info(f"找到 {len(data_files)} 个训练语料文件: {data_files}")
+
+file_ext = os.path.splitext(data_files[0])[1].lower()
+
+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=MAX_SEQ_LENGTH,
-        padding="max_length",
-        return_tensors="np"
+        max_length=1024,
+        padding='max_length',
+        return_tensors='pt'
    )
-    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}行/秒")
+tokenized_datasets = dataset.map(
+    preprocess_function,
+    batched=True,
+    remove_columns=dataset.column_names,
+    num_proc=16,
+    desc="Tokenizing dataset"
+)

-    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)
-    gc.collect()
-
-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("所有文件处理完成!")
+logger.info(f"保存tokenized数据集到: {TOKENS_DATA_DIR}")
+DatasetDict({"train": tokenized_datasets}).save_to_disk(TOKENS_DATA_DIR)
+logger.info(f"Tokenized数据已保存，大小: {len(tokenized_datasets)} 样本")
@@ -12,25 +12,27 @@ 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 first_time:
+    model = DollyForCausalLM(DollyConfig()) 
+else:
    if isinstance(resume_option, str) and os.path.exists(resume_option):
        logger.info(f"从指定检查点 {resume_option} 加载模型...")
        model = DollyForCausalLM.from_pretrained(resume_option)
@@ -38,30 +40,10 @@ 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",