Views
No views yet
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))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}分钟)")