[2025.06.19] update train code, pretoken corpus

src/configuration_dolly.py

@@ -28,9 +28,9 @@ class DollyConfig(PretrainedConfig):
 
     def __init__(
         self,
-        vocab_size=32000,
+        vocab_size=151936,
-        hidden_size=2048,
+        hidden_size=1024,
-        intermediate_size=6144,
+        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
 

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,

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,

src/step1.1_count_tokens_length.py

@@ -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
 

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"
+from transformers import AutoTokenizer, set_seed
-OUTPUT_DIR = "../corpus/tokens"  
+from datasets import load_dataset, DatasetDict
-MAX_SEQ_LENGTH = 1024       
+from accelerate import Accelerator
-NUM_PROC = min(32, psutil.cpu_count(logical=False))  
+import os
-CHUNK_SIZE = 50000           
+import logging
-MAX_MEMORY_PERCENT = 70   
+import glob
-os.makedirs(OUTPUT_DIR, exist_ok=True)
+
-tokenizer = AutoTokenizer.from_pretrained("../weights/weights_tokenizer")
+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"))
+data_files = []
-print(f"找到 {len(file_paths)} 个parquet文件")
+supported_formats = ['.jsonl', '.parquet', '.txt', '.json', '.csv']
+for format in supported_formats:
+    data_files.extend(glob.glob(f"../train/*{format}"))
 
-def memory_safe():
+if not data_files:
-    mem = psutil.virtual_memory()
+    logger.error("未找到任何支持的语料文件！支持的格式: .jsonl, .parquet, .txt, .json, .csv")
-    return mem.percent < MAX_MEMORY_PERCENT
+    exit(1)
     
-def process_chunk(texts, tokenizer):
+logger.info(f"找到 {len(data_files)} 个训练语料文件: {data_files}")
-    tokenized = tokenizer(
+
-        texts,
+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,
+        max_length=1024,
-        padding="max_length",
+        padding='max_length',
-        return_tensors="np"
+        return_tensors='pt'
     )
-    return tokenized["input_ids"], tokenized["attention_mask"]
 
-def process_file(file_path):
+tokenized_datasets = dataset.map(
-    try:
+    preprocess_function,
-        base_name = os.path.basename(file_path)
+    batched=True,
-        output_path = os.path.join(OUTPUT_DIR, f"tokenized_{base_name}")
+    remove_columns=dataset.column_names,
-        if os.path.exists(output_path):
+    num_proc=16,
-            print(f"文件已存在: {output_path}, 跳过")
+    desc="Tokenizing dataset"
-            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}行/秒")
 
-    except Exception as e:
+logger.info(f"保存tokenized数据集到: {TOKENS_DATA_DIR}")
-        print(f"处理文件 {file_path} 时出错: {str(e)}")
+DatasetDict({"train": tokenized_datasets}).save_to_disk(TOKENS_DATA_DIR)
-        if 'dataset' in locals():
+logger.info(f"Tokenized数据已保存，大小: {len(tokenized_datasets)} 样本")
-            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("所有文件处理完成!")

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(tokenizer_path)
-tokenizer = AutoTokenizer.from_pretrained(model_name)
 if tokenizer.pad_token is None:
     tokenizer.pad_token = tokenizer.eos_token
     
-if isinstance(resume_option, str) and os.path.exists(resume_option):
+if first_time:
-    logger.info(f"从指定检查点 {resume_option} 加载模型...")
+    model = DollyForCausalLM(DollyConfig()) 
-    model = DollyForCausalLM.from_pretrained(resume_option)
 else:
-    logger.info("未找到检查点，从预训练模型加载...")
+    if isinstance(resume_option, str) and os.path.exists(resume_option):
-    logger.info(f"从预训练模型 {model_name} 加载...")
+        logger.info(f"从指定检查点 {resume_option} 加载模型...")
-    model = DollyForCausalLM.from_pretrained(model_name)
+        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")
+tokenized_datasets = DatasetDict.load_from_disk(TOKENS_DATA_DIR)["train"]
-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  
-)
 
 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,
+    dataloader_num_workers=4,
-    gradient_accumulation_steps=2,
+    gradient_accumulation_steps=4,
     report_to="tensorboard",
     optim="adamw_torch",
     lr_scheduler_type="cosine",