admin/LLM-from-scratch
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

[2025.06.19] update train code, pretoken corpus

Browse Source
This commit is contained in:
root
2025-06-19 17:00:27 +08:00
parent f28b310bc3
commit bd75a6d3ac
6 changed files with 100 additions and 158 deletions
Download Patch File Download Diff File Expand all files Collapse all files
src/configuration_dolly.py
+5 -5
View File
@@ -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
src/step0_train_tokenizer_jsonl.py
+3 -3
View File
@@ -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
+3 -3
View File
@@ -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
+8 -9
View File
@@ -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']:,}")
src/step2_get_tokens_from_parquet.py
+61 -100
View File
@@ -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("所有文件处理完成!")
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)} 样本")
src/step3_train_llm.py
+20 -38
View File
@@ -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",