444 lines
14 KiB
Python
444 lines
14 KiB
Python
#!/usr/bin/env python3
|
||
"""
|
||
缺失字符补充工具
|
||
|
||
步骤 1: find-missing — 扫描已预处理数据，找出从未出现的 label ID，输出 JSON
|
||
步骤 2: generate-template — 根据 JSON 生成 JSONL 占位文件，供用户手动填入包含缺失字的真实文本
|
||
步骤 3: preprocess-supplement — 将填好的 JSONL 文本预处理为 .npz 分片，输出到独立目录
|
||
|
||
用法：
|
||
python -m model.supplement_missing find-missing \
|
||
--preprocessed-dir ./preprocessed/train \
|
||
--output missing_chars.json
|
||
|
||
python -m model.supplement_missing generate-template \
|
||
--missing-chars missing_chars.json \
|
||
--output supplement_texts.jsonl
|
||
|
||
python -m model.supplement_missing preprocess-supplement \
|
||
--missing-chars missing_chars.json \
|
||
--supplement-texts supplement_texts.jsonl \
|
||
--output-dir ./preprocessed/supplement \
|
||
--num-samples 100000
|
||
"""
|
||
|
||
import argparse
|
||
import json
|
||
from pathlib import Path
|
||
from typing import Set
|
||
|
||
import numpy as np
|
||
import torch
|
||
from loguru import logger
|
||
from rich.console import Console
|
||
from rich.table import Table
|
||
from torch.utils.data import DataLoader
|
||
from tqdm import tqdm
|
||
|
||
from .dataset import PinyinInputDataset
|
||
from .preprocess import collect_samples
|
||
from .query import QueryEngine
|
||
from .trainer import preprocess_collate_fn, worker_init_fn
|
||
|
||
|
||
def scan_labels(preprocessed_dir: Path) -> Set[int]:
|
||
"""扫描预处理目录中所有 .npz 分片，收集所有出现过的 label ID"""
|
||
appeared: Set[int] = set()
|
||
|
||
shard_files = sorted(preprocessed_dir.glob("shard_*.npz"))
|
||
if not shard_files:
|
||
logger.warning(f"未找到 .npz 分片文件: {preprocessed_dir}")
|
||
return appeared
|
||
|
||
for shard_path in tqdm(shard_files, desc="扫描分片", unit="shard"):
|
||
data = np.load(shard_path)
|
||
labels = data["labels"].astype(np.int64)
|
||
if labels.ndim > 1 and labels.shape[-1] == 1:
|
||
labels = labels.squeeze(-1)
|
||
unique_ids = np.unique(labels)
|
||
appeared.update(int(uid) for uid in unique_ids)
|
||
del data
|
||
|
||
return appeared
|
||
|
||
|
||
def cmd_find_missing(args):
|
||
console = Console()
|
||
preprocessed_dir = Path(args.preprocessed_dir)
|
||
|
||
if not preprocessed_dir.exists():
|
||
console.print(f"[bold red]目录不存在: {preprocessed_dir}[/bold red]")
|
||
return
|
||
|
||
metadata_path = preprocessed_dir / "metadata.json"
|
||
if not metadata_path.exists():
|
||
console.print(f"[bold red]未找到 metadata.json: {metadata_path}[/bold red]")
|
||
return
|
||
|
||
with open(metadata_path, "r", encoding="utf-8") as f:
|
||
metadata = json.load(f)
|
||
console.print(
|
||
f"[bold cyan]预处理数据: {metadata['num_samples']:,} 样本, {metadata['num_shards']} 分片[/bold cyan]"
|
||
)
|
||
|
||
console.print("[bold cyan]扫描 labels...[/bold cyan]")
|
||
appeared = scan_labels(preprocessed_dir)
|
||
|
||
console.print("[bold cyan]加载 QueryEngine...[/bold cyan]")
|
||
query_engine = QueryEngine()
|
||
query_engine.load()
|
||
|
||
all_ids = set(query_engine._id_to_info.keys())
|
||
missing_ids = all_ids - appeared
|
||
|
||
missing_chars = []
|
||
for mid in sorted(missing_ids):
|
||
if mid == 0:
|
||
continue
|
||
info = query_engine.query_by_id(mid)
|
||
if info is not None:
|
||
missing_chars.append(
|
||
{
|
||
"id": info.id,
|
||
"char": info.char,
|
||
"pinyin": info.pinyin,
|
||
"count": info.count,
|
||
}
|
||
)
|
||
|
||
result = {
|
||
"missing_count": len(missing_chars),
|
||
"missing_chars": missing_chars,
|
||
}
|
||
|
||
output_path = Path(args.output)
|
||
output_path.parent.mkdir(parents=True, exist_ok=True)
|
||
with open(output_path, "w", encoding="utf-8") as f:
|
||
json.dump(result, f, ensure_ascii=False, indent=2)
|
||
|
||
console.print(f"\n[bold green]=== 扫描完成 ===[/bold green]")
|
||
console.print(f"词表大小: {len(all_ids):,} (含 EOS)")
|
||
console.print(f"已出现标签: {len(appeared):,}")
|
||
console.print(
|
||
f"[bold red]缺失标签: {len(missing_ids):,}[/bold red] (其中非 EOS: {len(missing_chars)})"
|
||
)
|
||
|
||
if missing_chars:
|
||
table = Table(
|
||
title=f"缺失字符 (共 {len(missing_chars)} 个)",
|
||
show_header=True,
|
||
header_style="bold magenta",
|
||
)
|
||
table.add_column("ID", style="cyan", width=8)
|
||
table.add_column("字符", style="yellow", width=6)
|
||
table.add_column("拼音", style="green", width=12)
|
||
table.add_column("语料频次", style="red", width=12)
|
||
for entry in missing_chars:
|
||
table.add_row(
|
||
str(entry["id"]),
|
||
entry["char"],
|
||
entry["pinyin"],
|
||
f"{entry['count']:,}",
|
||
)
|
||
console.print(table)
|
||
|
||
console.print(f"\n已输出到: {output_path}")
|
||
|
||
|
||
def cmd_generate_template(args):
|
||
console = Console()
|
||
|
||
missing_path = Path(args.missing_chars)
|
||
if not missing_path.exists():
|
||
console.print(f"[bold red]文件不存在: {missing_path}[/bold red]")
|
||
return
|
||
|
||
with open(missing_path, "r", encoding="utf-8") as f:
|
||
data = json.load(f)
|
||
|
||
missing_chars = data.get("missing_chars", [])
|
||
if not missing_chars:
|
||
console.print("[bold green]没有缺失字符，无需生成模板[/bold green]")
|
||
return
|
||
|
||
num_entries = args.num_entries
|
||
total_lines = len(missing_chars) * num_entries
|
||
|
||
console.print(f"[bold cyan]缺失字符数: {len(missing_chars)}[/bold cyan]")
|
||
console.print(f"[bold cyan]每字符模板数: {num_entries}[/bold cyan]")
|
||
console.print(f"[bold cyan]总模板行数: {total_lines}[/bold cyan]")
|
||
|
||
output_path = Path(args.output)
|
||
output_path.parent.mkdir(parents=True, exist_ok=True)
|
||
|
||
with open(output_path, "w", encoding="utf-8") as f:
|
||
for entry in missing_chars:
|
||
for i in range(num_entries):
|
||
line = json.dumps(
|
||
{"text": f"请在这里输入包含「{entry['char']}」字的第{i + 1}条文本"},
|
||
ensure_ascii=False,
|
||
)
|
||
f.write(line + "\n")
|
||
|
||
console.print(f"[bold green]模板已生成: {output_path}[/bold green]")
|
||
console.print(
|
||
f"共 {total_lines} 条（{len(missing_chars)} 字符 × {num_entries} 条/字符），"
|
||
f"请手动编辑该文件，将占位文本替换为包含对应字符的真实文本。"
|
||
)
|
||
|
||
|
||
def cmd_preprocess_supplement(args):
|
||
console = Console()
|
||
|
||
# 加载缺失字符
|
||
missing_path = Path(args.missing_chars)
|
||
if not missing_path.exists():
|
||
console.print(f"[bold red]文件不存在: {missing_path}[/bold red]")
|
||
return
|
||
|
||
with open(missing_path, "r", encoding="utf-8") as f:
|
||
data = json.load(f)
|
||
|
||
missing_chars = data.get("missing_chars", [])
|
||
if not missing_chars:
|
||
console.print("[bold green]没有缺失字符，无需处理[/bold green]")
|
||
return
|
||
|
||
target_labels = {entry["id"] for entry in missing_chars}
|
||
target_labels.add(0) # 包含 EOS
|
||
|
||
# 解析参数
|
||
py_style_weight = tuple(int(x) for x in args.py_style_weight.split(","))
|
||
length_weights = {
|
||
int(k): int(v)
|
||
for k, v in (item.split(":") for item in args.length_weights.split(","))
|
||
}
|
||
|
||
output_dir = Path(args.output_dir)
|
||
output_dir.mkdir(parents=True, exist_ok=True)
|
||
|
||
max_iter = args.num_samples * 5
|
||
num_workers = args.num_workers
|
||
|
||
console.print("[bold cyan]=== 补充数据预处理 ===[/bold cyan]")
|
||
console.print(f"补充文本: {args.supplement_texts}")
|
||
console.print(f"缺失字符数: {len(missing_chars)}")
|
||
console.print(f"目标样本: {args.num_samples:,}")
|
||
console.print(f"输出目录: {output_dir}")
|
||
console.print(f"Worker 数: {num_workers}")
|
||
console.print()
|
||
|
||
torch.manual_seed(args.seed)
|
||
np.random.seed(args.seed)
|
||
|
||
console.print("[bold cyan]创建补充数据集...[/bold cyan]")
|
||
dataset = PinyinInputDataset(
|
||
data_path="json",
|
||
max_workers=num_workers,
|
||
max_iter_length=max_iter,
|
||
max_seq_length=args.max_seq_length,
|
||
text_field="text",
|
||
py_style_weight=py_style_weight,
|
||
shuffle_buffer_size=100,
|
||
length_weights=length_weights,
|
||
data_kwargs={
|
||
"data_files": args.supplement_texts,
|
||
"streaming": False,
|
||
},
|
||
target_labels=target_labels,
|
||
)
|
||
|
||
dataloader_kwargs = {
|
||
"batch_size": args.batch_size,
|
||
"num_workers": num_workers,
|
||
"pin_memory": False,
|
||
"worker_init_fn": worker_init_fn,
|
||
"collate_fn": preprocess_collate_fn(args.max_seq_length),
|
||
}
|
||
if num_workers > 0:
|
||
dataloader_kwargs["prefetch_factor"] = 2
|
||
dataloader_kwargs["persistent_workers"] = True
|
||
|
||
dataloader = DataLoader(dataset, **dataloader_kwargs)
|
||
|
||
logger.info("开始收集补充数据...")
|
||
count = collect_samples(
|
||
dataloader,
|
||
args.num_samples,
|
||
output_dir,
|
||
"supplement",
|
||
args.max_seq_length,
|
||
args.shard_size,
|
||
)
|
||
|
||
if count < args.num_samples:
|
||
logger.warning(f"补充样本不足: 目标 {args.num_samples}, 实际 {count}")
|
||
|
||
console.print("\n[bold green]=== 补充预处理完成 ===[/bold green]")
|
||
console.print(f"生成样本: {count:,}")
|
||
console.print(f"输出目录: {output_dir}")
|
||
|
||
total_size = sum(
|
||
f.stat().st_size for f in output_dir.iterdir() if f.suffix == ".npz"
|
||
)
|
||
console.print(f"总大小: {total_size / (1024**3):.2f} GB (compressed)")
|
||
console.print()
|
||
console.print(
|
||
"[bold yellow]提示[/bold yellow]: 请检查补充数据质量，清洗无误后手动将 shard_*.npz 合并到 train/ 目录并更新 metadata.json"
|
||
)
|
||
|
||
|
||
def main():
|
||
parser = argparse.ArgumentParser(
|
||
description="缺失字符补充工具",
|
||
formatter_class=argparse.RawDescriptionHelpFormatter,
|
||
epilog="""
|
||
子命令：
|
||
find-missing 扫描已预处理数据，找出从未出现的 label ID
|
||
generate-template 根据缺失字符 JSON 生成 JSONL 占位文件
|
||
preprocess-supplement 将填好的 JSONL 预处理为 .npz 分片（独立目录）
|
||
|
||
示例：
|
||
python -m model.supplement_missing find-missing \\
|
||
--preprocessed-dir ./preprocessed/train \\
|
||
--output missing_chars.json
|
||
|
||
python -m model.supplement_missing generate-template \\
|
||
--missing-chars missing_chars.json \\
|
||
--output supplement_texts.jsonl
|
||
|
||
python -m model.supplement_missing preprocess-supplement \\
|
||
--missing-chars missing_chars.json \\
|
||
--supplement-texts supplement_texts.jsonl \\
|
||
--output-dir ./preprocessed/supplement \\
|
||
--num-samples 100000
|
||
""",
|
||
)
|
||
subparsers = parser.add_subparsers(dest="command", help="子命令")
|
||
|
||
# find-missing
|
||
p_find = subparsers.add_parser("find-missing", help="扫描预处理数据，找出缺失标签")
|
||
p_find.add_argument(
|
||
"--preprocessed-dir",
|
||
type=str,
|
||
required=True,
|
||
help="预处理数据目录（包含 shard_*.npz 和 metadata.json）",
|
||
)
|
||
p_find.add_argument(
|
||
"--output",
|
||
type=str,
|
||
default="missing_chars.json",
|
||
help="输出 JSON 文件路径（默认: missing_chars.json）",
|
||
)
|
||
|
||
# generate-template
|
||
p_gen = subparsers.add_parser("generate-template", help="生成补充文本模板")
|
||
p_gen.add_argument(
|
||
"--missing-chars",
|
||
type=str,
|
||
required=True,
|
||
help="缺失字符 JSON 文件路径（由 find-missing 生成）",
|
||
)
|
||
p_gen.add_argument(
|
||
"--output",
|
||
type=str,
|
||
default="supplement_texts.jsonl",
|
||
help="输出 JSONL 文件路径（默认: supplement_texts.jsonl）",
|
||
)
|
||
p_gen.add_argument(
|
||
"--num-entries",
|
||
type=int,
|
||
default=3,
|
||
help="每个缺失字符生成的模板条数（默认: 3）",
|
||
)
|
||
|
||
# preprocess-supplement
|
||
p_pre = subparsers.add_parser(
|
||
"preprocess-supplement", help="将 JSONL 预处理为 .npz 分片"
|
||
)
|
||
p_pre.add_argument(
|
||
"--missing-chars",
|
||
type=str,
|
||
required=True,
|
||
help="缺失字符 JSON 文件路径（由 find-missing 生成）",
|
||
)
|
||
p_pre.add_argument(
|
||
"--supplement-texts",
|
||
type=str,
|
||
required=True,
|
||
help="已填写的补充文本 JSONL 文件路径",
|
||
)
|
||
p_pre.add_argument(
|
||
"--output-dir",
|
||
type=str,
|
||
required=True,
|
||
help="输出目录（独立目录，不会覆盖已有数据）",
|
||
)
|
||
p_pre.add_argument(
|
||
"--num-samples",
|
||
type=int,
|
||
required=True,
|
||
help="目标样本数量",
|
||
)
|
||
p_pre.add_argument(
|
||
"--batch-size",
|
||
type=int,
|
||
default=128,
|
||
help="批大小（默认: 128）",
|
||
)
|
||
p_pre.add_argument(
|
||
"--num-workers",
|
||
type=int,
|
||
default=0,
|
||
help="DataLoader worker 数量。本地 JSONL 小文件建议 0（默认: 0）",
|
||
)
|
||
p_pre.add_argument(
|
||
"--max-seq-length",
|
||
type=int,
|
||
default=128,
|
||
help="最大序列长度（默认: 128）",
|
||
)
|
||
p_pre.add_argument(
|
||
"--seed",
|
||
type=int,
|
||
default=42,
|
||
help="随机种子（默认: 42）",
|
||
)
|
||
p_pre.add_argument(
|
||
"--shard-size",
|
||
type=int,
|
||
default=5_000_000,
|
||
help="分片大小（样本数），控制内存峰值（默认: 500万）",
|
||
)
|
||
p_pre.add_argument(
|
||
"--py-style-weight",
|
||
type=str,
|
||
default="9,2,1",
|
||
help="拼音风格权重（逗号分隔，默认: 9,2,1）",
|
||
)
|
||
p_pre.add_argument(
|
||
"--length-weights",
|
||
type=str,
|
||
default="1:10,2:50,3:50,4:40,5:15,6:10,7:5,8:2",
|
||
help="词长权重（默认: 1:10,2:50,3:50,4:40,5:15,6:10,7:5,8:2）",
|
||
)
|
||
|
||
args = parser.parse_args()
|
||
|
||
if args.command is None:
|
||
parser.print_help()
|
||
return
|
||
|
||
if args.command == "find-missing":
|
||
cmd_find_missing(args)
|
||
elif args.command == "generate-template":
|
||
cmd_generate_template(args)
|
||
elif args.command == "preprocess-supplement":
|
||
cmd_preprocess_supplement(args)
|
||
|
||
|
||
app = main
|
||
|
||
if __name__ == "__main__":
|
||
main()
|