feat(confusion-finetune): 引入拼音混淆样本生成与三路交叉注意力架构
This commit is contained in:
parent
88955bcfdd
commit
bfff409b35
|
|
@ -0,0 +1,289 @@
|
|||
# 混淆样本生成与候选词解码 (Confusion Sample Generation & Word Candidate Decoding)
|
||||
|
||||
> 分支: `feat/confusion-finetune`
|
||||
> 包管理: [uv](https://github.com/astral-sh/uv)
|
||||
|
||||
---
|
||||
|
||||
## 1. 背景与动机
|
||||
|
||||
### 1.1 原架构的问题
|
||||
|
||||
原始训练管线中 `part4`（候选词提示）位于 Encoder 输入：
|
||||
|
||||
```
|
||||
[CLS] part4 | part1 [SEP] part3 [SEP] → BERT Encoder → H, P
|
||||
```
|
||||
|
||||
- **训练-推理不一致**：训练时 30% 样本包含 `part4`，推理时完全不包含
|
||||
- **Encoder 不可复用**：`part4` 改变会污染 prefix 的上下文编码 H，ONNX 拆分后 encoder 无法独立于候选词运行
|
||||
|
||||
### 1.2 目标
|
||||
|
||||
- 将 `part4` 从 Encoder 移至 Decoder，实现三路 CrossAttention
|
||||
- 用"拼音相似词混淆"替代随机词，增强模型区分同音/近音字的能力
|
||||
- Encoder 序列长度从 128 缩减至 88（移除 part4 后实际最大需求为 67 tokens）
|
||||
|
||||
---
|
||||
|
||||
## 2. 架构变更
|
||||
|
||||
### 2.1 CrossAttention 从双路扩展为三路
|
||||
|
||||
```
|
||||
原架构（双路）:
|
||||
slots_S(Q) × context_H(K,V) → text_attn
|
||||
slots_S(Q) × pinyin_P(K,V) → pinyin_attn
|
||||
combined = text_attn + pinyin_attn
|
||||
|
||||
新架构（三路）:
|
||||
slots_S(Q) × context_H(K,V) → text_attn
|
||||
slots_S(Q) × pinyin_P(K,V) → pinyin_attn
|
||||
slots_S(Q) × part4_E(K,V) → part4_attn ← 新增
|
||||
combined = text_attn + pinyin_attn + part4_attn
|
||||
```
|
||||
|
||||
每路有独立的 K/V 投影层：
|
||||
- `k_text_proj` / `v_text_proj`（已有）
|
||||
- `k_pinyin_proj` / `v_pinyin_proj`（已有）
|
||||
- `k_part4_proj` / `v_part4_proj`（新增）
|
||||
|
||||
文件: `src/model/components.py` → `CrossAttentionFusion`
|
||||
|
||||
### 2.2 Part4 编码 (`_encode_part4`)
|
||||
|
||||
```
|
||||
part4_ids [B, N, L] 例: [128, 5, 6]
|
||||
→ text_emb (frozen, BERT wordpiece) → [B*N, L, dim]
|
||||
→ avg_pool (over valid tokens) → [B*N, dim]
|
||||
→ Linear(dim, dim) (可训练) → [B*N, dim]
|
||||
→ LayerNorm(dim) → [B*N, dim]
|
||||
→ reshape → [B, N, dim] = part4_E
|
||||
```
|
||||
|
||||
关键设计决策:
|
||||
- **复用 Encoder 的 BERT text embedding（`text_emb`）**——与 context_H 在同一语义空间
|
||||
- **可训练的投影层**（`part4_proj` + `part4_ln`）——允许模型自适应地将候选词映射到 CrossAttention 的有效空间
|
||||
|
||||
文件: `src/model/model.py` → `InputMethodEngine._encode_part4()`
|
||||
|
||||
### 2.3 Encoder 序列长度 128 → 88
|
||||
|
||||
| 项 | 旧值 | 新值 |
|
||||
|------|------|------|
|
||||
| `max_seq_len` | 128 | 88 |
|
||||
| 实际最大 token 数 | part4(3词/~9t) + prefix(48) + suffix(16) + spec(3) ≈ 76 | prefix(48) + suffix(16) + spec(3) = 67 |
|
||||
| pos_emb 行数 | 128 | 88 |
|
||||
| Encoder 自注意力 O(L²) | 128² | 88² (−53%) |
|
||||
|
||||
### 2.4 数据流总结
|
||||
|
||||
```
|
||||
┌─ part4_ids ──→ text_emb → avg_pool → Linear → part4_E ─┐
|
||||
│ │
|
||||
encoder (冻结/训练): │ CrossAttention (三路)
|
||||
input_ids ──→ H ─────────┼──→ text_attn ──────────────────────────────────────→ ─ 想
|
||||
pinyin_ids ─→ P ─────────┼──→ pinyin_attn ────────────────────────────────────→ ─ 想
|
||||
│ part4_attn → fused → MoE → logits
|
||||
│ │
|
||||
└──→ part4_E ─────────────────────────────────────────────┘
|
||||
```
|
||||
|
||||
---
|
||||
|
||||
## 3. 混淆数据流水线 (新方案)
|
||||
|
||||
### 3.1 概述
|
||||
|
||||
新方案**完全在 CPU 上运行**，无需 GPU 参与模型推理。流程为两步：
|
||||
|
||||
```
|
||||
┌──────────────────┐ ┌───────────────────────┐ ┌──────────────────────┐
|
||||
│ pinyin-index │ │ preprocess-model │ │ train-model │
|
||||
│ │ │ │ │ │
|
||||
│ 语料分词 + │ → │ 加载 pinyin_index │ → │ 加载 .npz (含 │
|
||||
│ 姓名词典 + │ │ 低频词生成混淆 part4 │ │ part4_ids) │
|
||||
│ 随机连接 │ │ 保存为 .npz 分片 │ │ 从头训练 / 微调 │
|
||||
└──────────────────┘ └───────────────────────┘ └──────────────────────┘
|
||||
```
|
||||
|
||||
### 3.2 步骤 1: pinyin-index
|
||||
|
||||
**前置**: 将 [Chinese-Names-Corpus](https://github.com/wainshine/Chinese-Names-Corpus) 克隆到本地:
|
||||
```bash
|
||||
git clone https://github.com/wainshine/Chinese-Names-Corpus.git /home/songsenand/DataSet/Chinese-Names-Corpus
|
||||
```
|
||||
|
||||
**输入**: 文本语料（HuggingFace 格式）+ Chinese-Names-Corpus 路径
|
||||
**输出**: `pinyin_index.msgpack`
|
||||
|
||||
```
|
||||
流程:
|
||||
1. 从 Chinese-Names-Corpus 加载词典:
|
||||
- 现代人名 (120W): 1,144,226 条
|
||||
- 性别标注人名 (120W): 1,144,226 条
|
||||
- 古代人名 (25W): 255,268 条
|
||||
- 成语 (5W): 50,372 条
|
||||
- 内嵌地名 (424 个省级/地级/主要城市)
|
||||
- 去重总计 ~142W 条 → 注入 jieba
|
||||
2. 流式读取语料，jieba 分词（HMM=False, 已注词典）
|
||||
3. 每个词 → 统计词频 + 计算拼音 → 加入 word_freq, word_to_pinyin
|
||||
4. 连续 2~4 个单字中文词 → 50% 概率随机合并 → 补充罕见多字词
|
||||
5. 过滤: 词频 ≥ 2 保留
|
||||
6. pinyin_to_words: 拼音 → 词集合
|
||||
7. 按拼音长度分桶 → Levenshtein 编辑距离 → 保留 sim ≥ 0.4 的 top-5 → similar_pinyin
|
||||
8. 序列化为 msgpack
|
||||
```
|
||||
|
||||
**输出格式**:
|
||||
```python
|
||||
{
|
||||
"pinyin_to_words": {
|
||||
"zhangwei": ["张伟", "章伟", "张维", ...],
|
||||
...
|
||||
},
|
||||
"similar_pinyin": {
|
||||
"zhangwei": ["zhangwei", "zangwei", "zhangwen", ...],
|
||||
...
|
||||
},
|
||||
"word_to_pinyin": {
|
||||
"张伟": "zhangwei", ...
|
||||
},
|
||||
"word_freq": {
|
||||
"张伟": 15, "中国": 50000, ...
|
||||
},
|
||||
"meta": { "num_words": ..., "num_pinyins": ... }
|
||||
}
|
||||
```
|
||||
|
||||
文件: `src/model/confusion/pinyin_index.py` → `PinyinIndexBuilder`
|
||||
|
||||
**使用**:
|
||||
```bash
|
||||
pinyin-index build \
|
||||
--corpus "path/to/hf_dataset" \
|
||||
--name-corpus /home/songsenand/DataSet/Chinese-Names-Corpus \
|
||||
--output ./banks/pinyin_index.msgpack \
|
||||
--max-samples 10000000
|
||||
```
|
||||
|
||||
**效果**: jieba 加载词典后单字词占比从 27.04% → 26.19% (降 0.85%), 超越 HanLP (26.55%). 词典修复了 "张志恒" (张|志|恒→张志恒)、"蔡畅" (蔡|畅→蔡畅) 等历史人名拆分问题.
|
||||
|
||||
### 3.3 步骤 2: preprocess-model (带混淆)
|
||||
|
||||
**输入**: 文本语料 + `pinyin_index.msgpack`
|
||||
**输出**: .npz 分片（含 part4_ids 字段）
|
||||
|
||||
```
|
||||
流程:
|
||||
1. 加载 pinyin_index
|
||||
2. PinyinInputDataset 生成基础样本（含 target_word, word_pinyin）
|
||||
3. 对每个样本:
|
||||
a. 查询 word_freq[target_word]
|
||||
b. 若 freq ≤ confusion_max_freq (默认 100) → 低频词, 触发混淆:
|
||||
- 查 word_to_pinyin[target_word] → pinyin
|
||||
- 查 pinyin_to_words[pinyin] → 同音词集合
|
||||
- 查 similar_pinyin[pinyin] → 近音词集合
|
||||
- 汇总所有候选词
|
||||
- 随机抽 0~4 个 (60% 概率有候选, 40% 空)
|
||||
- 50% 概率移除 target_word (防止模型走捷径)
|
||||
- _build_part4_ids(candidates) → part4_ids tensor
|
||||
c. 若 freq > confusion_max_freq → 高频词, part4_ids 为空
|
||||
4. 写入 .npz 分片 (FIELDS + part4_ids)
|
||||
```
|
||||
|
||||
文件: `src/model/preprocess.py` → `--pinyin-index`, `--confusion-max-freq`
|
||||
|
||||
**使用**:
|
||||
```bash
|
||||
preprocess-model \
|
||||
--train-data-path "path/to/hf_dataset" \
|
||||
--eval-data-path "path/to/eval/dataset" \
|
||||
--pinyin-index ./banks/pinyin_index.msgpack \
|
||||
--confusion-max-freq 100 \
|
||||
--output-dir ./samples \
|
||||
--num-train-samples 10000000 \
|
||||
--num-eval-samples 500000 \
|
||||
--max-seq-length 88
|
||||
```
|
||||
|
||||
### 3.5 编码器冻结 (`--freeze-encoder`)
|
||||
|
||||
启用时冻结 `ContextEncoder` 所有参数（包括 `text_emb`, `pinyin_emb`, `TransformerEncoder`, `PinyinLSTMEncoder`），仅训练 Decoder 部分:
|
||||
|
||||
- `part4_proj` / `part4_ln`（新增）
|
||||
- `cross_attn`（含新加的 `k_part4_proj` / `v_part4_proj`）
|
||||
- `slot_memory`
|
||||
- `moe`
|
||||
- `slot_attention` / `classifier`
|
||||
|
||||
建议学习率: 新模块 `5e-4`，已有模块 `1e-5`（分层 LR）。
|
||||
|
||||
---
|
||||
|
||||
## 4. ONNX 导出变更
|
||||
|
||||
### 4.1 Decoder 接口
|
||||
|
||||
```
|
||||
原:
|
||||
decoder(context_H, pinyin_P, history_slot_ids, context_mask, pinyin_mask) → logits
|
||||
|
||||
新:
|
||||
decoder(context_H, pinyin_P, history_slot_ids, context_mask, pinyin_mask,
|
||||
part4_ids) → logits
|
||||
```
|
||||
|
||||
### 4.2 动态维度
|
||||
|
||||
新增: `"part4_ids": {0: "batch_size", 1: "num_words"}`
|
||||
|
||||
文件: `src/model/export_models.py` → `DecoderExport`, `src/model/onnx_export.py`
|
||||
|
||||
---
|
||||
|
||||
## 5. CLI 命令参考
|
||||
|
||||
| 命令 | 用途 |
|
||||
|------|------|
|
||||
| `train-model train` | 从头训练 / 微调（自动检测 .npz → PreProcessedDataset） |
|
||||
| `pinyin-index build` | 从语料构建拼音索引（词→拼音→相似拼音） |
|
||||
| `preprocess-model` | 生成 .npz 分片（支持 `--pinyin-index` 生成混淆 part4） |
|
||||
| `train-model export` | 导出 ONNX（含 part4_ids 输入） |
|
||||
|
||||
---
|
||||
|
||||
## 6. 文件清单
|
||||
|
||||
### 新建
|
||||
|
||||
| 文件 | 职责 |
|
||||
|------|------|
|
||||
| `src/model/confusion/__init__.py` | 包导出 |
|
||||
| `src/model/confusion/pinyin_index.py` | `PinyinIndexBuilder` — 拼音索引构建（词→拼音→相似拼音） |
|
||||
| `src/model/confusion/cli.py` | CLI: `pinyin-index` |
|
||||
|
||||
### 修改
|
||||
|
||||
| 文件 | 关键变更 |
|
||||
|------|----------|
|
||||
| `src/model/components.py` | `CrossAttentionFusion` + 第三路 attention |
|
||||
| `src/model/model.py` | `_encode_part4()`, `part4_ids` forward, `freeze_encoder`, `max_seq_len=88` |
|
||||
| `src/model/dataset.py` | `target_word`/`word_pinyin`/`part4_ids` 字段, 移除 encoder 中 part4, `max_seq_length=88` |
|
||||
| `src/model/trainer.py` | `collate_fn` + part4_ids, `train_step` 传 part4, checkpoint 含完整模型 config |
|
||||
| `src/model/export_models.py` | `DecoderExport` + `_encode_part4` + `part4_ids` |
|
||||
| `src/model/onnx_export.py` | decoder 导出 + `part4_ids` |
|
||||
| `src/model/preprocess.py` | `--pinyin-index`, `--confusion-max-freq`, 生成含 part4_ids 的 .npz |
|
||||
| `src/model/preprocessed_dataset.py` | `part4_ids` 字段 + collate 支持 |
|
||||
| `scripts/finetune_slots.py` | `max_seq_len=88` |
|
||||
| `pyproject.toml` | CLI 入口 `pinyin-index`, `python-levenshtein` 依赖 |
|
||||
| `README.md` | §8 混淆数据文档（新流水线） |
|
||||
|
||||
---
|
||||
|
||||
## 7. 依赖
|
||||
|
||||
```bash
|
||||
uv pip install -e .
|
||||
uv pip install python-levenshtein
|
||||
```
|
||||
146
README.md
146
README.md
|
|
@ -53,7 +53,7 @@
|
|||
|
||||
| 参数项 | 推荐值 | 说明 |
|
||||
| :--- | :--- | :--- |
|
||||
| **序列长度 (L)** | 128 | 上下文窗口大小 [1] |
|
||||
| **序列长度 (L)** | 88 | 上下文窗口大小 [1] |
|
||||
| **隐藏层维度** | 512 | Embedding 及 Transformer 内部维度 [1] |
|
||||
| **Transformer 层数** | 4 | 轻量级骨干，降低延迟 [1] |
|
||||
| **注意力头数** | 4 | 适配 512 维度的高效配置 [1] |
|
||||
|
|
@ -171,7 +171,7 @@ config = {
|
|||
"n_layers": 4,
|
||||
"n_heads": 4,
|
||||
"num_experts": 20,
|
||||
"max_seq_len": 128,
|
||||
"max_seq_len": 88,
|
||||
|
||||
# 训练参数
|
||||
"batch_size": 64, # 根据GPU内存调整
|
||||
|
|
@ -437,6 +437,7 @@ def load_trained_model(checkpoint_path):
|
|||
2. **安装项目依赖**：
|
||||
```bash
|
||||
uv pip install -e .
|
||||
uv pip install python-levenshtein
|
||||
```
|
||||
|
||||
#### 使用传统 pip
|
||||
|
|
@ -837,5 +838,146 @@ train-model export \
|
|||
- 支持在不同推理引擎上部署
|
||||
- 提供优化后的推理模型
|
||||
|
||||
## 8. 混淆样本生成 (Pinyin Index & Confusion Sample Generation)
|
||||
|
||||
### 8.1 设计动机
|
||||
|
||||
原始训练中 `part4` 随机从文本分词结果中选取，训练时位于 Encoder 输入层，推理时不包含 `part4`，存在训练/推理分布不一致的问题。
|
||||
|
||||
本方案核心改动：
|
||||
- **part4 从 Encoder 移除**：Encoder 只处理光标前后文本，输出纯上下文编码 H
|
||||
- **part4 进入 Decoder**：通过独立的 CrossAttention 第三路（`k_part4_proj` / `v_part4_proj`）引入候选词信息
|
||||
- **序列长度从 128 缩减至 88**
|
||||
- **拼音相似词混淆**替代随机词：用同音/近音词作为 part4，增强模型区分同音字的能力
|
||||
|
||||
```
|
||||
Encoder (max_seq_len=88):
|
||||
input_ids → TransformerEncoder → H, P
|
||||
|
||||
Decoder:
|
||||
H, P, history_slot_ids, part4_ids → CrossAttention(三路) → MoE → logits
|
||||
text_attn + pinyin_attn + part4_attn
|
||||
```
|
||||
|
||||
### 8.2 混淆数据构建流程 (新流水线)
|
||||
|
||||
新方案移除对 GPU 的依赖，全部在 CPU 上完成：
|
||||
|
||||
```
|
||||
┌────────────────────┐ ┌──────────────────────┐ ┌──────────────────────┐
|
||||
│ pinyin-index build │ │ preprocess-model │ │ train-model train │
|
||||
│ │ │ │ │ │
|
||||
│ 142W姓名词典注入 │ → │ 加载 pinyin_index │ → │ 加载 .npz (含 │
|
||||
│ jieba + 语料分词 │ │ 低频词生成混淆 part4 │ │ part4_ids) │
|
||||
│ + 随机连接 + │ │ 保存为 .npz 分片 │ │ 从头训练 / 微调 │
|
||||
│ Levenshtein相似度 │ │ │ │ │
|
||||
└────────────────────┘ └──────────────────────┘ └──────────────────────┘
|
||||
```
|
||||
|
||||
#### 步骤 1: 构建拼音索引 (pinyin-index)
|
||||
|
||||
从文本语料构建拼音→词映射，无需模型推理。
|
||||
|
||||
**词典准备 (前置)**：将 [Chinese-Names-Corpus](https://github.com/wainshine/Chinese-Names-Corpus) 克隆至本地：
|
||||
|
||||
```bash
|
||||
git clone https://github.com/wainshine/Chinese-Names-Corpus.git /home/songsenand/DataSet/Chinese-Names-Corpus
|
||||
```
|
||||
|
||||
**构建索引**：
|
||||
|
||||
```bash
|
||||
pinyin-index build \
|
||||
--corpus "path/to/hf_dataset" \
|
||||
--name-corpus /home/songsenand/DataSet/Chinese-Names-Corpus \
|
||||
--output ./banks/pinyin_index.msgpack \
|
||||
--max-samples 10000000
|
||||
```
|
||||
|
||||
核心策略：
|
||||
1. **142W 外挂词典**：从 `Chinese-Names-Corpus` 加载 142 万人名 + 5W 成语 + 424 内嵌地名，注入 jieba，大幅减少人名/地名拆分（单字词占比从 27.04% 降至 26.19%，超越 HanLP 的 26.55%）
|
||||
2. **随机连接补位**：词典之外，连续 2~4 个单字词以 50% 概率随机合并，补充罕见多字词
|
||||
3. **拼音相似度**：按拼音长度分桶，Levenshtein 编辑距离计算 top-5 相似拼音
|
||||
|
||||
产出 `pinyin_index.msgpack`，包含：
|
||||
- `pinyin_to_words`: 拼音 → 词集合
|
||||
- `similar_pinyin`: 拼音 → 相似拼音列表
|
||||
- `word_to_pinyin`: 词 → 拼音
|
||||
- `word_freq`: 词 → 频率
|
||||
|
||||
#### 步骤 2: 生成带混淆的 .npz 样本
|
||||
|
||||
```bash
|
||||
preprocess-model \
|
||||
--train-data-path "path/to/hf_dataset" \
|
||||
--eval-data-path "path/to/eval/dataset" \
|
||||
--pinyin-index ./banks/pinyin_index.msgpack \
|
||||
--confusion-max-freq 100 \
|
||||
--output-dir ./samples \
|
||||
--num-train-samples 10000000 \
|
||||
--num-eval-samples 500000 \
|
||||
--max-seq-length 88
|
||||
```
|
||||
|
||||
`--confusion-max-freq` 控制混淆频率阈值：只有词频 ≤100 的低频词才生成混淆 part4，高频词不提供 part4（因为模型已从海量样本中学会）。
|
||||
|
||||
#### 步骤 3: 训练
|
||||
|
||||
```bash
|
||||
# 从头训练（使用 .npz 分片，自动加载 part4_ids）
|
||||
train-model train \
|
||||
--train-data-path ./samples/train \
|
||||
--eval-data-path ./samples/eval \
|
||||
--output-dir ./output \
|
||||
--batch-size 128 \
|
||||
--num-epochs 10 \
|
||||
--learning-rate 2e-4
|
||||
|
||||
# 微调（冻结 Encoder）
|
||||
train-model train \
|
||||
--train-data-path ./samples/train \
|
||||
--eval-data-path ./samples/eval \
|
||||
--resume-from ./output/checkpoints/best_model.pt \
|
||||
--freeze-encoder \
|
||||
--output-dir ./finetune_output \
|
||||
--learning-rate 1e-5 \
|
||||
--num-epochs 5
|
||||
```
|
||||
|
||||
### 8.3 混淆生成策略
|
||||
|
||||
| 参数 | 默认值 | 说明 |
|
||||
|------|--------|------|
|
||||
| 混淆比例 | 60% | 低频词样本中提供混淆 part4 的概率 |
|
||||
| 空比例 | 40% | 低频词样本中 part4 为空的概率 |
|
||||
| 目标剔除 | 50% | 若目标词在候选列表，一半概率移除，防止模型走捷径 |
|
||||
| 最大候选数 | 5 | part4 最多 5 个候选词 |
|
||||
| 频率阈值 | 100 | 词频 ≤100 才触发混淆，高频词跳过 |
|
||||
|
||||
### 8.4 编码器冻结 (--freeze-encoder)
|
||||
|
||||
启用时冻结 ContextEncoder 的所有参数，仅训练 Decoder 部分（包括 `part4_proj`、`cross_attn.k_part4_proj/v_part4_proj` 及后续模块）。
|
||||
|
||||
### 8.5 CLI 命令参考
|
||||
|
||||
| 命令 | 用途 |
|
||||
|------|------|
|
||||
| `train-model train` | 从头训练 / 微调 |
|
||||
| `pinyin-index build` | 构建拼音索引（词→拼音→相似拼音） |
|
||||
| `preprocess-model` | 生成 .npz 分片（支持 `--pinyin-index` 生成混淆 part4） |
|
||||
| `train-model export` | 导出 ONNX（含 part4_ids 输入） |
|
||||
|
||||
### 8.6 文件清单
|
||||
|
||||
| 文件 | 职责 |
|
||||
|------|------|
|
||||
| `src/model/confusion/pinyin_index.py` | `PinyinIndexBuilder` — 拼音索引构建 |
|
||||
| `src/model/confusion/cli.py` | CLI: `pinyin-index` |
|
||||
| `src/model/preprocess.py` | `--pinyin-index` 参数，生成含 part4_ids 的 .npz |
|
||||
| `src/model/preprocessed_dataset.py` | `part4_ids` 字段支持 |
|
||||
| `src/model/components.py` | `CrossAttentionFusion` 三路 attention |
|
||||
| `src/model/model.py` | `_encode_part4()`, `freeze_encoder`, `max_seq_len=88` |
|
||||
| `src/model/trainer.py` | `collate_fn` + part4_ids, `--freeze-encoder` |
|
||||
|
||||
## 7. 总结
|
||||
本方案通过**单流 Transformer 编码**结合**结构化槽位交叉注意力**，并引入**20个专家的 MoE 模块** [1]，在保证模型轻量（4层 Transformer）的同时，有效利用了历史输入习惯并提升了模型表达上限。相比暴力拼接或双流架构，该设计在工程实现上更优雅，在推理效率上更高效，是轻量级输入法模型的局部最优解。
|
||||
|
|
@ -0,0 +1,26 @@
|
|||
{
|
||||
"jieba": {
|
||||
"status": "ok",
|
||||
"cold_start_s": 0.219,
|
||||
"memory_mb": 39,
|
||||
"peak_memory_mb": 95,
|
||||
"chars_per_sec": 1167606,
|
||||
"description": "jieba (基线, HMM=False)"
|
||||
},
|
||||
"LTP": {
|
||||
"status": "ok",
|
||||
"cold_start_s": 4.811,
|
||||
"memory_mb": 839,
|
||||
"peak_memory_mb": 936,
|
||||
"chars_per_sec": 2748,
|
||||
"description": "LTP (pipeline, cws)"
|
||||
},
|
||||
"HanLP": {
|
||||
"status": "ok",
|
||||
"cold_start_s": 0.161,
|
||||
"memory_mb": 95,
|
||||
"peak_memory_mb": 984,
|
||||
"chars_per_sec": 2523,
|
||||
"description": "HanLP (COARSE_ELECTRA_SMALL)"
|
||||
}
|
||||
}
|
||||
|
|
@ -0,0 +1,161 @@
|
|||
# 混淆样本生成流水线 (Confusion Sample Generation)
|
||||
|
||||
## 概述
|
||||
|
||||
本流水线用于生成带混淆候选词（part4）的训练样本，全部在 CPU 上完成，无需 GPU。
|
||||
|
||||
核心思路：用拼音相似词作为 part4 候选，帮助模型学会区分同音/近音字。
|
||||
|
||||
## 三步流程
|
||||
|
||||
```
|
||||
Step 1 (CPU): 语料 + 142W 词典 → pinyin-index → pinyin_index.msgpack
|
||||
Step 2 (CPU): 语料 + pinyin_index.msgpack → preprocess-model → .npz 分片
|
||||
Step 3 (GPU): .npz 分片 → train-model → 模型
|
||||
```
|
||||
|
||||
---
|
||||
|
||||
## Step 1: 构建拼音索引
|
||||
|
||||
### 前置准备
|
||||
|
||||
将 [Chinese-Names-Corpus](https://github.com/wainshine/Chinese-Names-Corpus) 克隆至本地：
|
||||
|
||||
```bash
|
||||
git clone https://github.com/wainshine/Chinese-Names-Corpus.git ~/DataSet/Chinese-Names-Corpus
|
||||
```
|
||||
|
||||
### 执行
|
||||
|
||||
```bash
|
||||
pinyin-index build \
|
||||
--corpus "path/to/hf_dataset" \
|
||||
--name-corpus ~/DataSet/Chinese-Names-Corpus \
|
||||
--output ./banks/pinyin_index.msgpack \
|
||||
--max-samples 10000000
|
||||
```
|
||||
|
||||
### 参数说明
|
||||
|
||||
| 参数 | 默认值 | 说明 |
|
||||
|------|--------|------|
|
||||
| `--corpus` | (必需) | HuggingFace 格式语料路径 |
|
||||
| `--name-corpus` | None | Chinese-Names-Corpus 仓库路径。不提供则仅用内嵌 424 地名 |
|
||||
| `--output` | `./banks/pinyin_index.msgpack` | 输出路径 |
|
||||
| `--max-samples` | 10,000,000 | 最大处理文本数 |
|
||||
| `--top-k-similar` | 5 | 每个拼音保留的相似拼音数 |
|
||||
| `--min-similarity` | 0.4 | Levenshtein 最小相似度阈值 |
|
||||
| `--merge-prob` | 0.5 | 连续单字词随机合并概率 |
|
||||
| `--min-word-freq` | 2 | 最小词频过滤 |
|
||||
|
||||
### 词典构成 (~142 万条)
|
||||
|
||||
| 来源 | 词条数 | 说明 |
|
||||
|------|--------|------|
|
||||
| Chinese_Names_Corpus (120W) | 1,144,226 | 现代中文人名 |
|
||||
| Chinese_Names_Corpus_Gender (120W) | 1,144,226 | 性别标注人名 |
|
||||
| Ancient_Names_Corpus (25W) | 255,268 | 古代人名 |
|
||||
| ChengYu_Corpus (5W) | 50,372 | 成语 |
|
||||
| 内嵌地名 | 424 | 省级/地级/省会/主要城市/大学 |
|
||||
|
||||
**效果**: jieba 加载词典后，单字词占比从 27.04% 降至 26.19%（降 0.85%），超越 HanLP（26.55%）。实际修复了 "张志恒"、"蔡畅" 等历史人名拆分问题。
|
||||
|
||||
### 输出格式
|
||||
|
||||
`pinyin_index.msgpack` 包含四个核心 dict：
|
||||
|
||||
```python
|
||||
{
|
||||
"pinyin_to_words": {"zhangwei": ["张伟", "章伟", ...]}, # pinyin → words
|
||||
"similar_pinyin": {"zhangwei": ["zangwei", ...]}, # pinyin → similar pinyins
|
||||
"word_to_pinyin": {"张伟": "zhangwei"}, # word → pinyin
|
||||
"word_freq": {"张伟": 15, "中国": 50000}, # word → frequency
|
||||
"meta": {"num_words": ..., "num_pinyins": ...}
|
||||
}
|
||||
```
|
||||
|
||||
---
|
||||
|
||||
## Step 2: 生成带混淆的 .npz 样本
|
||||
|
||||
```bash
|
||||
preprocess-model \
|
||||
--train-data-path "path/to/hf_dataset" \
|
||||
--eval-data-path "path/to/eval/dataset" \
|
||||
--pinyin-index ./banks/pinyin_index.msgpack \
|
||||
--confusion-max-freq 100 \
|
||||
--output-dir ./samples \
|
||||
--num-train-samples 10000000 \
|
||||
--num-eval-samples 500000 \
|
||||
--max-seq-length 88
|
||||
```
|
||||
|
||||
### 混淆生成逻辑
|
||||
|
||||
对每个训练样本：
|
||||
|
||||
```
|
||||
target_word → 查 word_freq
|
||||
├─ freq ≤ confusion_max_freq → 低频词 → 触发混淆:
|
||||
│ 1. 查 word_to_pinyin[target_word] → pinyin
|
||||
│ 2. 查 pinyin_to_words[pinyin] → 同音词集
|
||||
│ 3. 查 similar_pinyin[pinyin] → 近音词集
|
||||
│ 4. 汇总所有候选, 随机抽 1-4 个 (60% 概率)
|
||||
│ 5. 50% 概率移除 target_word (防捷径)
|
||||
│ 6. 构建 part4_ids → 写入 .npz
|
||||
│
|
||||
└─ freq > confusion_max_freq → 高频词 → part4_ids 为空
|
||||
```
|
||||
|
||||
### 关键参数
|
||||
|
||||
| 参数 | 默认值 | 说明 |
|
||||
|------|--------|------|
|
||||
| `--pinyin-index` | None | pinyin_index.msgpack 路径，不提供则不生成混淆 |
|
||||
| `--confusion-max-freq` | 100 | 词频阈值，≤ N 才触发混淆 |
|
||||
| `--max-seq-length` | 88 | 序列长度 |
|
||||
|
||||
---
|
||||
|
||||
## Step 3: 训练
|
||||
|
||||
### 从头训练
|
||||
|
||||
```bash
|
||||
train-model train \
|
||||
--train-data-path ./samples/train \
|
||||
--eval-data-path ./samples/eval \
|
||||
--output-dir ./output \
|
||||
--batch-size 128 \
|
||||
--num-epochs 10 \
|
||||
--learning-rate 2e-4
|
||||
```
|
||||
|
||||
### 微调 (冻结 Encoder)
|
||||
|
||||
```bash
|
||||
train-model train \
|
||||
--train-data-path ./samples/train \
|
||||
--eval-data-path ./samples/eval \
|
||||
--resume-from ./output/checkpoints/best_model.pt \
|
||||
--freeze-encoder \
|
||||
--output-dir ./finetune_output \
|
||||
--learning-rate 1e-5 \
|
||||
--num-epochs 5
|
||||
```
|
||||
|
||||
---
|
||||
|
||||
## 补充词典
|
||||
|
||||
### 添加新词典源
|
||||
|
||||
如果要添加更多专业术语词典（如法律、医学、IT），只需在文本文件中每行一个词，然后用 `jieba.add_word()` 加载即可。
|
||||
|
||||
也可以修改 `src/model/confusion/pinyin_index.py` 中的 `load_name_dict()` 函数添加新的词典文件。
|
||||
|
||||
参考资源:
|
||||
- [Chinese-Tripitaka/Chinese-medical-dictionary](https://github.com/Chinese-Tripitaka/Chinese-medical-dictionary) — 医学
|
||||
- [fighting41love/funNLP](https://github.com/fighting41love/funNLP) — 多领域词库
|
||||
- [THUOCL](http://thuocl.thunlp.org/) — 清华开放中文词库 (IT/财经/医疗)
|
||||
|
|
@ -16,6 +16,7 @@ dependencies = [
|
|||
"plotly>=5.0.0",
|
||||
"jieba>=0.42.1",
|
||||
"pypinyin>=0.55.0",
|
||||
"python-levenshtein>=0.25.0",
|
||||
"requests>=2.32.5",
|
||||
"rich>=14.3.1",
|
||||
"flask>=3.1.0",
|
||||
|
|
@ -32,6 +33,7 @@ train-model = "model.trainer:app"
|
|||
monitor-training = "model.monitor:app"
|
||||
preprocess-model = "model.preprocess:main"
|
||||
inspect-preprocessed = "model.inspect_preprocessed:main"
|
||||
pinyin-index = "model.confusion.cli:pinyin_index_app"
|
||||
|
||||
[tool.uv]
|
||||
# 设置当前项目的默认索引源
|
||||
|
|
|
|||
|
|
@ -0,0 +1,228 @@
|
|||
#!/usr/bin/env python3
|
||||
"""
|
||||
真实数据分词对比: jieba(基线) vs jieba+姓名词典 vs HanLP
|
||||
|
||||
用法:
|
||||
HF_ENDPOINT=https://hf-mirror.com \
|
||||
python scripts/benchmark_real_data.py \
|
||||
--data-dir /home/songsenand/DataSet/data/data \
|
||||
--num-samples 200
|
||||
"""
|
||||
|
||||
import argparse
|
||||
import json
|
||||
import os
|
||||
import random
|
||||
import sys
|
||||
import time
|
||||
|
||||
import numpy as np
|
||||
|
||||
parser = argparse.ArgumentParser()
|
||||
parser.add_argument("--data-dir", required=True, help="CCI-Data data 目录路径")
|
||||
parser.add_argument("--num-samples", type=int, default=200, help="抽样数量")
|
||||
parser.add_argument("--seed", type=int, default=42)
|
||||
args = parser.parse_args()
|
||||
|
||||
random.seed(args.seed)
|
||||
np.random.seed(args.seed)
|
||||
|
||||
# ── 加载文本 ──────────────────────────────────────────────────────────
|
||||
data_dir = args.data_dir
|
||||
jsonl_files = sorted([f for f in os.listdir(data_dir) if f.endswith(".jsonl")])
|
||||
num_files_to_read = min(50, len(jsonl_files))
|
||||
selected_files = random.sample(jsonl_files, num_files_to_read)
|
||||
print(f"从 {len(jsonl_files)} 个文件中随机选取 {num_files_to_read} 个读取...")
|
||||
|
||||
all_texts = []
|
||||
for fname in selected_files:
|
||||
with open(os.path.join(data_dir, fname), encoding="utf-8") as f:
|
||||
for line in f:
|
||||
obj = json.loads(line)
|
||||
text = obj.get("text", "")
|
||||
if text and len(text) >= 20:
|
||||
all_texts.append(text)
|
||||
if len(all_texts) >= args.num_samples * 10:
|
||||
break
|
||||
|
||||
sampled = random.sample(all_texts, min(args.num_samples, len(all_texts)))
|
||||
sample_chars = sum(len(t) for t in sampled)
|
||||
print(f"抽样 {len(sampled):,} 条, {sample_chars:,} 字符")
|
||||
print(f"平均长度: {sample_chars / len(sampled):.0f} 字/条\n")
|
||||
|
||||
# ── 1. jieba (基线, HMM=False) ───────────────────────────────────────
|
||||
print("=" * 60)
|
||||
print(" 1. jieba (HMM=False, 基线)")
|
||||
print("=" * 60)
|
||||
|
||||
import jieba
|
||||
|
||||
jieba.setLogLevel(jieba.logging.INFO)
|
||||
t0 = time.perf_counter()
|
||||
jieba_results = [list(jieba.cut(t, HMM=False)) for t in sampled]
|
||||
t1 = time.perf_counter()
|
||||
jieba_time = t1 - t0
|
||||
|
||||
# ── 2. jieba + 姓名词典 (Full Corpus) ───────────────────────────────
|
||||
print()
|
||||
print("=" * 60)
|
||||
print(" 2. jieba + 完整词典 (142W 人名 + 地名 + 成语)")
|
||||
print("=" * 60)
|
||||
|
||||
CORPUS_PATH = "/home/songsenand/DataSet/Chinese-Names-Corpus"
|
||||
|
||||
|
||||
def _load_txt(path, min_len=2, max_len=8, skip=4):
|
||||
items = set()
|
||||
with open(path, encoding="utf-8-sig") as f:
|
||||
for i, line in enumerate(f):
|
||||
if i < skip:
|
||||
continue
|
||||
w = line.strip()
|
||||
if w and min_len <= len(w) <= max_len:
|
||||
items.add(w)
|
||||
return items
|
||||
|
||||
|
||||
def _load_csv(path, col=0, min_len=2, max_len=8, skip=4):
|
||||
items = set()
|
||||
with open(path, encoding="utf-8-sig") as f:
|
||||
for i, line in enumerate(f):
|
||||
if i < skip:
|
||||
continue
|
||||
parts = line.strip().split(",")
|
||||
if len(parts) > col:
|
||||
w = parts[col].strip()
|
||||
if w and min_len <= len(w) <= max_len:
|
||||
items.add(w)
|
||||
return items
|
||||
|
||||
|
||||
# 内嵌地名 (与 pinyin_index.py 一致)
|
||||
_PLACES = [
|
||||
"北京", "天津", "上海", "重庆", "河北", "山西", "辽宁", "吉林", "黑龙江",
|
||||
"江苏", "浙江", "安徽", "福建", "江西", "山东", "河南", "湖北", "湖南",
|
||||
"广东", "海南", "四川", "贵州", "云南", "陕西", "甘肃", "青海", "台湾",
|
||||
"内蒙古", "广西", "西藏", "宁夏", "新疆", "香港", "澳门",
|
||||
"石家庄", "太原", "沈阳", "长春", "哈尔滨", "南京", "杭州", "合肥",
|
||||
"福州", "南昌", "济南", "郑州", "武汉", "长沙", "广州", "南宁",
|
||||
"成都", "贵阳", "昆明", "拉萨", "西安", "兰州", "西宁", "银川",
|
||||
"乌鲁木齐", "呼和浩特", "大连", "青岛", "厦门", "深圳", "苏州", "无锡",
|
||||
"宁波", "温州", "佛山", "东莞", "常州", "南通", "徐州", "烟台", "威海",
|
||||
"洛阳", "开封", "岳阳", "株洲", "湘潭", "珠海", "汕头", "三亚",
|
||||
"桂林", "丽江", "大理", "张家界", "黄山", "武夷山", "泰山", "华山",
|
||||
"清华", "北大", "复旦", "浙大", "南大",
|
||||
]
|
||||
|
||||
all_names = set(_PLACES)
|
||||
base = f"{CORPUS_PATH}/Chinese_Names_Corpus"
|
||||
dict_base = f"{CORPUS_PATH}/Chinese_Dict_Corpus"
|
||||
|
||||
for label, loader, path in [
|
||||
("Names(120W)", _load_txt, f"{base}/Chinese_Names_Corpus（120W）.txt"),
|
||||
("Names(Gender)", _load_csv, f"{base}/Chinese_Names_Corpus_Gender（120W）.txt"),
|
||||
("Names(Ancient)", _load_txt, f"{base}/Ancient_Names_Corpus（25W）.txt"),
|
||||
("Chengyu(5W)", _load_txt, f"{dict_base}/ChengYu_Corpus（5W）.txt"),
|
||||
]:
|
||||
items = loader(path, min_len=2, max_len=12)
|
||||
all_names.update(items)
|
||||
print(f" {label}: {len(items):,}")
|
||||
|
||||
print(f" 加载 {len(all_names):,} 个词条到 jieba 词典...")
|
||||
for name in all_names:
|
||||
jieba.add_word(name)
|
||||
|
||||
t0 = time.perf_counter()
|
||||
jieba_dict_results = [list(jieba.cut(t, HMM=False)) for t in sampled]
|
||||
t2 = time.perf_counter()
|
||||
jieba_dict_time = t2 - t0
|
||||
|
||||
# ── 3. HanLP ─────────────────────────────────────────────────────────
|
||||
print()
|
||||
print("=" * 60)
|
||||
print(" 3. HanLP (COARSE_ELECTRA_SMALL)")
|
||||
print("=" * 60)
|
||||
|
||||
os.environ["CUDA_VISIBLE_DEVICES"] = ""
|
||||
import hanlp
|
||||
from loguru import logger
|
||||
|
||||
logger.disable("hanlp")
|
||||
tok = hanlp.load(hanlp.pretrained.tok.COARSE_ELECTRA_SMALL_ZH)
|
||||
tok(["热身"])
|
||||
|
||||
t0 = time.perf_counter()
|
||||
hanlp_results = []
|
||||
for i in range(0, len(sampled), 32):
|
||||
batch = sampled[i : i + 32]
|
||||
hanlp_results.extend(tok(batch))
|
||||
pct = min(i + 32, len(sampled)) / len(sampled) * 100
|
||||
elapsed = time.perf_counter() - t0
|
||||
print(f" 进度: {pct:.0f}% ({min(i+32, len(sampled))}/{len(sampled)}), 耗时 {elapsed:.0f}s", end="\r")
|
||||
print()
|
||||
t3 = time.perf_counter()
|
||||
hanlp_time = t3 - t0
|
||||
|
||||
# ── 统计 ──────────────────────────────────────────────────────────────
|
||||
def stats(results, label):
|
||||
single = 0
|
||||
total = 0
|
||||
for words in results:
|
||||
for w in words:
|
||||
if len(w) == 1 and "\u4e00" <= w <= "\u9fff":
|
||||
single += 1
|
||||
total += 1
|
||||
ratio = single / total if total else 0
|
||||
return single, total, ratio
|
||||
|
||||
j_sc, j_tc, j_ratio = stats(jieba_results, "jieba")
|
||||
jd_sc, jd_tc, jd_ratio = stats(jieba_dict_results, "jieba+dict")
|
||||
h_sc, h_tc, h_ratio = stats(hanlp_results, "HanLP")
|
||||
|
||||
j_cps = sample_chars / jieba_time
|
||||
jd_cps = sample_chars / jieba_dict_time
|
||||
h_cps = sample_chars / hanlp_time
|
||||
|
||||
# ── 对比表 ────────────────────────────────────────────────────────────
|
||||
print()
|
||||
print("=" * 80)
|
||||
print(" 📊 对比")
|
||||
print("=" * 80)
|
||||
print(f" {'指标':<20s} {'jieba':>15s} {'jieba+词典':>15s} {'HanLP':>15s}")
|
||||
print(f" {'─'*20} {'─'*15} {'─'*15} {'─'*15}")
|
||||
print(f" {'速度 (字/秒)':<20s} {j_cps:>15,.0f} {jd_cps:>15,.0f} {h_cps:>15,.0f}")
|
||||
print(f" {'耗时 (s)':<20s} {jieba_time:>15.2f} {jieba_dict_time:>15.2f} {hanlp_time:>15.2f}")
|
||||
print(f" {'单字词占比':<20s} {j_ratio:>14.2%} {jd_ratio:>14.2%} {h_ratio:>14.2%}")
|
||||
print(f" {'单字词数':<20s} {j_sc:>15,} {jd_sc:>15,} {h_sc:>15,}")
|
||||
print(f" {'总词数':<20s} {j_tc:>15,} {jd_tc:>15,} {h_tc:>15,}")
|
||||
jd_vs_h = jd_cps / h_cps
|
||||
print(f" {'速度比 (dict/HanLP)':<20s} {'—':>15s} {'—':>15s} {jd_vs_h:>14.0f}x")
|
||||
print(f" {'单字率降幅 (vs jieba)':<20s} {'—':>15s} {j_ratio - jd_ratio:>14.2%} {j_ratio - h_ratio:>14.2%}")
|
||||
|
||||
# ── 差异案例 ──────────────────────────────────────────────────────────
|
||||
print()
|
||||
print("=" * 80)
|
||||
print(" 🔍 差异案例 (jieba 拆开 → jieba+词典 不拆)")
|
||||
print("=" * 80)
|
||||
|
||||
shown = 0
|
||||
for i in range(len(sampled)):
|
||||
if shown >= 10:
|
||||
break
|
||||
j_words = jieba_results[i]
|
||||
jd_words = jieba_dict_results[i]
|
||||
if j_words == jd_words:
|
||||
continue
|
||||
shown += 1
|
||||
print(f"\n [{shown}] 原文: {sampled[i][:100]}...")
|
||||
print(f" jieba : {' | '.join(j_words[:25])}")
|
||||
print(f" jieba+词典: {' | '.join(jd_words[:25])}")
|
||||
|
||||
# ── 汇总 ──────────────────────────────────────────────────────────────
|
||||
print()
|
||||
print("=" * 80)
|
||||
print(f" ✅ 结论:")
|
||||
print(f" jieba+词典 速度: {jd_cps:,.0f} 字/秒 (几乎无损)")
|
||||
print(f" 单字词占比: jieba {j_ratio:.2%} → jieba+词典 {jd_ratio:.2%} (降 {j_ratio - jd_ratio:.2%})")
|
||||
print(f" jieba+词典 vs HanLP: {jd_vs_h:.0f}x 更快")
|
||||
print("=" * 80)
|
||||
|
|
@ -0,0 +1,226 @@
|
|||
#!/usr/bin/env python3
|
||||
"""
|
||||
Chinese Word Segmentation Benchmark
|
||||
比较 jieba / LTP / HanLP（LAC 不兼容 Python 3.12）
|
||||
|
||||
用法:
|
||||
HF_ENDPOINT=https://hf-mirror.com python scripts/benchmark_segmentation.py
|
||||
|
||||
依赖安装（独立 venv）:
|
||||
uv venv /tmp/segbench-venv --python 3.12
|
||||
uv pip install --python /tmp/segbench-venv/bin/python \
|
||||
psutil ltp hanlp "huggingface-hub<1.0,>=0.26" "transformers>=4.40,<5.0"
|
||||
"""
|
||||
|
||||
import os
|
||||
import sys
|
||||
import time
|
||||
import json
|
||||
import textwrap
|
||||
from typing import Dict, List, Tuple, Callable, Optional
|
||||
|
||||
import psutil
|
||||
|
||||
# ── Test Cases ──────────────────────────────────────────────────────────
|
||||
|
||||
TEST_CASES: List[Tuple[str, str, str]] = [
|
||||
# (类别, 原文, 关注点)
|
||||
("常见人名", "张伟去了上海找李明", "张伟、李明为人名，不应拆开"),
|
||||
("常见人名2", "王芳和李刚一起去北京", "王芳、李刚为人名"),
|
||||
("罕见复姓", "欧阳修与司马光论道饮酒", "欧阳修、司马光为复姓人名"),
|
||||
("三字人名", "诸葛亮借东风火烧赤壁", "诸葛亮为历史人名"),
|
||||
("长地名", "乌鲁木齐的天气比呼和浩特好多了", "乌鲁木齐、呼和浩特为地名"),
|
||||
("品牌名", "使用华为手机拍摄视频发朋友圈", "华为为品牌名"),
|
||||
("公司名", "字节跳动发布了豆包大模型产品", "字节跳动为公司名"),
|
||||
("罕见词", "这道菜真是饕餮盛宴令人垂涎", "饕餮为罕见词"),
|
||||
("谚语", "不到长城非好汉屈指行程二万", "长城为地名"),
|
||||
("经典歧义", "武汉市长江大桥通车了", "歧义: 武汉/市长/江大桥 vs 武汉市/长江/大桥"),
|
||||
("混合场景", "张三在北京用华为手机给李四发微信", "综合测试"),
|
||||
("常见词对照", "今天天气很好中国人民站起来了", "常见词基线，应正常分词"),
|
||||
("机构名", "中国科学院计算技术研究所发布报告", "机构名不应过度拆分"),
|
||||
("外来词", "使用ChatGPT生成文本然后发表", "中英混合"),
|
||||
]
|
||||
|
||||
# ── Library Wrappers ────────────────────────────────────────────────────
|
||||
|
||||
def setup_jieba():
|
||||
import jieba
|
||||
jieba.setLogLevel(jieba.logging.INFO)
|
||||
return lambda text: list(jieba.cut(text, HMM=False))
|
||||
|
||||
|
||||
def setup_ltp():
|
||||
from ltp import LTP
|
||||
model = LTP()
|
||||
model.pipeline(["热身"], tasks=["cws"])
|
||||
return lambda text: model.pipeline([text], tasks=["cws"]).cws[0]
|
||||
|
||||
|
||||
def setup_hanlp():
|
||||
import hanlp
|
||||
tok = hanlp.load(hanlp.pretrained.tok.COARSE_ELECTRA_SMALL_ZH)
|
||||
tok("热身")
|
||||
return lambda text: tok(text)
|
||||
|
||||
|
||||
def setup_lac():
|
||||
"""LAC 2.1.2 requires paddle.fluid which is removed in Python 3.12."""
|
||||
from LAC import LAC
|
||||
lac = LAC(mode="seg")
|
||||
lac.run("热身")
|
||||
return lambda text: lac.run(text)
|
||||
|
||||
|
||||
LIBRARIES: Dict[str, Tuple[Callable[[], Optional[Callable]], str]] = {
|
||||
"jieba": (setup_jieba, "jieba (基线, HMM=False)"),
|
||||
"LTP": (setup_ltp, "LTP (pipeline, cws)"),
|
||||
"HanLP": (setup_hanlp, "HanLP (COARSE_ELECTRA_SMALL)"),
|
||||
}
|
||||
|
||||
# ── Benchmark Logic ─────────────────────────────────────────────────────
|
||||
|
||||
def measure_memory() -> int:
|
||||
return psutil.Process(os.getpid()).memory_info().rss // (1024 * 1024)
|
||||
|
||||
|
||||
def run_benchmark() -> Dict:
|
||||
results = {}
|
||||
|
||||
for name, (setup_fn, description) in LIBRARIES.items():
|
||||
print(f"\n{'='*60}")
|
||||
print(f" 测试 {name}: {description}")
|
||||
print(f"{'='*60}")
|
||||
|
||||
mem_before = measure_memory()
|
||||
|
||||
# Cold start
|
||||
t0 = time.perf_counter()
|
||||
try:
|
||||
seg_fn = setup_fn()
|
||||
except Exception as e:
|
||||
print(f" ❌ 初始化失败: {e}")
|
||||
results[name] = {"status": "failed", "error": str(e)}
|
||||
continue
|
||||
cold_time = time.perf_counter() - t0
|
||||
|
||||
mem_after = measure_memory()
|
||||
|
||||
# Segmentation results
|
||||
print(f"\n {'原文':<24s} │ 分词结果")
|
||||
print(f" {'─'*24}─┼─{'─'*50}")
|
||||
|
||||
seg_outputs = {}
|
||||
for cat, text, note in TEST_CASES:
|
||||
words = seg_fn(text)
|
||||
seg_outputs[text] = words
|
||||
result_str = " | ".join(words)
|
||||
if len(result_str) > 50:
|
||||
result_str = result_str[:47] + "..."
|
||||
print(f" {text:<24s} │ {result_str}")
|
||||
|
||||
# Speed benchmark (warm runs)
|
||||
speed_texts = [text for _, text, _ in TEST_CASES]
|
||||
warm_runs = 5
|
||||
bench_runs = 100
|
||||
|
||||
for _ in range(warm_runs):
|
||||
for t in speed_texts:
|
||||
seg_fn(t)
|
||||
|
||||
t0 = time.perf_counter()
|
||||
for _ in range(bench_runs):
|
||||
for t in speed_texts:
|
||||
seg_fn(t)
|
||||
total_time = time.perf_counter() - t0
|
||||
total_chars = sum(len(t) for t in speed_texts) * bench_runs
|
||||
chars_per_sec = total_chars / total_time
|
||||
|
||||
results[name] = {
|
||||
"status": "ok",
|
||||
"cold_start_s": round(cold_time, 3),
|
||||
"memory_mb": mem_after - mem_before,
|
||||
"peak_memory_mb": measure_memory(),
|
||||
"chars_per_sec": int(chars_per_sec),
|
||||
"seg_outputs": seg_outputs,
|
||||
"description": description,
|
||||
}
|
||||
|
||||
print(f"\n 📊 冷启动: {cold_time:.2f}s | 速度: {chars_per_sec:,.0f} 字/秒 | 内存增量: {mem_after - mem_before} MB")
|
||||
|
||||
return results
|
||||
|
||||
|
||||
# ── Comparison Table ────────────────────────────────────────────────────
|
||||
|
||||
def print_comparison(results: Dict):
|
||||
print(f"\n\n{'='*80}")
|
||||
print(f" 📋 分词结果对照表")
|
||||
print(f"{'='*80}")
|
||||
|
||||
# Per-sentence comparison
|
||||
for cat, text, note in TEST_CASES:
|
||||
print(f"\n [{cat}] {text}")
|
||||
print(f" 关注: {note}")
|
||||
for name in results:
|
||||
if results[name].get("status") != "ok":
|
||||
continue
|
||||
words = results[name]["seg_outputs"].get(text, [])
|
||||
print(f" {name:<6s}: {' | '.join(words)}")
|
||||
|
||||
# Performance comparison
|
||||
print(f"\n\n{'='*80}")
|
||||
print(f" 📊 性能对比")
|
||||
print(f"{'='*80}")
|
||||
print(f" {'库':<10s} {'冷启动(s)':>10s} {'速度(字/s)':>14s} {'内存增量(MB)':>14s} {'状态':>8s}")
|
||||
print(f" {'─'*10} {'─'*10} {'─'*14} {'─'*14} {'─'*8}")
|
||||
|
||||
for name in ["jieba", "LTP", "HanLP"]:
|
||||
r = results.get(name, {})
|
||||
if r.get("status") != "ok":
|
||||
print(f" {name:<10s} {'─':>10s} {'─':>14s} {'─':>14s} {'失败':>8s}")
|
||||
continue
|
||||
print(
|
||||
f" {name:<10s} {r['cold_start_s']:>10.2f} {r['chars_per_sec']:>14,} "
|
||||
f"{r['memory_mb']:>14} {'OK':>8s}"
|
||||
)
|
||||
|
||||
# LAC note
|
||||
print(f"\n ℹ️ LAC 不兼容 Python 3.12（依赖已移除的 paddle.fluid 模块），无法参与对比。")
|
||||
|
||||
# Summary
|
||||
print(f"\n{'='*80}")
|
||||
print(f" 🏆 结论")
|
||||
print(f"{'='*80}")
|
||||
|
||||
for name in ["jieba", "LTP", "HanLP"]:
|
||||
r = results.get(name, {})
|
||||
if r.get("status") != "ok":
|
||||
continue
|
||||
desc = r.get("description", name)
|
||||
print(f" {name}: {desc}")
|
||||
|
||||
|
||||
# ── Main ────────────────────────────────────────────────────────────────
|
||||
|
||||
def main():
|
||||
print("=" * 60)
|
||||
print(" 中文分词库基准测试")
|
||||
print(" jieba(基线) vs LTP vs HanLP")
|
||||
print("=" * 60)
|
||||
|
||||
if "HF_ENDPOINT" in os.environ:
|
||||
print(f" HF_ENDPOINT={os.environ['HF_ENDPOINT']}")
|
||||
|
||||
results = run_benchmark()
|
||||
print_comparison(results)
|
||||
|
||||
# Save JSON
|
||||
out = {k: {kk: vv for kk, vv in v.items() if kk != "seg_outputs"} for k, v in results.items()}
|
||||
json_path = os.path.join(os.path.dirname(__file__), "..", "benchmark_segmentation.json")
|
||||
with open(json_path, "w", encoding="utf-8") as f:
|
||||
json.dump(out, f, ensure_ascii=False, indent=2)
|
||||
print(f"\n 💾 结果已保存到 {json_path}")
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
main()
|
||||
|
|
@ -128,7 +128,7 @@ def main():
|
|||
n_layers = 4
|
||||
n_heads = 4
|
||||
num_experts = 10
|
||||
max_seq_len = 128
|
||||
max_seq_len = 88
|
||||
|
||||
# ================================================================
|
||||
# 打印配置
|
||||
|
|
|
|||
|
|
@ -240,12 +240,21 @@ class CrossAttentionFusion(nn.Module):
|
|||
self.v_text_proj = nn.Linear(dim, dim, bias=False)
|
||||
self.k_pinyin_proj = nn.Linear(dim, dim, bias=False)
|
||||
self.v_pinyin_proj = nn.Linear(dim, dim, bias=False)
|
||||
self.k_part4_proj = nn.Linear(dim, dim, bias=False)
|
||||
self.v_part4_proj = nn.Linear(dim, dim, bias=False)
|
||||
|
||||
self.out_proj = nn.Linear(dim, dim, bias=False)
|
||||
self.ln = nn.LayerNorm(dim)
|
||||
|
||||
def forward(
|
||||
self, slots_S, context_H, pinyin_P, context_mask=None, pinyin_mask=None
|
||||
self,
|
||||
slots_S,
|
||||
context_H,
|
||||
pinyin_P,
|
||||
context_mask=None,
|
||||
pinyin_mask=None,
|
||||
part4_E=None,
|
||||
part4_mask=None,
|
||||
):
|
||||
"""
|
||||
Args:
|
||||
|
|
@ -254,6 +263,8 @@ class CrossAttentionFusion(nn.Module):
|
|||
pinyin_P: [batch, pinyin_len, dim] 拼音序列编码
|
||||
context_mask: [batch, ctx_len] 文本 padding mask (True for padding)
|
||||
pinyin_mask: [batch, pinyin_len] 拼音 padding mask (True for padding)
|
||||
part4_E: [batch, num_part4, dim] 候选词编码 (optional)
|
||||
part4_mask: [batch, num_part4] 候选词 padding mask (True for padding)
|
||||
Returns:
|
||||
fused: [batch, num_slots, dim]
|
||||
"""
|
||||
|
|
@ -306,6 +317,30 @@ class CrossAttentionFusion(nn.Module):
|
|||
|
||||
combined_attn = text_attn + pinyin_attn
|
||||
|
||||
if part4_E is not None:
|
||||
part4_len = part4_E.size(1)
|
||||
K_part4 = self.k_part4_proj(part4_E)
|
||||
V_part4 = self.v_part4_proj(part4_E)
|
||||
K_part4 = K_part4.view(
|
||||
batch_size, part4_len, self.n_heads, self.head_dim
|
||||
).transpose(1, 2)
|
||||
V_part4 = V_part4.view(
|
||||
batch_size, part4_len, self.n_heads, self.head_dim
|
||||
).transpose(1, 2)
|
||||
|
||||
part4_attn_mask = None
|
||||
if part4_mask is not None:
|
||||
part4_attn_mask = (
|
||||
part4_mask[:, None, None, :]
|
||||
.float()
|
||||
.masked_fill(part4_mask[:, None, None, :], -1e9)
|
||||
)
|
||||
|
||||
part4_attn = F.scaled_dot_product_attention(
|
||||
Q, K_part4, V_part4, attn_mask=part4_attn_mask
|
||||
)
|
||||
combined_attn = combined_attn + part4_attn
|
||||
|
||||
combined_attn = (
|
||||
combined_attn.transpose(1, 2)
|
||||
.contiguous()
|
||||
|
|
|
|||
|
|
@ -0,0 +1,6 @@
|
|||
from .pinyin_index import PinyinIndexBuilder, load_pinyin_index
|
||||
|
||||
__all__ = [
|
||||
"PinyinIndexBuilder",
|
||||
"load_pinyin_index",
|
||||
]
|
||||
|
|
@ -0,0 +1,35 @@
|
|||
#!/usr/bin/env python3
|
||||
import typer
|
||||
from loguru import logger
|
||||
|
||||
from .pinyin_index import PinyinIndexBuilder
|
||||
|
||||
pinyin_index_app = typer.Typer(name="pinyin-index")
|
||||
|
||||
|
||||
@pinyin_index_app.command("build")
|
||||
def build_index(
|
||||
corpus: str = typer.Option(..., "--corpus", help="文本语料路径 (HuggingFace 格式)"),
|
||||
output: str = typer.Option("./banks/pinyin_index.msgpack", "--output", "-o", help="输出路径"),
|
||||
name_corpus: str = typer.Option(None, "--name-corpus", help="Chinese-Names-Corpus 仓库路径"),
|
||||
max_samples: int = typer.Option(10_000_000, "--max-samples", help="最大处理文本数"),
|
||||
max_seq_length: int = typer.Option(88, "--max-seq-length", help="最大序列长度"),
|
||||
top_k_similar: int = typer.Option(5, "--top-k-similar", help="每个拼音保留的相似拼音数"),
|
||||
min_similarity: float = typer.Option(0.4, "--min-similarity", help="最小相似度阈值"),
|
||||
merge_prob: float = typer.Option(0.5, "--merge-prob", help="相邻单字词随机合并概率"),
|
||||
min_word_freq: int = typer.Option(2, "--min-word-freq", help="最小词频过滤阈值"),
|
||||
):
|
||||
builder = PinyinIndexBuilder(
|
||||
max_seq_length=max_seq_length,
|
||||
top_k_similar=top_k_similar,
|
||||
min_similarity=min_similarity,
|
||||
merge_prob=merge_prob,
|
||||
min_word_freq=min_word_freq,
|
||||
)
|
||||
result = builder.build(
|
||||
corpus_path=corpus,
|
||||
output_path=output,
|
||||
max_samples=max_samples,
|
||||
name_corpus=name_corpus,
|
||||
)
|
||||
logger.info(f"拼音索引已保存: {result}")
|
||||
|
|
@ -0,0 +1,343 @@
|
|||
"""拼音索引构建器 — 替代原 build-bank + build-confusion 流水线。
|
||||
|
||||
词典来源:
|
||||
1. Chinese-Names-Corpus (142W 人名 + 5W 成语)
|
||||
2. 内嵌地名库 (省级/地级/省会/主要城市, ~600 条)
|
||||
3. 随机连接补位 (连续单字词合并)
|
||||
|
||||
用法:
|
||||
pinyin-index build \
|
||||
--corpus "path/to/hf_dataset" \
|
||||
--name-corpus /home/songsenand/DataSet/Chinese-Names-Corpus \
|
||||
--output ./banks/pinyin_index.msgpack \
|
||||
--max-samples 10000000
|
||||
"""
|
||||
|
||||
import random
|
||||
from collections import defaultdict
|
||||
from pathlib import Path
|
||||
from typing import Dict, List, Optional, Tuple
|
||||
|
||||
import msgpack
|
||||
from loguru import logger
|
||||
from pypinyin import lazy_pinyin
|
||||
from tqdm import tqdm
|
||||
|
||||
try:
|
||||
from Levenshtein import distance as levenshtein
|
||||
except ImportError:
|
||||
def levenshtein(a: str, b: str) -> int:
|
||||
m, n = len(a), len(b)
|
||||
dp = list(range(n + 1))
|
||||
for i in range(1, m + 1):
|
||||
prev = dp[0]
|
||||
dp[0] = i
|
||||
for j in range(1, n + 1):
|
||||
tmp = dp[j]
|
||||
dp[j] = min(prev + (0 if a[i - 1] == b[j - 1] else 1), dp[j - 1] + 1, dp[j] + 1)
|
||||
prev = tmp
|
||||
return dp[n]
|
||||
|
||||
# ── 内嵌地名库 ──────────────────────────────────────────────────────
|
||||
|
||||
_PLACES = [
|
||||
# 省级 (34)
|
||||
"北京", "天津", "上海", "重庆", "河北", "山西", "辽宁", "吉林", "黑龙江",
|
||||
"江苏", "浙江", "安徽", "福建", "江西", "山东", "河南", "湖北", "湖南",
|
||||
"广东", "海南", "四川", "贵州", "云南", "陕西", "甘肃", "青海", "台湾",
|
||||
"内蒙古", "广西", "西藏", "宁夏", "新疆", "香港", "澳门",
|
||||
# 省会/直辖市 (31)
|
||||
"石家庄", "太原", "沈阳", "长春", "哈尔滨", "南京", "杭州", "合肥",
|
||||
"福州", "南昌", "济南", "郑州", "武汉", "长沙", "广州", "南宁",
|
||||
"成都", "贵阳", "昆明", "拉萨", "西安", "兰州", "西宁", "银川",
|
||||
"乌鲁木齐", "呼和浩特",
|
||||
# 地级市/重要城市 (~280)
|
||||
"唐山", "秦皇岛", "邯郸", "邢台", "保定", "张家口", "承德", "沧州", "廊坊", "衡水",
|
||||
"大同", "阳泉", "长治", "晋城", "朔州", "晋中", "运城", "忻州", "临汾", "吕梁",
|
||||
"大连", "鞍山", "抚顺", "本溪", "丹东", "锦州", "营口", "阜新", "辽阳", "盘锦", "铁岭", "朝阳", "葫芦岛",
|
||||
"吉林市", "四平", "辽源", "通化", "白山", "松原", "白城", "延边",
|
||||
"齐齐哈尔", "鸡西", "鹤岗", "双鸭山", "大庆", "伊春", "佳木斯", "七台河", "牡丹江", "黑河", "绥化",
|
||||
"无锡", "徐州", "常州", "苏州", "南通", "连云港", "淮安", "盐城", "扬州", "镇江", "泰州", "宿迁",
|
||||
"宁波", "温州", "嘉兴", "湖州", "绍兴", "金华", "衢州", "舟山", "台州", "丽水",
|
||||
"芜湖", "蚌埠", "淮南", "马鞍山", "淮北", "铜陵", "安庆", "黄山", "滁州", "阜阳", "宿州", "六安", "亳州", "池州", "宣城",
|
||||
"厦门", "莆田", "三明", "泉州", "漳州", "南平", "龙岩", "宁德",
|
||||
"景德镇", "萍乡", "九江", "新余", "鹰潭", "赣州", "吉安", "宜春", "抚州", "上饶",
|
||||
"青岛", "淄博", "枣庄", "东营", "烟台", "潍坊", "济宁", "泰安", "威海", "日照", "临沂", "德州", "聊城", "滨州", "菏泽",
|
||||
"开封", "洛阳", "平顶山", "安阳", "鹤壁", "新乡", "焦作", "濮阳", "许昌", "漯河", "三门峡", "南阳", "商丘", "信阳", "周口", "驻马店",
|
||||
"黄石", "十堰", "宜昌", "襄阳", "鄂州", "荆门", "孝感", "荆州", "黄冈", "咸宁", "随州", "恩施",
|
||||
"株洲", "湘潭", "衡阳", "邵阳", "岳阳", "常德", "张家界", "益阳", "郴州", "永州", "怀化", "娄底", "湘西",
|
||||
"韶关", "深圳", "珠海", "汕头", "佛山", "江门", "湛江", "茂名", "肇庆", "惠州", "梅州", "汕尾", "河源", "阳江", "清远", "东莞", "中山", "潮州", "揭阳", "云浮",
|
||||
"柳州", "桂林", "梧州", "北海", "防城港", "钦州", "贵港", "玉林", "百色", "贺州", "河池", "来宾", "崇左",
|
||||
"自贡", "攀枝花", "泸州", "德阳", "绵阳", "广元", "遂宁", "内江", "乐山", "南充", "眉山", "宜宾", "广安", "达州", "雅安", "巴中", "资阳",
|
||||
"六盘水", "遵义", "安顺", "毕节", "铜仁", "黔西南", "黔东南", "黔南",
|
||||
"曲靖", "玉溪", "保山", "昭通", "丽江", "普洱", "临沧", "楚雄", "红河", "文山", "西双版纳", "大理", "德宏", "怒江", "迪庆",
|
||||
"铜川", "宝鸡", "咸阳", "渭南", "延安", "汉中", "榆林", "安康", "商洛",
|
||||
"嘉峪关", "金昌", "白银", "天水", "武威", "张掖", "平凉", "酒泉", "庆阳", "定西", "陇南", "临夏", "甘南",
|
||||
"海东", "海北", "黄南", "海南州", "果洛", "玉树", "海西",
|
||||
"石嘴山", "吴忠", "固原", "中卫",
|
||||
"克拉玛依", "吐鲁番", "哈密", "昌吉", "博尔塔拉", "巴音郭楞", "阿克苏", "克孜勒苏", "喀什", "和田", "伊犁", "塔城", "阿勒泰",
|
||||
# 常见地名补充
|
||||
"浦东", "滨海", "雄安", "郑东", "福田", "南山", "海淀", "朝阳区", "徐汇",
|
||||
"大兴", "顺义", "通州", "闵行", "宝山", "嘉定", "松江", "黄埔", "白云",
|
||||
"番禺", "花都", "龙岗", "宝安", "龙华", "坪山", "高新", "天府", "两江",
|
||||
"洞庭湖", "太湖", "鄱阳湖", "青海湖", "千岛湖", "泸沽湖",
|
||||
"黄山", "泰山", "华山", "衡山", "恒山", "嵩山", "峨眉山", "五台山", "九华山", "普陀山", "武当山",
|
||||
"长江", "黄河", "珠江", "淮河", "海河", "松花江", "辽河",
|
||||
"故宫", "天坛", "颐和园", "长城", "兵马俑", "西湖", "周庄", "凤凰",
|
||||
# 大学名称
|
||||
"清华", "北大", "复旦", "浙大", "南大", "上交", "中科大", "哈工大", "西安交大",
|
||||
"武汉大学", "中山大学", "四川大学", "山东大学", "厦门大学", "同济", "南开", "天津大学",
|
||||
"北师大", "华科", "东南大学", "人大", "北航", "北理工",
|
||||
]
|
||||
|
||||
# ── 从 Chinese-Names-Corpus 加载词典 ──────────────────────────────
|
||||
|
||||
def _load_txt_names(path: str, min_len: int = 2, max_len: int = 8, skip_lines: int = 4) -> set:
|
||||
"""从 txt 文件加载词条, 跳过头部注释行"""
|
||||
items = set()
|
||||
try:
|
||||
with open(path, encoding="utf-8-sig") as f:
|
||||
for i, line in enumerate(f):
|
||||
if i < skip_lines:
|
||||
continue
|
||||
w = line.strip()
|
||||
if w and min_len <= len(w) <= max_len:
|
||||
items.add(w)
|
||||
except FileNotFoundError:
|
||||
logger.warning(f"File not found: {path}")
|
||||
return items
|
||||
|
||||
|
||||
def _load_csv_names(path: str, col: int = 0, min_len: int = 2, max_len: int = 8, skip_lines: int = 4) -> set:
|
||||
"""从 csv 文件加载指定列词条"""
|
||||
items = set()
|
||||
try:
|
||||
with open(path, encoding="utf-8-sig") as f:
|
||||
for i, line in enumerate(f):
|
||||
if i < skip_lines:
|
||||
continue
|
||||
parts = line.strip().split(",")
|
||||
if len(parts) > col:
|
||||
w = parts[col].strip()
|
||||
if w and min_len <= len(w) <= max_len:
|
||||
items.add(w)
|
||||
except FileNotFoundError:
|
||||
logger.warning(f"File not found: {path}")
|
||||
return items
|
||||
|
||||
|
||||
def load_name_dict(corpus_path: Optional[str] = None) -> List[str]:
|
||||
"""加载中文姓名+地名+成语词典, 返回所有词条列表"""
|
||||
all_words = set(_PLACES)
|
||||
|
||||
if corpus_path:
|
||||
base = Path(corpus_path) / "Chinese_Names_Corpus"
|
||||
dict_base = Path(corpus_path) / "Chinese_Dict_Corpus"
|
||||
|
||||
# 120W 现代人名
|
||||
names = _load_txt_names(str(base / "Chinese_Names_Corpus（120W）.txt"))
|
||||
all_words.update(names)
|
||||
logger.info(f" Names (120W): {len(names):,}")
|
||||
|
||||
# 120W 性别标注人名 (第0列)
|
||||
names_g = _load_csv_names(str(base / "Chinese_Names_Corpus_Gender（120W）.txt"), col=0)
|
||||
all_words.update(names_g)
|
||||
logger.info(f" Names (Gender): {len(names_g):,}")
|
||||
|
||||
# 25W 古代人名
|
||||
ancient = _load_txt_names(str(base / "Ancient_Names_Corpus（25W）.txt"))
|
||||
all_words.update(ancient)
|
||||
logger.info(f" Names (Ancient): {len(ancient):,}")
|
||||
|
||||
# 5W 成语
|
||||
cy = _load_txt_names(str(dict_base / "ChengYu_Corpus（5W）.txt"), min_len=4, max_len=12)
|
||||
all_words.update(cy)
|
||||
logger.info(f" Chengyu: {len(cy):,}")
|
||||
else:
|
||||
# 无外部词典时使用内嵌地名
|
||||
logger.info(" No name corpus provided, using embedded places only")
|
||||
|
||||
result = sorted(all_words)
|
||||
logger.info(f" Places: {len(_PLACES):,}")
|
||||
logger.info(f" Total: {len(result):,} unique words")
|
||||
|
||||
return result
|
||||
|
||||
|
||||
# ── 拼音工具 ──────────────────────────────────────────────────────────
|
||||
|
||||
def _is_chinese(c: str) -> bool:
|
||||
return "\u4e00" <= c <= "\u9fff"
|
||||
|
||||
|
||||
def _word_to_pinyin(word: str) -> str:
|
||||
return "".join(lazy_pinyin(word))
|
||||
|
||||
|
||||
def _pinyin_similarity(a: str, b: str) -> float:
|
||||
if not a or not b:
|
||||
return 0.0
|
||||
d = levenshtein(a, b)
|
||||
return 1.0 - d / max(len(a), len(b))
|
||||
|
||||
|
||||
# ── PinyinIndexBuilder ─────────────────────────────────────────────────
|
||||
|
||||
class PinyinIndexBuilder:
|
||||
def __init__(
|
||||
self,
|
||||
max_seq_length: int = 88,
|
||||
top_k_similar: int = 5,
|
||||
min_similarity: float = 0.4,
|
||||
merge_prob: float = 0.5,
|
||||
min_word_freq: int = 2,
|
||||
):
|
||||
self.max_seq_length = max_seq_length
|
||||
self.top_k_similar = top_k_similar
|
||||
self.min_similarity = min_similarity
|
||||
self.merge_prob = merge_prob
|
||||
self.min_word_freq = min_word_freq
|
||||
|
||||
def build(
|
||||
self,
|
||||
corpus_path: str,
|
||||
output_path: str,
|
||||
max_samples: int = 10_000_000,
|
||||
name_corpus: Optional[str] = None,
|
||||
):
|
||||
import jieba
|
||||
from datasets import load_dataset
|
||||
|
||||
# 加载姓名+地名+成语词典
|
||||
logger.info("Loading name/place/chengyu dictionaries...")
|
||||
names = load_name_dict(name_corpus)
|
||||
for name in names:
|
||||
jieba.add_word(name)
|
||||
logger.info(f"Loaded {len(names):,} entries into jieba user dict")
|
||||
|
||||
# 流式读取语料
|
||||
dataset = load_dataset(corpus_path, split="train", streaming=True)
|
||||
word_freq: Dict[str, int] = defaultdict(int)
|
||||
word_pinyin_map: Dict[str, str] = {}
|
||||
n_total = 0
|
||||
merge_added = 0
|
||||
|
||||
logger.info(f"Streaming corpus (max {max_samples:,} samples)...")
|
||||
pbar = tqdm(total=max_samples, desc="Building index", unit="samples")
|
||||
|
||||
for sample in dataset:
|
||||
if n_total >= max_samples:
|
||||
break
|
||||
text = sample.get("text", "")
|
||||
if not text or len(text) < 10:
|
||||
continue
|
||||
|
||||
words = list(jieba.cut(text, HMM=False))
|
||||
|
||||
for w in words:
|
||||
w = w.strip()
|
||||
if not w or not any(_is_chinese(c) for c in w):
|
||||
continue
|
||||
word_freq[w] += 1
|
||||
if w not in word_pinyin_map:
|
||||
word_pinyin_map[w] = _word_to_pinyin(w)
|
||||
|
||||
# 连续单字中文词随机合并
|
||||
i = 0
|
||||
while i < len(words):
|
||||
if not (_is_chinese(words[i]) and len(words[i]) == 1):
|
||||
i += 1
|
||||
continue
|
||||
j = i + 1
|
||||
while j < len(words) and j - i < 4:
|
||||
if not (_is_chinese(words[j]) and len(words[j]) == 1):
|
||||
break
|
||||
j += 1
|
||||
run_len = j - i
|
||||
if run_len >= 2 and random.random() < self.merge_prob:
|
||||
max_merge = min(run_len, 4)
|
||||
merge_len = random.randint(2, max_merge)
|
||||
merged = "".join(words[i : i + merge_len])
|
||||
if merged not in word_freq:
|
||||
word_freq[merged] += 1
|
||||
word_pinyin_map[merged] = _word_to_pinyin(merged)
|
||||
merge_added += 1
|
||||
i = max(i + 1, j)
|
||||
|
||||
n_total += 1
|
||||
if n_total % 1000 == 0:
|
||||
pbar.update(1000)
|
||||
pbar.set_postfix(words=len(word_freq), merged=merge_added)
|
||||
|
||||
pbar.close()
|
||||
|
||||
words_to_keep = {w for w, f in word_freq.items() if f >= self.min_word_freq}
|
||||
logger.info(
|
||||
f"Corpus: {n_total:,} texts, {len(words_to_keep):,} unique words "
|
||||
f"(min_freq={self.min_word_freq}), merge_added={merge_added:,}"
|
||||
)
|
||||
|
||||
pinyin_to_words: Dict[str, List[str]] = defaultdict(list)
|
||||
for w in words_to_keep:
|
||||
py = word_pinyin_map[w]
|
||||
pinyin_to_words[py].append(w)
|
||||
|
||||
unique_pinyins = list(pinyin_to_words.keys())
|
||||
logger.info(f"Computing pinyin similarity for {len(unique_pinyins):,} pinyin strings...")
|
||||
|
||||
buckets: Dict[int, List[Tuple[str, List[str]]]] = defaultdict(list)
|
||||
for py in unique_pinyins:
|
||||
buckets[len(py)].append((py, pinyin_to_words[py]))
|
||||
|
||||
similar_pinyin: Dict[str, List[str]] = {}
|
||||
bucket_keys = sorted(buckets.keys())
|
||||
|
||||
for k in tqdm(bucket_keys, desc="Pinyin similarity", unit="buckets"):
|
||||
current = buckets[k]
|
||||
candidate_keys = [k2 for k2 in bucket_keys if abs(k2 - k) <= 1]
|
||||
candidates = []
|
||||
for ck in candidate_keys:
|
||||
candidates.extend(buckets[ck])
|
||||
for py, _ in current:
|
||||
sims = []
|
||||
for cand_py, _ in candidates:
|
||||
if cand_py == py:
|
||||
continue
|
||||
sim = _pinyin_similarity(py, cand_py)
|
||||
if sim >= self.min_similarity:
|
||||
sims.append((cand_py, sim))
|
||||
sims.sort(key=lambda x: x[1], reverse=True)
|
||||
similar_pinyin[py] = [s[0] for s in sims[:self.top_k_similar]]
|
||||
|
||||
output = {
|
||||
"pinyin_to_words": {k: v for k, v in pinyin_to_words.items()},
|
||||
"similar_pinyin": similar_pinyin,
|
||||
"word_to_pinyin": word_pinyin_map,
|
||||
"word_freq": {k: v for k, v in word_freq.items() if k in words_to_keep},
|
||||
"meta": {
|
||||
"num_texts": n_total,
|
||||
"num_words": len(words_to_keep),
|
||||
"num_pinyins": len(unique_pinyins),
|
||||
"merge_added": merge_added,
|
||||
"top_k_similar": self.top_k_similar,
|
||||
"min_similarity": self.min_similarity,
|
||||
},
|
||||
}
|
||||
|
||||
Path(output_path).parent.mkdir(parents=True, exist_ok=True)
|
||||
with open(output_path, "wb") as f:
|
||||
f.write(msgpack.packb(output))
|
||||
|
||||
total_entries = sum(len(v) for v in similar_pinyin.values())
|
||||
size_mb = Path(output_path).stat().st_size / (1024 * 1024)
|
||||
logger.info(
|
||||
f"Saved: {output_path} ({size_mb:.1f} MB), "
|
||||
f"{len(words_to_keep):,} words, {len(unique_pinyins):,} pinyins, "
|
||||
f"{total_entries:,} similarity edges"
|
||||
)
|
||||
return output_path
|
||||
|
||||
|
||||
def load_pinyin_index(path: str) -> Dict:
|
||||
with open(path, "rb") as f:
|
||||
return msgpack.unpackb(f.read(), raw=False)
|
||||
|
|
@ -63,7 +63,7 @@ class PinyinInputDataset(IterableDataset):
|
|||
data_path: str,
|
||||
max_workers: int = -1,
|
||||
max_iter_length=1e6,
|
||||
max_seq_length=128,
|
||||
max_seq_length=88,
|
||||
text_field: str = "text",
|
||||
py_style_weight=(9, 2, 1),
|
||||
shuffle_buffer_size: int = 100000,
|
||||
|
|
@ -76,6 +76,10 @@ class PinyinInputDataset(IterableDataset):
|
|||
high_freq_repeat: float = 0.1,
|
||||
data_kwargs: Optional[Dict] = None,
|
||||
target_labels: Optional[Set[int]] = None,
|
||||
pinyin_index: Optional[Dict] = None,
|
||||
confusion_max_freq: int = 100,
|
||||
part4_ratio: float = 0.6,
|
||||
part4_target_remove: float = 0.5,
|
||||
):
|
||||
# 频率调整参数 - 幂律平滑方案
|
||||
self.min_freq = 109
|
||||
|
|
@ -91,6 +95,11 @@ class PinyinInputDataset(IterableDataset):
|
|||
self.data_kwargs = data_kwargs or {}
|
||||
self.target_labels = target_labels
|
||||
|
||||
self.pinyin_index = pinyin_index
|
||||
self.confusion_max_freq = confusion_max_freq
|
||||
self.part4_ratio = part4_ratio
|
||||
self.part4_target_remove = part4_target_remove
|
||||
|
||||
jieba.initialize()
|
||||
|
||||
self.tokenizer = AutoTokenizer.from_pretrained(
|
||||
|
|
@ -245,6 +254,17 @@ class PinyinInputDataset(IterableDataset):
|
|||
pinyin_ids = pinyin_ids[:24]
|
||||
return torch.tensor(pinyin_ids, dtype=torch.long)
|
||||
|
||||
def _build_part4_ids(self, part4_str: str, max_words: int = 5, max_word_len: int = 6) -> torch.Tensor:
|
||||
part4_ids = torch.zeros(max_words, max_word_len, dtype=torch.long)
|
||||
if not part4_str:
|
||||
return part4_ids
|
||||
words = part4_str.split("|")
|
||||
for i, w in enumerate(words[:max_words]):
|
||||
tokens = self.tokenizer.encode(w, add_special_tokens=False)[:max_word_len]
|
||||
for j, t in enumerate(tokens):
|
||||
part4_ids[i, j] = t
|
||||
return part4_ids
|
||||
|
||||
def _build_single_sample(
|
||||
self,
|
||||
label: int,
|
||||
|
|
@ -255,34 +275,51 @@ class PinyinInputDataset(IterableDataset):
|
|||
part2: str,
|
||||
pinyin_ids: torch.Tensor,
|
||||
words: list,
|
||||
target_word: str = "",
|
||||
word_pinyin: str = "",
|
||||
) -> dict:
|
||||
"""构造单条样本，每次调用都会重新随机采样上下文"""
|
||||
|
||||
# part1 长度：高斯分布 N(36, 6^2)，截断 [0, min(48, word_start)]
|
||||
part1_len = min(max(int(random.gauss(36, 6)), 0), 48, word_start)
|
||||
part1 = text[word_start - part1_len : word_start]
|
||||
|
||||
# part3：每次重新 roll
|
||||
part3 = ""
|
||||
if random.random() > 0.7:
|
||||
part3 = text[word_end : word_end + random.randint(1, 16)]
|
||||
|
||||
# part4：每次重新 roll
|
||||
part4 = ""
|
||||
if random.random() > 0.7 and words:
|
||||
if self.pinyin_index is not None and target_word and word_pinyin:
|
||||
freq = self.pinyin_index.get("word_freq", {}).get(target_word, 0)
|
||||
if freq <= self.confusion_max_freq:
|
||||
if random.random() < self.part4_ratio:
|
||||
candidates = []
|
||||
seen = {target_word}
|
||||
pinyin_to_words = self.pinyin_index.get("pinyin_to_words", {})
|
||||
similar_pinyin = self.pinyin_index.get("similar_pinyin", {})
|
||||
for py in similar_pinyin.get(word_pinyin, [word_pinyin]):
|
||||
for w in pinyin_to_words.get(py, []):
|
||||
if w not in seen:
|
||||
candidates.append(w)
|
||||
seen.add(w)
|
||||
if target_word in candidates and random.random() < self.part4_target_remove:
|
||||
candidates.remove(target_word)
|
||||
if candidates:
|
||||
num = random.randint(1, min(4, len(candidates)))
|
||||
selected = random.sample(candidates, num)
|
||||
part4 = "|".join(selected)
|
||||
elif random.random() > 0.7 and words:
|
||||
num_words = random.randint(1, 3)
|
||||
selected_words = random.sample(words, min(num_words, len(words)))
|
||||
part4 = "|".join(selected_words)
|
||||
|
||||
part4_ids = self._build_part4_ids(part4)
|
||||
|
||||
encoded = self.tokenizer(
|
||||
f"{part4}|{part1}",
|
||||
part3,
|
||||
part1 if part1 else "",
|
||||
part3 if part3 else None,
|
||||
max_length=self.max_seq_length,
|
||||
truncation=True,
|
||||
return_token_type_ids=True,
|
||||
)
|
||||
|
||||
# 确保 history 长度为 8
|
||||
hist = list(history)
|
||||
if len(hist) > 8:
|
||||
hist = hist[-8:]
|
||||
|
|
@ -299,6 +336,9 @@ class PinyinInputDataset(IterableDataset):
|
|||
"suffix": part3,
|
||||
"pinyin": part2,
|
||||
"pinyin_ids": pinyin_ids,
|
||||
"target_word": target_word,
|
||||
"word_pinyin": word_pinyin,
|
||||
"part4_ids": part4_ids,
|
||||
}
|
||||
|
||||
def __iter__(self):
|
||||
|
|
@ -448,7 +488,13 @@ class PinyinInputDataset(IterableDataset):
|
|||
if not should_break and random.random() <= 0.1:
|
||||
labels.append(0)
|
||||
|
||||
# 逐个 label 处理，削峰填谷前置，每次重复重新采样上下文
|
||||
target_word = text[word_start:word_end]
|
||||
word_pinyin = "".join(
|
||||
pinyin_list[i]
|
||||
for i in range(word_start, word_end)
|
||||
if self.query_engine.is_chinese_char(text[i])
|
||||
)
|
||||
|
||||
processed_history = []
|
||||
for label_idx, label in enumerate(labels):
|
||||
base_repeats = self.adjust_frequency(
|
||||
|
|
@ -481,6 +527,8 @@ class PinyinInputDataset(IterableDataset):
|
|||
part2=part2,
|
||||
pinyin_ids=pinyin_ids,
|
||||
words=words,
|
||||
target_word=target_word,
|
||||
word_pinyin=word_pinyin,
|
||||
)
|
||||
batch_samples.append(sample)
|
||||
|
||||
|
|
@ -549,6 +597,12 @@ class PinyinInputDataset(IterableDataset):
|
|||
# 逐个 label 处理，削峰填谷前置，每次重复重新采样上下文
|
||||
cont_processed_history = []
|
||||
cont_end = cont_positions[-1] + 1
|
||||
cont_target_word = text[cont_start:cont_end]
|
||||
cont_word_pinyin = "".join(
|
||||
pinyin_list[i]
|
||||
for i in range(cont_start, cont_end)
|
||||
if self.query_engine.is_chinese_char(text[i])
|
||||
)
|
||||
for label_idx, label in enumerate(cont_labels):
|
||||
base_repeats = self.adjust_frequency(
|
||||
self.sample_freqs.get(label, 0)
|
||||
|
|
@ -580,6 +634,8 @@ class PinyinInputDataset(IterableDataset):
|
|||
part2=part2_cont,
|
||||
pinyin_ids=pinyin_ids_cont,
|
||||
words=words,
|
||||
target_word=cont_target_word,
|
||||
word_pinyin=cont_word_pinyin,
|
||||
)
|
||||
batch_samples.append(sample)
|
||||
|
||||
|
|
|
|||
|
|
@ -144,7 +144,7 @@ class ContextEncoderExport(nn.Module):
|
|||
class DecoderExport(nn.Module):
|
||||
"""
|
||||
解码器导出模型
|
||||
输入: context_H, pinyin_P, history_slot_ids, context_mask, pinyin_mask
|
||||
输入: context_H, pinyin_P, history_slot_ids, context_mask, pinyin_mask, part4_ids
|
||||
输出: logits
|
||||
"""
|
||||
|
||||
|
|
@ -155,6 +155,9 @@ class DecoderExport(nn.Module):
|
|||
moe: MoELayer,
|
||||
slot_attention: nn.Module,
|
||||
classifier: nn.Module,
|
||||
text_emb: nn.Module,
|
||||
part4_proj: nn.Module,
|
||||
part4_ln: nn.Module,
|
||||
num_slots: int = 8,
|
||||
dim: int = 512,
|
||||
):
|
||||
|
|
@ -164,10 +167,23 @@ class DecoderExport(nn.Module):
|
|||
self.moe = moe
|
||||
self.slot_attention = slot_attention
|
||||
self.classifier = classifier
|
||||
self.text_emb = text_emb
|
||||
self.part4_proj = part4_proj
|
||||
self.part4_ln = part4_ln
|
||||
self.num_slots = num_slots
|
||||
self.dim = dim
|
||||
|
||||
def forward(self, context_H, pinyin_P, history_slot_ids, context_mask, pinyin_mask):
|
||||
def _encode_part4(self, part4_ids):
|
||||
B, N, L = part4_ids.shape
|
||||
flat = part4_ids.view(B * N, L)
|
||||
emb = self.text_emb(flat)
|
||||
valid_mask = (flat != 0).unsqueeze(-1).float()
|
||||
pooled = (emb * valid_mask).sum(dim=1) / valid_mask.sum(dim=1).clamp(min=1)
|
||||
pooled = self.part4_ln(self.part4_proj(pooled))
|
||||
return pooled.view(B, N, self.dim)
|
||||
|
||||
def forward(self, context_H, pinyin_P, history_slot_ids, context_mask, pinyin_mask,
|
||||
part4_ids=None):
|
||||
"""
|
||||
Args:
|
||||
context_H: [batch_size, seq_len, dim]
|
||||
|
|
@ -175,20 +191,23 @@ class DecoderExport(nn.Module):
|
|||
history_slot_ids: [batch_size, num_slots]
|
||||
context_mask: [batch_size, seq_len] (int32, 1表示padding)
|
||||
pinyin_mask: [batch_size, pinyin_len] (int32, 1表示padding)
|
||||
part4_ids: [batch_size, num_words, max_word_len] (optional)
|
||||
|
||||
Returns:
|
||||
logits: [batch_size, vocab_size]
|
||||
"""
|
||||
batch_size = context_H.size(0)
|
||||
|
||||
# 确保history_slot_ids形状正确
|
||||
history_slot_ids = history_slot_ids.view(batch_size, self.num_slots)
|
||||
|
||||
# 1. 槽位记忆
|
||||
S = self.slot_memory(history_slot_ids) # [batch_size, num_slots, dim]
|
||||
S = self.slot_memory(history_slot_ids)
|
||||
|
||||
part4_E = None
|
||||
part4_mask = None
|
||||
if part4_ids is not None and part4_ids.size(1) > 0:
|
||||
part4_E = self._encode_part4(part4_ids)
|
||||
part4_mask = (part4_ids.sum(dim=-1) == 0).to(torch.bool)
|
||||
|
||||
# 2. 交叉注意力融合
|
||||
# 转换mask: int32 -> bool (ONNX中需要bool类型)
|
||||
context_mask_bool = context_mask.to(torch.bool)
|
||||
pinyin_mask_bool = pinyin_mask.to(torch.bool)
|
||||
|
||||
|
|
@ -198,20 +217,17 @@ class DecoderExport(nn.Module):
|
|||
pinyin_P,
|
||||
context_mask=context_mask_bool,
|
||||
pinyin_mask=pinyin_mask_bool,
|
||||
) # [batch_size, num_slots, dim]
|
||||
part4_E=part4_E,
|
||||
part4_mask=part4_mask,
|
||||
)
|
||||
|
||||
# 3. MoE层
|
||||
moe_out = self.moe(fused) # [batch_size, num_slots, dim]
|
||||
moe_out = self.moe(fused)
|
||||
|
||||
# 4. 槽位注意力池化
|
||||
slot_scores = self.slot_attention(moe_out).squeeze(
|
||||
-1
|
||||
) # [batch_size, num_slots]
|
||||
slot_weights = torch.softmax(slot_scores, dim=1) # [batch_size, num_slots]
|
||||
pooled = (moe_out * slot_weights.unsqueeze(-1)).sum(dim=1) # [batch_size, dim]
|
||||
slot_scores = self.slot_attention(moe_out).squeeze(-1)
|
||||
slot_weights = torch.softmax(slot_scores, dim=1)
|
||||
pooled = (moe_out * slot_weights.unsqueeze(-1)).sum(dim=1)
|
||||
|
||||
# 5. 分类头
|
||||
logits = self.classifier(pooled) # [batch_size, vocab_size]
|
||||
logits = self.classifier(pooled)
|
||||
|
||||
return logits
|
||||
|
||||
|
|
@ -245,7 +261,7 @@ def create_export_models_from_checkpoint(checkpoint_path, device="cpu"):
|
|||
"n_layers": 4,
|
||||
"n_heads": 4,
|
||||
"num_experts": 10,
|
||||
"max_seq_len": 128,
|
||||
"max_seq_len": 88,
|
||||
}
|
||||
|
||||
# 创建原始模型
|
||||
|
|
@ -280,6 +296,9 @@ def create_export_models_from_checkpoint(checkpoint_path, device="cpu"):
|
|||
moe=model.moe,
|
||||
slot_attention=model.slot_attention,
|
||||
classifier=model.classifier,
|
||||
text_emb=model.context_encoder.text_emb,
|
||||
part4_proj=model.part4_proj,
|
||||
part4_ln=model.part4_ln,
|
||||
num_slots=model.num_slots,
|
||||
dim=model.dim,
|
||||
)
|
||||
|
|
|
|||
|
|
@ -39,17 +39,17 @@ class InputMethodEngine(nn.Module):
|
|||
n_layers: int = 4,
|
||||
n_heads: int = 4,
|
||||
num_experts: int = 10,
|
||||
max_seq_len: int = 128,
|
||||
max_seq_len: int = 88,
|
||||
compile: bool = False,
|
||||
moe_mode: str = "all", # "all" / "sparse" / "sparse_allow_graph"
|
||||
freeze_encoder: bool = False,
|
||||
):
|
||||
super().__init__()
|
||||
self.dim = dim
|
||||
self.num_slots = num_slots
|
||||
self.vocab_size = vocab_size
|
||||
self.max_seq_len = max_seq_len
|
||||
|
||||
# 1. 上下文编码器 (ContextEncoder)
|
||||
# 若 use_pinyin=False，则传入 pinyin_vocab_size=1 并固定嵌入为零
|
||||
self.context_encoder = ContextEncoder(
|
||||
vocab_size=vocab_size,
|
||||
pinyin_vocab_size=pinyin_vocab_size,
|
||||
|
|
@ -59,20 +59,15 @@ class InputMethodEngine(nn.Module):
|
|||
max_len=max_seq_len,
|
||||
)
|
||||
|
||||
# 2. 槽位记忆模块 (SlotMemory)
|
||||
# 适配历史槽位数量为 num_slots（每个槽位对应一个词，而非多步）
|
||||
self.slot_memory = SlotMemory(
|
||||
vocab_size=vocab_size,
|
||||
max_slots=num_slots,
|
||||
steps_per_slot=1, # 每个槽位只占一步
|
||||
steps_per_slot=1,
|
||||
dim=dim,
|
||||
)
|
||||
|
||||
# 3. 交叉注意力融合 (CrossAttentionFusion)
|
||||
# 使用 F.scaled_dot_product_attention 实现的版本
|
||||
self.cross_attn = CrossAttentionFusion(dim=dim, n_heads=n_heads)
|
||||
|
||||
# 4. 混合专家层 (MoE)
|
||||
self.moe = MoELayer(
|
||||
dim=dim,
|
||||
num_experts=num_experts,
|
||||
|
|
@ -81,12 +76,17 @@ class InputMethodEngine(nn.Module):
|
|||
moe_mode=moe_mode,
|
||||
)
|
||||
|
||||
# 5. 槽位注意力池化
|
||||
self.slot_attention = nn.Linear(dim, 1)
|
||||
|
||||
# 6. 分类头
|
||||
self.classifier = nn.Linear(dim, vocab_size)
|
||||
|
||||
self.part4_proj = nn.Linear(dim, dim, bias=False)
|
||||
self.part4_ln = nn.LayerNorm(dim)
|
||||
|
||||
if freeze_encoder:
|
||||
for p in self.context_encoder.parameters():
|
||||
p.requires_grad_(False)
|
||||
|
||||
if compile:
|
||||
self.forward = torch.compile(
|
||||
self.forward,
|
||||
|
|
@ -102,6 +102,7 @@ class InputMethodEngine(nn.Module):
|
|||
attention_mask: torch.Tensor,
|
||||
pinyin_ids: torch.Tensor,
|
||||
history_slot_ids: torch.Tensor,
|
||||
part4_ids: Optional[torch.Tensor] = None,
|
||||
) -> torch.Tensor:
|
||||
"""
|
||||
前向传播。
|
||||
|
|
@ -116,12 +117,24 @@ class InputMethodEngine(nn.Module):
|
|||
|
||||
H, P = self.context_encoder(input_ids, pinyin_ids, mask=attention_mask)
|
||||
|
||||
part4_E = None
|
||||
part4_mask = None
|
||||
if part4_ids is not None and part4_ids.size(1) > 0:
|
||||
part4_E = self._encode_part4(part4_ids)
|
||||
part4_mask = part4_ids.sum(dim=-1) == 0
|
||||
|
||||
S = self.slot_memory(history_slot_ids)
|
||||
|
||||
context_mask = attention_mask == 0
|
||||
pinyin_mask = pinyin_ids == 0
|
||||
fused = self.cross_attn(
|
||||
S, H, P, context_mask=context_mask, pinyin_mask=pinyin_mask
|
||||
S,
|
||||
H,
|
||||
P,
|
||||
context_mask=context_mask,
|
||||
pinyin_mask=pinyin_mask,
|
||||
part4_E=part4_E,
|
||||
part4_mask=part4_mask,
|
||||
)
|
||||
|
||||
fused = fused * slot_mask.unsqueeze(-1)
|
||||
|
|
@ -137,3 +150,12 @@ class InputMethodEngine(nn.Module):
|
|||
|
||||
logits = self.classifier(pooled)
|
||||
return logits
|
||||
|
||||
def _encode_part4(self, part4_ids: torch.Tensor) -> torch.Tensor:
|
||||
B, N, L = part4_ids.shape
|
||||
flat = part4_ids.view(B * N, L)
|
||||
emb = self.context_encoder.text_emb(flat)
|
||||
valid_mask = (flat != 0).unsqueeze(-1).float()
|
||||
pooled = (emb * valid_mask).sum(dim=1) / valid_mask.sum(dim=1).clamp(min=1)
|
||||
pooled = self.part4_ln(self.part4_proj(pooled))
|
||||
return pooled.view(B, N, self.dim)
|
||||
|
|
|
|||
|
|
@ -101,6 +101,9 @@ def export_decoder(
|
|||
history_slot_ids = torch.randint(
|
||||
0, 100, (batch_size, num_slots), dtype=torch.long
|
||||
)
|
||||
part4_ids = torch.randint(
|
||||
1, 1000, (batch_size, 5, 6), dtype=torch.long
|
||||
)
|
||||
else:
|
||||
context_H = torch.randn(batch_size, seq_len, dim, dtype=torch.float32)
|
||||
pinyin_P = torch.randn(batch_size, pinyin_len, dim, dtype=torch.float32)
|
||||
|
|
@ -109,6 +112,9 @@ def export_decoder(
|
|||
history_slot_ids = torch.randint(
|
||||
0, 100, (batch_size, num_slots), dtype=torch.long
|
||||
)
|
||||
part4_ids = torch.randint(
|
||||
1, 1000, (batch_size, 5, 6), dtype=torch.long
|
||||
)
|
||||
|
||||
dynamic_axes = {
|
||||
"context_H": {0: "batch_size", 1: "seq_len"},
|
||||
|
|
@ -116,12 +122,13 @@ def export_decoder(
|
|||
"history_slot_ids": {0: "batch_size"},
|
||||
"context_mask": {0: "batch_size", 1: "seq_len"},
|
||||
"pinyin_mask": {0: "batch_size"},
|
||||
"part4_ids": {0: "batch_size", 1: "num_words"},
|
||||
"logits": {0: "batch_size"},
|
||||
}
|
||||
|
||||
torch.onnx.export(
|
||||
model,
|
||||
(context_H, pinyin_P, history_slot_ids, context_mask, pinyin_mask),
|
||||
(context_H, pinyin_P, history_slot_ids, context_mask, pinyin_mask, part4_ids),
|
||||
output_path,
|
||||
input_names=[
|
||||
"context_H",
|
||||
|
|
@ -129,6 +136,7 @@ def export_decoder(
|
|||
"history_slot_ids",
|
||||
"context_mask",
|
||||
"pinyin_mask",
|
||||
"part4_ids",
|
||||
],
|
||||
output_names=["logits"],
|
||||
dynamic_axes=dynamic_axes,
|
||||
|
|
@ -147,7 +155,7 @@ def export_decoder(
|
|||
except Exception as e:
|
||||
print(f" ONNX 模型验证警告: {e}")
|
||||
|
||||
return context_H, pinyin_P, history_slot_ids, context_mask, pinyin_mask
|
||||
return context_H, pinyin_P, history_slot_ids, context_mask, pinyin_mask, part4_ids
|
||||
|
||||
|
||||
def save_example_inputs(output_dir: str, example_inputs_dict: Dict) -> None:
|
||||
|
|
|
|||
|
|
@ -43,6 +43,7 @@ from torch.utils.data import DataLoader
|
|||
from tqdm import tqdm
|
||||
|
||||
from .dataset import PinyinInputDataset
|
||||
from .confusion.pinyin_index import load_pinyin_index
|
||||
from .trainer import preprocess_collate_fn, worker_init_fn
|
||||
|
||||
FIELDS = [
|
||||
|
|
@ -53,12 +54,16 @@ FIELDS = [
|
|||
"history_slot_ids",
|
||||
"pinyin_ids",
|
||||
]
|
||||
FIELDS_PART4 = FIELDS + ["part4_ids"]
|
||||
|
||||
MAX_PART4_WORDS = 5
|
||||
MAX_WORD_LEN = 6
|
||||
|
||||
|
||||
def _extract_batch(batch: dict, take: int) -> Dict[str, np.ndarray]:
|
||||
def _extract_batch(batch: dict, take: int, active_fields: list) -> Dict[str, np.ndarray]:
|
||||
"""从 DataLoader batch 中提取指定数量的样本，转为 int16 numpy 数组"""
|
||||
result = {}
|
||||
for f in FIELDS:
|
||||
for f in active_fields:
|
||||
tensor = batch[f][:take]
|
||||
arr = tensor.numpy().astype(np.int16)
|
||||
if f == "labels" and arr.ndim > 1 and arr.shape[-1] == 1:
|
||||
|
|
@ -72,8 +77,9 @@ def collect_samples(
|
|||
num_samples: int,
|
||||
output_dir: Path,
|
||||
split_name: str,
|
||||
max_seq_length: int = 128,
|
||||
max_seq_length: int = 88,
|
||||
shard_size: int = 5_000_000,
|
||||
has_part4: bool = False,
|
||||
) -> int:
|
||||
"""
|
||||
分片流式收集样本，每累积 shard_size 个样本保存为一个压缩 .npz 分片。
|
||||
|
|
@ -82,8 +88,9 @@ def collect_samples(
|
|||
"""
|
||||
split_dir = output_dir / split_name
|
||||
split_dir.mkdir(parents=True, exist_ok=True)
|
||||
active_fields = FIELDS_PART4 if has_part4 else FIELDS
|
||||
|
||||
shard_buffers: Dict[str, List[np.ndarray]] = {f: [] for f in FIELDS}
|
||||
shard_buffers: Dict[str, List[np.ndarray]] = {f: [] for f in active_fields}
|
||||
shard_count = 0
|
||||
shard_idx = 0
|
||||
total = 0
|
||||
|
|
@ -97,8 +104,8 @@ def collect_samples(
|
|||
break
|
||||
take = min(batch_size, remaining)
|
||||
|
||||
extracted = _extract_batch(batch, take)
|
||||
for f in FIELDS:
|
||||
extracted = _extract_batch(batch, take, active_fields)
|
||||
for f in active_fields:
|
||||
shard_buffers[f].append(extracted[f])
|
||||
|
||||
shard_count += take
|
||||
|
|
@ -107,13 +114,13 @@ def collect_samples(
|
|||
|
||||
if shard_count >= shard_size:
|
||||
merged = {}
|
||||
for f in FIELDS:
|
||||
for f in active_fields:
|
||||
merged[f] = np.concatenate(shard_buffers[f], axis=0)
|
||||
np.savez_compressed(split_dir / f"shard_{shard_idx:06d}.npz", **merged)
|
||||
logger.debug(f"Saved {split_name} shard {shard_idx}: {shard_count} samples")
|
||||
|
||||
shard_idx += 1
|
||||
shard_buffers = {f: [] for f in FIELDS}
|
||||
shard_buffers = {f: [] for f in active_fields}
|
||||
shard_count = 0
|
||||
del merged
|
||||
gc.collect()
|
||||
|
|
@ -124,7 +131,7 @@ def collect_samples(
|
|||
# 写入最后一个不满的分片
|
||||
if shard_count > 0:
|
||||
merged = {}
|
||||
for f in FIELDS:
|
||||
for f in active_fields:
|
||||
merged[f] = np.concatenate(shard_buffers[f], axis=0)
|
||||
np.savez_compressed(split_dir / f"shard_{shard_idx:06d}.npz", **merged)
|
||||
logger.debug(f"Saved {split_name} shard {shard_idx}: {shard_count} samples")
|
||||
|
|
@ -143,7 +150,7 @@ def collect_samples(
|
|||
"num_samples": actual_count,
|
||||
"max_seq_length": max_seq_length,
|
||||
"dtype": "int16",
|
||||
"fields": FIELDS,
|
||||
"fields": active_fields,
|
||||
"shard_size": shard_size,
|
||||
"num_shards": num_shards,
|
||||
"shard_sizes": shard_sizes,
|
||||
|
|
@ -189,8 +196,15 @@ def main():
|
|||
parser.add_argument(
|
||||
"--num-workers", type=int, default=2, help="DataLoader worker数量"
|
||||
)
|
||||
parser.add_argument("--max-seq-length", type=int, default=128, help="最大序列长度")
|
||||
parser.add_argument("--max-seq-length", type=int, default=88, help="最大序列长度")
|
||||
parser.add_argument("--seed", type=int, default=42, help="随机种子")
|
||||
parser.add_argument(
|
||||
"--pinyin-index", type=str, default=None, help="拼音索引文件路径（pinyin_index.msgpack），启用混淆 part4 生成"
|
||||
)
|
||||
parser.add_argument(
|
||||
"--confusion-max-freq", type=int, default=100,
|
||||
help="混淆频率阈值: 词频 ≤ N 才触发混淆，默认 100"
|
||||
)
|
||||
parser.add_argument(
|
||||
"--shard-size",
|
||||
type=int,
|
||||
|
|
@ -227,6 +241,13 @@ def main():
|
|||
for k, v in (item.split(":") for item in args.length_weights.split(","))
|
||||
}
|
||||
|
||||
pinyin_index = None
|
||||
if args.pinyin_index:
|
||||
pinyin_index = load_pinyin_index(args.pinyin_index)
|
||||
console.print(f"[bold cyan]拼音索引已加载: {args.pinyin_index}[/bold cyan]")
|
||||
console.print(f" 词数: {pinyin_index['meta'].get('num_words', '?'):,}")
|
||||
console.print(f" 混淆频率阈值: ≤{args.confusion_max_freq}")
|
||||
|
||||
output_dir = Path(args.output_dir)
|
||||
output_dir.mkdir(parents=True, exist_ok=True)
|
||||
|
||||
|
|
@ -257,6 +278,8 @@ def main():
|
|||
py_style_weight=py_style_weight,
|
||||
shuffle_buffer_size=100,
|
||||
length_weights=length_weights,
|
||||
pinyin_index=pinyin_index,
|
||||
confusion_max_freq=args.confusion_max_freq,
|
||||
)
|
||||
|
||||
train_dataloader = DataLoader(
|
||||
|
|
@ -294,6 +317,7 @@ def main():
|
|||
)
|
||||
|
||||
logger.info("开始收集训练数据...")
|
||||
has_p4 = pinyin_index is not None
|
||||
train_count = collect_samples(
|
||||
train_dataloader,
|
||||
args.num_train_samples,
|
||||
|
|
@ -301,6 +325,7 @@ def main():
|
|||
"train",
|
||||
args.max_seq_length,
|
||||
args.shard_size,
|
||||
has_part4=has_p4,
|
||||
)
|
||||
|
||||
if train_count < args.num_train_samples:
|
||||
|
|
@ -316,6 +341,7 @@ def main():
|
|||
"eval",
|
||||
args.max_seq_length,
|
||||
args.shard_size,
|
||||
has_part4=False,
|
||||
)
|
||||
|
||||
if eval_count < args.num_eval_samples:
|
||||
|
|
|
|||
|
|
@ -31,6 +31,8 @@ FIELDS = [
|
|||
"pinyin_ids",
|
||||
]
|
||||
|
||||
_PART4_FIELDS = FIELDS + ["part4_ids"]
|
||||
|
||||
|
||||
def _read_shard_size(npz_path: Path) -> int:
|
||||
with zipfile.ZipFile(npz_path, "r") as z:
|
||||
|
|
@ -190,6 +192,7 @@ class PreProcessedDataset(Dataset):
|
|||
shard_idx = lo - 1
|
||||
local_idx = idx - self._shard_offsets[shard_idx]
|
||||
shard_data = self._cache.get(shard_idx, self._load_shard)
|
||||
has_p4 = "part4_ids" in shard_data
|
||||
result = {
|
||||
"input_ids": torch.from_numpy(
|
||||
shard_data["input_ids"][local_idx].astype(np.int64)
|
||||
|
|
@ -210,6 +213,10 @@ class PreProcessedDataset(Dataset):
|
|||
shard_data["pinyin_ids"][local_idx].astype(np.int64)
|
||||
),
|
||||
}
|
||||
if has_p4:
|
||||
result["part4_ids"] = torch.from_numpy(
|
||||
shard_data["part4_ids"][local_idx].astype(np.int64)
|
||||
)
|
||||
else:
|
||||
result = {
|
||||
"input_ids": torch.from_numpy(self.input_ids[idx].astype(np.int64)),
|
||||
|
|
@ -233,7 +240,7 @@ def preprocessed_collate_fn(batch):
|
|||
预处理数据的 collate 函数。
|
||||
不含 string 字段（prefix/suffix/pinyin），仅处理 tensor 字段。
|
||||
"""
|
||||
return {
|
||||
result = {
|
||||
"input_ids": torch.stack([item["input_ids"] for item in batch]),
|
||||
"token_type_ids": torch.stack([item["token_type_ids"] for item in batch]),
|
||||
"attention_mask": torch.stack([item["attention_mask"] for item in batch]),
|
||||
|
|
@ -241,3 +248,6 @@ def preprocessed_collate_fn(batch):
|
|||
"history_slot_ids": torch.stack([item["history_slot_ids"] for item in batch]),
|
||||
"pinyin_ids": torch.stack([item["pinyin_ids"] for item in batch]),
|
||||
}
|
||||
if "part4_ids" in batch[0]:
|
||||
result["part4_ids"] = torch.stack([item["part4_ids"] for item in batch])
|
||||
return result
|
||||
|
|
|
|||
|
|
@ -265,6 +265,7 @@ class Trainer:
|
|||
attention_mask=attention_mask,
|
||||
pinyin_ids=pinyin_ids,
|
||||
history_slot_ids=history_slot_ids,
|
||||
part4_ids=batch.get("part4_ids"),
|
||||
)
|
||||
|
||||
# 计算损失
|
||||
|
|
@ -362,11 +363,20 @@ class Trainer:
|
|||
"scaler_state_dict": self.scaler.state_dict(),
|
||||
"best_eval_loss": self.best_eval_loss,
|
||||
"config": {
|
||||
"vocab_size": self.model.vocab_size,
|
||||
"pinyin_vocab_size": 30,
|
||||
"dim": self.model.dim,
|
||||
"num_slots": self.model.num_slots,
|
||||
"n_layers": self.model.context_encoder.transformer.num_layers,
|
||||
"n_heads": self.model.cross_attn.n_heads,
|
||||
"num_experts": self.model.moe.num_experts,
|
||||
"max_seq_len": self.model.max_seq_len,
|
||||
"learning_rate": self.learning_rate,
|
||||
"weight_decay": self.weight_decay,
|
||||
"warmup_ratio": self.warmup_ratio,
|
||||
"label_smoothing": self.label_smoothing,
|
||||
"total_steps": self.total_steps,
|
||||
"moe_mode": self.model.moe.moe_mode,
|
||||
},
|
||||
}
|
||||
|
||||
|
|
@ -436,11 +446,20 @@ class Trainer:
|
|||
"scaler_state_dict": self.scaler.state_dict(),
|
||||
"best_eval_loss": self.best_eval_loss,
|
||||
"config": {
|
||||
"vocab_size": self.model.vocab_size,
|
||||
"pinyin_vocab_size": 30,
|
||||
"dim": self.model.dim,
|
||||
"num_slots": self.model.num_slots,
|
||||
"n_layers": self.model.context_encoder.transformer.num_layers,
|
||||
"n_heads": self.model.cross_attn.n_heads,
|
||||
"num_experts": self.model.moe.num_experts,
|
||||
"max_seq_len": self.model.max_seq_len,
|
||||
"learning_rate": self.learning_rate,
|
||||
"weight_decay": self.weight_decay,
|
||||
"warmup_ratio": self.warmup_ratio,
|
||||
"label_smoothing": self.label_smoothing,
|
||||
"total_steps": self.total_steps,
|
||||
"moe_mode": self.model.moe.moe_mode,
|
||||
},
|
||||
}
|
||||
|
||||
|
|
@ -1089,9 +1108,32 @@ def collate_fn(batch: List[Dict[str, Any]], max_seq_length: int = 0) -> Dict[str
|
|||
pinyin_ids = torch.stack([item["pinyin_ids"] for item in batch])
|
||||
|
||||
prefixes = [item["prefix"] for item in batch]
|
||||
suffixes = [item["suffix"] for item in batch]
|
||||
suffixes = [item.get("suffix", "") for item in batch]
|
||||
pinyins = [item["pinyin"] for item in batch]
|
||||
|
||||
target_words = [item.get("target_word", "") for item in batch]
|
||||
word_pinyins = [item.get("word_pinyin", "") for item in batch]
|
||||
|
||||
part4_ids_list = [item.get("part4_ids", torch.zeros(0)) for item in batch]
|
||||
has_part4 = any(p.numel() > 0 and p.sum() > 0 for p in part4_ids_list)
|
||||
if has_part4:
|
||||
part4_n = max(p.shape[0] for p in part4_ids_list if p.numel() > 0)
|
||||
part4_l = max(p.shape[1] for p in part4_ids_list if p.numel() > 0)
|
||||
padded_part4_ids = []
|
||||
for pids in part4_ids_list:
|
||||
if pids.numel() == 0 or pids.shape[0] < part4_n:
|
||||
if pids.numel() == 0:
|
||||
pad = torch.zeros(part4_n, part4_l, dtype=torch.long)
|
||||
else:
|
||||
pw = torch.zeros(part4_n - pids.shape[0], pids.shape[1], dtype=torch.long)
|
||||
pad = torch.cat([pids, pw], dim=0)
|
||||
padded_part4_ids.append(pad)
|
||||
else:
|
||||
padded_part4_ids.append(pids[:part4_n])
|
||||
part4_ids = torch.stack(padded_part4_ids)
|
||||
else:
|
||||
part4_ids = None
|
||||
|
||||
return {
|
||||
"input_ids": input_ids,
|
||||
"token_type_ids": token_type_ids,
|
||||
|
|
@ -1102,6 +1144,9 @@ def collate_fn(batch: List[Dict[str, Any]], max_seq_length: int = 0) -> Dict[str
|
|||
"suffix": suffixes,
|
||||
"pinyin": pinyins,
|
||||
"pinyin_ids": pinyin_ids,
|
||||
"target_word": target_words,
|
||||
"word_pinyin": word_pinyins,
|
||||
"part4_ids": part4_ids,
|
||||
}
|
||||
|
||||
|
||||
|
|
@ -1143,7 +1188,7 @@ def create_dataloader(
|
|||
)
|
||||
|
||||
logger.info(f"📊 使用标准DataLoader，worker数量: {num_workers}")
|
||||
fixed_max_seq_length = getattr(dataset, "max_seq_length", 128)
|
||||
fixed_max_seq_length = getattr(dataset, "max_seq_length", 88)
|
||||
dataloader = DataLoader(
|
||||
dataset,
|
||||
batch_size=batch_size,
|
||||
|
|
@ -1222,6 +1267,11 @@ def train(
|
|||
"--moe-mode",
|
||||
help="MoE 计算策略: all（全量计算）, sparse（稀疏计算）, sparse_allow_graph（稀疏+allow_in_graph）",
|
||||
),
|
||||
freeze_encoder: bool = typer.Option(
|
||||
False,
|
||||
"--freeze-encoder/--no-freeze-encoder",
|
||||
help="是否冻结编码器（微调模式）",
|
||||
),
|
||||
):
|
||||
"""
|
||||
训练输入法模型
|
||||
|
|
@ -1245,7 +1295,7 @@ def train(
|
|||
n_layers = 4
|
||||
n_heads = 4
|
||||
num_experts = 10
|
||||
max_seq_len = 128
|
||||
max_seq_len = 88
|
||||
use_pinyin = True # 始终使用拼音
|
||||
|
||||
console = Console()
|
||||
|
|
@ -1436,6 +1486,7 @@ def train(
|
|||
max_seq_len=max_seq_len,
|
||||
compile=compile,
|
||||
moe_mode=moe_mode,
|
||||
freeze_encoder=freeze_encoder,
|
||||
)
|
||||
|
||||
console.print(
|
||||
|
|
|
|||
Loading…
Reference in New Issue