Upload prepare_wiki_data.py with huggingface_hub

prepare_wiki_data.py

@@ -0,0 +1,197 @@
+#!/usr/bin/env python3
+"""下载维基百科中文数据 + 分词 + 生成 LAL 训练 .bin 文件.
+
+用法:
+  python3 scripts/prepare_wiki_data.py [--n_articles 10000] [--out data/wiki_bpe.bin]
+
+输出格式: LALT 二进制 (与 large_bpe_v3.bin 兼容)
+  [magic: 4 bytes = "LALT"]
+  [n_samples: int32]
+  [n_vocab: int32]
+  then n_samples records:
+    [n_tokens: int32]
+    [token_ids: int32 * n_tokens]
+"""
+import os
+import sys
+import struct
+import subprocess
+import argparse
+import time
+
+# === 固化配置 ===
+TOKENIZER_MODEL = "tokenizer/chinese_bpe.model"
+WIKI_DUMP_URL = "https://dumps.wikimedia.org/zhwiki/latest/zhwiki-latest-pages-articles-multistream.xml.bz2"
+DEFAULT_OUT = "data/wiki_bpe.bin"
+DEFAULT_N_ARTICLES = 10000  # 先取 1 万篇, 约 500-1000 万 token
+
+
+def log(msg):
+    print(f"[WIKI] {msg}", flush=True)
+
+
+def download_wiki_dump(local_path, max_articles=None):
+    """下载维基百科 dump (流式解压 + 解析, 避免下载整个 2GB+ bz2)."""
+    import bz2
+    import xml.etree.ElementTree as ET
+    import urllib.request
+    
+    log(f"下载+解析维基百科 (最多 {max_articles or '全部'} 篇)...")
+    
+    articles = []
+    in_text = False
+    in_title = False
+    current_title = ""
+    current_text = ""
+    article_count = 0
+    
+    # 流式下载 + bz2 解压
+    req = urllib.request.Request(WIKI_DUMP_URL, headers={"User-Agent": "LAL-Data-Prep/1.0"})
+    with urllib.request.urlopen(req) as resp:
+        with bz2.open(resp, "rt", encoding="utf-8") as f:
+            for line in f:
+                if "<title>" in line:
+                    start = line.index("<title>") + 7
+                    end = line.index("</title>")
+                    current_title = line[start:end].strip()
+                elif "<text" in line:
+                    in_text = True
+                    # 提取 text 标签内的内容
+                    if ">" in line:
+                        start = line.index(">") + 1
+                        current_text = line[start:]
+                    else:
+                        current_text = ""
+                elif "</text>" in line:
+                    in_text = False
+                    end = line.index("</text>")
+                    current_text += line[:end]
+                    
+                    # 过滤: 跳过重定向、空页面
+                    if current_text and not current_text.startswith("#REDIRECT"):
+                        # 清理 wiki 标记 (简单版)
+                        text = clean_wiki_text(current_text)
+                        if len(text) > 100:  # 太短的文章跳过
+                            articles.append(text)
+                            article_count += 1
+                            if article_count % 1000 == 0:
+                                log(f"  已收集 {article_count} 篇文章")
+                    
+                    current_text = ""
+                    
+                    if max_articles and article_count >= max_articles:
+                        break
+            
+            # 处理最后一篇
+            if in_text and current_text:
+                text = clean_wiki_text(current_text)
+                if len(text) > 100:
+                    articles.append(text)
+                    article_count += 1
+    
+    log(f"共收集 {len(articles)} 篇文章")
+    return articles
+
+
+def clean_wiki_text(text):
+    """简单清理 wiki 标记."""
+    import re
+    # 去掉 wiki 模板 {{...}}
+    text = re.sub(r'\{\{[^}]*\}\}', '', text)
+    # 去掉 wiki 链接 [[...]]
+    text = re.sub(r'\[\[([^|\]]*\|)?([^\]]*)\]\]', r'\2', text)
+    # 去掉 HTML 标签
+    text = re.sub(r'<[^>]+>', '', text)
+    # 去掉 wiki 标题标记 ==
+    text = re.sub(r'^=+\s*([^=]+)\s*=+$', r'\1', text, flags=re.MULTILINE)
+    # 去掉引用
+    text = re.sub(r'<ref[^>]*>.*?</ref>', '', text, flags=re.DOTALL)
+    text = re.sub(r'<ref[^>]*/>', '', text)
+    # 去掉多余空行
+    text = re.sub(r'\n{3,}', '\n\n', text)
+    return text.strip()
+
+
+def tokenize_with_bpe(texts, tokenizer_model):
+    """用 sentencepiece BPE 分词."""
+    try:
+        import sentencepiece as spm
+    except ImportError:
+        log("安装 sentencepiece...")
+        subprocess.run([sys.executable, "-m", "pip", "install", "-q", "sentencepiece"], check=True)
+        import sentencepiece as spm
+    
+    sp = spm.SentenceProcessor()
+    sp.Load(tokenizer_model)
+    
+    all_samples = []
+    total_tokens = 0
+    
+    for i, text in enumerate(texts):
+        # 分词 (每篇文章作为一个 sample)
+        tokens = sp.EncodeAsIds(text)
+        if len(tokens) > 10:  # 太短的跳过
+            all_samples.append(tokens)
+            total_tokens += len(tokens)
+        
+        if (i + 1) % 1000 == 0:
+            log(f"  分词 {i+1}/{len(texts)} 篇, 总 token {total_tokens}")
+    
+    log(f"分词完成: {len(all_samples)} samples, {total_tokens} tokens ({total_tokens/10000:.1f}万)")
+    return all_samples
+
+
+def write_lalt_bin(samples, out_path, n_vocab=32768):
+    """写 LALT 二进制格式."""
+    log(f"写入 {out_path}...")
+    with open(out_path, "wb") as f:
+        # header
+        f.write(b"LALT")
+        f.write(struct.pack("<i", len(samples)))
+        f.write(struct.pack("<i", n_vocab))
+        
+        # samples
+        for tokens in samples:
+            f.write(struct.pack("<i", len(tokens)))
+            for tok in tokens:
+                f.write(struct.pack("<i", tok))
+    
+    size = os.path.getsize(out_path)
+    log(f"完成: {out_path} ({size / 1024 / 1024:.1f} MB)")
+
+
+def main():
+    parser = argparse.ArgumentParser(description="准备维基百科训练数据")
+    parser.add_argument("--n_articles", type=int, default=DEFAULT_N_ARTICLES,
+                        help=f"文章数 (默认 {DEFAULT_N_ARTICLES})")
+    parser.add_argument("--out", type=str, default=DEFAULT_OUT,
+                        help=f"输出路径 (默认 {DEFAULT_OUT})")
+    args = parser.parse_args()
+    
+    # 检查 tokenizer
+    if not os.path.exists(TOKENIZER_MODEL):
+        log(f"[!] tokenizer 不存在: {TOKENIZER_MODEL}")
+        sys.exit(1)
+    
+    # 确保输出目录存在
+    os.makedirs(os.path.dirname(args.out), exist_ok=True)
+    
+    start = time.time()
+    
+    # 1. 下载 + 解析维基百科
+    articles = download_wiki_dump(args.out + ".tmp", max_articles=args.n_articles)
+    
+    # 2. 分词
+    samples = tokenize_with_bpe(articles, TOKENIZER_MODEL)
+    
+    # 3. 写 .bin
+    write_lalt_bin(samples, args.out)
+    
+    elapsed = time.time() - start
+    log(f"总计耗时 {elapsed:.0f}s")
+    log(f"数据文件: {args.out}")
+    log(f"替换训练数据: cp {args.out} data/large_bpe_v3.bin")
+
+
+if __name__ == "__main__":
+    main()