# 预训练数据收集脚本 import os import json import pandas as pd from datasets import load_dataset, concatenate_datasets from typing import List, Dict import re from tqdm import tqdm class PretrainingDataBuilder: def __init__(self, output_dir: str = "./pretraining_data"): self.output_dir = output_dir os.makedirs(output_dir, exist_ok=True) def collect_datasets(self): """收集多个数据源""" print("📚 收集预训练数据源...") datasets = {} # 1. WikiText-103：高质量的维基百科文本 print(" 下载 WikiText-103...") try: wikitext = load_dataset("wikitext", "wikitext-103-raw-v1") datasets["wikitext"] = wikitext print(f" ✅ WikiText-103: {len(wikitext['train'])} 条") except Exception as e: print(f" ❌ WikiText-103 下载失败: {e}") # 2. BookCorpus：书籍文本（使用子集） print(" 下载 BookCorpus 子集...") try: bookcorpus = load_dataset("bookcorpus", split="train[:10000]") # 取前1万条 datasets["bookcorpus"] = {"train": bookcorpus} print(f" ✅ BookCorpus: {len(bookcorpus)} 条") except Exception as e: print(f" ❌ BookCorpus 下载失败: {e}") # 3. OpenWebText：网页文本子集 print(" 下载 OpenWebText 子集...") try: openwebtext = load_dataset("openwebtext", split="train[:5000]") # 取前5千条 datasets["openwebtext"] = {"train": openwebtext} print(f" ✅ OpenWebText: {len(openwebtext)} 条") except Exception as e: print(f" ❌ OpenWebText 下载失败: {e}") # 4. 扩展SST-2相关的电影评论数据（无标签版本） print(" 收集电影评论文本...") try: # IMDB数据集的无监督部分 imdb = load_dataset("imdb", split="unsupervised[:10000]") datasets["imdb_unsup"] = {"train": imdb} print(f" ✅ IMDB无监督: {len(imdb)} 条") except Exception as e: print(f" ❌ IMDB数据下载失败: {e}") return datasets def clean_text(self, text: str) -> str: """清洗文本数据""" if not text or not isinstance(text, str): return "" # 移除多余的空白字符 text = re.sub(r'\s+', ' ', text) # 移除特殊字符（保留基本标点） text = re.sub(r'[^\w\s\.\,\!\?\;\:\-\(\)\"\']+', '', text) # 移除过短的文本 if len(text.split()) < 10: return "" # 移除过长的文本（可能是噪声） if len(text.split()) > 1000: text = ' '.join(text.split()[:1000]) return text.strip() def process_datasets(self, datasets: Dict): """处理和合并数据集""" print("\n🔧 处理和清洗数据...") all_texts = [] for dataset_name, dataset in datasets.items(): print(f" 处理 {dataset_name}...") if dataset_name == "wikitext": # WikiText格式处理 for split in ["train", "validation", "test"]: if split in dataset: for item in tqdm(dataset[split], desc=f"处理{split}"): text = self.clean_text(item["text"]) if text: all_texts.append(text) elif dataset_name in ["bookcorpus", "openwebtext"]: # 书籍和网页文本处理 for item in tqdm(dataset["train"], desc=f"处理{dataset_name}"): text = self.clean_text(item["text"]) if text: all_texts.append(text) elif dataset_name == "imdb_unsup": # IMDB无监督数据处理 for item in tqdm(dataset["train"], desc="处理IMDB"): text = self.clean_text(item["text"]) if text: all_texts.append(text) print(f"\n✅ 数据处理完成") print(f" 总文本数: {len(all_texts):,}") print(f" 平均长度: {sum(len(t.split()) for t in all_texts) / len(all_texts):.1f} 词") return all_texts def create_pretraining_corpus(self, texts: List[str]): """创建预训练语料库""" print("\n📝 创建预训练语料库...") # 随机打乱 import random random.shuffle(texts) # 分割数据集 total_size = len(texts) train_size = int(0.9 * total_size) val_size = int(0.05 * total_size) train_texts = texts[:train_size] val_texts = texts[train_size:train_size + val_size] test_texts = texts[train_size + val_size:] # 保存为文本文件 splits = { "train": train_texts, "validation": val_texts, "test": test_texts } for split_name, split_texts in splits.items(): file_path = os.path.join(self.output_dir, f"{split_name}.txt") with open(file_path, "w", encoding="utf-8") as f: for text in split_texts: f.write(text + "\n") print(f" {split_name}: {len(split_texts):,} 条 -> {file_path}") # 生成统计报告 self.generate_corpus_report(splits) return splits def generate_corpus_report(self, splits: Dict[str, List[str]]): """生成语料库统计报告""" report = { "corpus_name": "Mini-BERT Pretraining Corpus", "creation_date": pd.Timestamp.now().isoformat(), "splits": {}, "vocabulary_stats": {}, "text_stats": {} } all_texts = [] for split_name, texts in splits.items(): all_texts.extend(texts) # 计算统计信息 word_counts = [len(text.split()) for text in texts] char_counts = [len(text) for text in texts] report["splits"][split_name] = { "num_documents": len(texts), "avg_words": sum(word_counts) / len(word_counts), "avg_chars": sum(char_counts) / len(char_counts), "total_words": sum(word_counts), "total_chars": sum(char_counts) } # 整体统计 all_words = [] for text in all_texts: all_words.extend(text.lower().split()) from collections import Counter word_freq = Counter(all_words) report["vocabulary_stats"] = { "total_tokens": len(all_words), "unique_tokens": len(word_freq), "vocab_size": len(word_freq), "most_common_words": word_freq.most_common(20) } # 保存报告 report_path = os.path.join(self.output_dir, "corpus_report.json") with open(report_path, "w", encoding="utf-8") as f: json.dump(report, f, ensure_ascii=False, indent=2) print(f"\n📊 语料库统计报告:") print(f" 总文档数: {len(all_texts):,}") print(f" 总词数: {report['vocabulary_stats']['total_tokens']:,}") print(f" 词汇表大小: {report['vocabulary_stats']['vocab_size']:,}") print(f" 报告已保存: {report_path}") def main(): """主函数""" print("🏗️ 构建预训练语料库") print("=" * 60) builder = PretrainingDataBuilder() # 收集数据集 datasets = builder.collect_datasets() if not datasets: print("❌ 没有成功下载任何数据集") return # 处理数据 texts = builder.process_datasets(datasets) if not texts: print("❌ 没有有效的文本数据") return # 创建语料库 corpus = builder.create_pretraining_corpus(texts) print("\n🎉 预训练语料库构建完成！") print("接下来可以用于BERT预训练...") if __name__ == "__main__": main()