README.md · xingyu1996/tiger-gpt2 at main

metadata

license: apache-2.0
datasets:
  - TigerResearch/pretrain_zh
language:
  - zh
tags:
  - gpt2
pipeline_tag: text-generation

Model Name/ID：tiger-gpt2

模型描述

这是一个基于 GPT-2 架构，在中文数据上预训练的语言模型。

如何使用

from transformers import AutoModelForCausalLM, AutoTokenizer

model_name = "xingyu1996/tiger-gpt2" # 例如 "your_username/my-awesome-gpt2-minimind"
tokenizer = AutoTokenizer.from_pretrained(model_name)
model = AutoModelForCausalLM.from_pretrained(model_name)

# 这里可以加一个简单的生成示例
input_text = "今天天气"
input_ids = tokenizer.encode(input_text, return_tensors='pt')
output = model.generate(input_ids, max_length=50)
print(tokenizer.decode(output[0], skip_special_tokens=True))

训练数据

原始数据集：https://huggingface.co/datasets/TigerResearch/pretrain_zh/viewer/default/train?p=169048

原始数据集共有16,905,023条记录，预算较低，因此抽取其中5%数据集。

分词器采用gpt2的bpe分词器，词表大小50257。一共 1,987,309,561 tokens。

抽取代码：

import datasets
import torch
from tqdm import tqdm
from transformers import AutoTokenizer
import os

# 设置固定种子确保可复现
SEED = 42
SAMPLE_RATIO = 0.05  # 抽取5%的数据

# 创建输出目录
os.makedirs("tiger_data", exist_ok=True)

print("加载数据集...")
ds = datasets.load_dataset("TigerResearch/pretrain_zh")
train_ds = ds["train"]


train_ds = train_ds.filter(lambda x: len(x["content"]) >= 50)
print(f"过滤后剩余: {len(train_ds)}条 (过滤掉 {len(ds['train'])-len(train_ds)} 条)")

# 计算目标抽样大小
total_size = len(train_ds)
target_size = int(total_size * SAMPLE_RATIO)
print(f"数据集总大小: {total_size}条")
print(f"目标抽样大小: {target_size}条 ({SAMPLE_RATIO*100:.1f}%)")

# 随机抽样
print("执行随机抽样...")
sampled_ds = train_ds.shuffle(seed=SEED).select(range(target_size))

# 打印抽样结果的统计信息
data_types = sampled_ds["dataType"]
from collections import Counter
type_counts = Counter(data_types)
print("\n抽样后dataType分布:")
for dtype, count in type_counts.most_common():
    print(f"{dtype}: {count}条 ({count/target_size*100:.2f}%)")

# 计算内容长度统计
print("\n计算内容长度统计...")
content_lengths = []
for i in tqdm(range(min(1000, len(sampled_ds)))):  # 抽样1000条计算
    content_lengths.append(len(sampled_ds[i]["content"]))

avg_length = sum(content_lengths) / len(content_lengths)
print(f"内容平均长度: {avg_length:.1f}字符")
print(f"内容长度范围: {min(content_lengths)} - {max(content_lengths)}字符")

# 打乱数据集
sampled_ds = sampled_ds.shuffle(seed=42)

# 保存抽样数据集(可选)
print("\n保存抽样数据集...")
sampled_ds.save_to_disk("tiger_data/sampled_dataset")
print(f"已保存到 tiger_data/sampled_dataset")

# 加载tokenizer
print("\n加载tokenizer...")
tokenizer = AutoTokenizer.from_pretrained("gpt2")

# 定义一个处理函数
def tokenize_function(examples):
    # examples 是一个包含多条数据的字典，例如 {'content': ['文本1', '文本2', ...]}
    # 我们需要对 'content' 字段中的每个文本进行处理
    # 添加特殊标记并编码
    # 注意：tokenizer可以直接处理文本列表
    texts = [f"{content}{tokenizer.eos_token}" for content in examples["content"]]
    # return tokenizer(texts) # 直接返回编码结果，datasets会自动处理
    # 为了得到一个扁平的token列表，我们稍微复杂一点处理：
    all_tokens = []
    for text in texts:
        tokens = tokenizer.encode(text)
        all_tokens.extend(tokens)
    # .map期望返回一个字典，key是新的列名或要修改的列名
    # 但这里我们想汇总所有tokens，所以先用map做初步编码，再汇总
    return tokenizer(texts, truncation=False, add_special_tokens=False) # 返回编码结果，包含 input_ids

# 使用 .map() 进行并行处理
# TODO1: 了解 .map() 的参数，特别是 batched 和 num_proc
# batched=True 让 tokenize_function 一次接收多条数据，效率更高
# num_proc 设置并行处理的核心数，可以设置为 os.cpu_count()
print("开始并行tokenize处理...")
tokenized_datasets = sampled_ds.map(
    tokenize_function,
    batched=True,
    num_proc=4, 
    remove_columns=sampled_ds.column_names 
)

# TODO3: 从 tokenized_datasets 中提取所有 token 并合并成一个列表
#       (提示: tokenized_datasets['input_ids'] 可能是一个嵌套列表，需要展平)
print("合并所有tokens...")
all_tokens = []
for tokens_list in tqdm(tokenized_datasets['input_ids']):
     all_tokens.extend(tokens_list)


# 转换为tensor并保存
print("转换为Tensor并保存...")
all_tokens_tensor = torch.tensor(all_tokens, dtype=torch.long)
# ... 后续保存和统计代码不变 ...
torch.save(all_tokens_tensor, "tiger_data/tokens.pt")

# 打印token统计信息
token_count = len(all_tokens_tensor)
print(f"\n总token数: {token_count}")
# 注意：平均token数计算方式可能需要调整，因为上面合并了所有tokens
# print(f"每条数据平均token数: {token_count/len(sampled_ds):.1f}") # 这个计算不再准确

# 估算训练时间(基于之前427000 tokens/分钟的速度)
est_hours = token_count / (427000 * 60)
print(f"预估训练时间: {est_hours:.2f}小时 ({est_hours*60:.1f}分钟)")

转换后的tokens.pt，保存在modelscope，方便在不同空间迁移：https://modelscope.cn/datasets/xiaoyuer2025/tiger_pretrain_tokens.pt

训练过程

因为是学习过程中制作的大模型，中间学习率策略换了好几次。主要用了余弦退火策略和三步下降策略。也就是验证集的loss在三次测试都没有下降的话，学习率减半。

但是这个策略太过激进，建议后续使用https://www.alphaxiv.org/abs/2408.13359该论文的学习率控制方式。

评估结果

学习率最低到了1e-6，大约30000步左右，val_loss不再下降（大约1.5左右）。因此采取早停法。