Add Streamlit preview frontend and normalize line endings

.gitattributes

@@ -1,35 +1,35 @@
-*.7z filter=lfs diff=lfs merge=lfs -text
-*.arrow filter=lfs diff=lfs merge=lfs -text
-*.bin filter=lfs diff=lfs merge=lfs -text
-*.bz2 filter=lfs diff=lfs merge=lfs -text
-*.ckpt filter=lfs diff=lfs merge=lfs -text
-*.ftz filter=lfs diff=lfs merge=lfs -text
-*.gz filter=lfs diff=lfs merge=lfs -text
-*.h5 filter=lfs diff=lfs merge=lfs -text
-*.joblib filter=lfs diff=lfs merge=lfs -text
-*.lfs.* filter=lfs diff=lfs merge=lfs -text
-*.mlmodel filter=lfs diff=lfs merge=lfs -text
-*.model filter=lfs diff=lfs merge=lfs -text
-*.msgpack filter=lfs diff=lfs merge=lfs -text
-*.npy filter=lfs diff=lfs merge=lfs -text
-*.npz filter=lfs diff=lfs merge=lfs -text
-*.onnx filter=lfs diff=lfs merge=lfs -text
-*.ot filter=lfs diff=lfs merge=lfs -text
-*.parquet filter=lfs diff=lfs merge=lfs -text
-*.pb filter=lfs diff=lfs merge=lfs -text
-*.pickle filter=lfs diff=lfs merge=lfs -text
-*.pkl filter=lfs diff=lfs merge=lfs -text
-*.pt filter=lfs diff=lfs merge=lfs -text
-*.pth filter=lfs diff=lfs merge=lfs -text
-*.rar filter=lfs diff=lfs merge=lfs -text
-*.safetensors filter=lfs diff=lfs merge=lfs -text
-saved_model/**/* filter=lfs diff=lfs merge=lfs -text
-*.tar.* filter=lfs diff=lfs merge=lfs -text
-*.tar filter=lfs diff=lfs merge=lfs -text
-*.tflite filter=lfs diff=lfs merge=lfs -text
-*.tgz filter=lfs diff=lfs merge=lfs -text
-*.wasm filter=lfs diff=lfs merge=lfs -text
-*.xz filter=lfs diff=lfs merge=lfs -text
-*.zip filter=lfs diff=lfs merge=lfs -text
-*.zst filter=lfs diff=lfs merge=lfs -text
-*tfevents* filter=lfs diff=lfs merge=lfs -text
+*.7z filter=lfs diff=lfs merge=lfs -text
+*.arrow filter=lfs diff=lfs merge=lfs -text
+*.bin filter=lfs diff=lfs merge=lfs -text
+*.bz2 filter=lfs diff=lfs merge=lfs -text
+*.ckpt filter=lfs diff=lfs merge=lfs -text
+*.ftz filter=lfs diff=lfs merge=lfs -text
+*.gz filter=lfs diff=lfs merge=lfs -text
+*.h5 filter=lfs diff=lfs merge=lfs -text
+*.joblib filter=lfs diff=lfs merge=lfs -text
+*.lfs.* filter=lfs diff=lfs merge=lfs -text
+*.mlmodel filter=lfs diff=lfs merge=lfs -text
+*.model filter=lfs diff=lfs merge=lfs -text
+*.msgpack filter=lfs diff=lfs merge=lfs -text
+*.npy filter=lfs diff=lfs merge=lfs -text
+*.npz filter=lfs diff=lfs merge=lfs -text
+*.onnx filter=lfs diff=lfs merge=lfs -text
+*.ot filter=lfs diff=lfs merge=lfs -text
+*.parquet filter=lfs diff=lfs merge=lfs -text
+*.pb filter=lfs diff=lfs merge=lfs -text
+*.pickle filter=lfs diff=lfs merge=lfs -text
+*.pkl filter=lfs diff=lfs merge=lfs -text
+*.pt filter=lfs diff=lfs merge=lfs -text
+*.pth filter=lfs diff=lfs merge=lfs -text
+*.rar filter=lfs diff=lfs merge=lfs -text
+*.safetensors filter=lfs diff=lfs merge=lfs -text
+saved_model/**/* filter=lfs diff=lfs merge=lfs -text
+*.tar.* filter=lfs diff=lfs merge=lfs -text
+*.tar filter=lfs diff=lfs merge=lfs -text
+*.tflite filter=lfs diff=lfs merge=lfs -text
+*.tgz filter=lfs diff=lfs merge=lfs -text
+*.wasm filter=lfs diff=lfs merge=lfs -text
+*.xz filter=lfs diff=lfs merge=lfs -text
+*.zip filter=lfs diff=lfs merge=lfs -text
+*.zst filter=lfs diff=lfs merge=lfs -text
+*tfevents* filter=lfs diff=lfs merge=lfs -text

.gitignore

@@ -1,57 +1,57 @@
-# Byte-compiled / optimized / DLL files
-__pycache__/
-*.py[cod]
-*$py.class
-
-# Distribution
-dist/
-build/
-*.egg-info/
-*.egg
-
-# Virtual environments
-.venv/
-venv/
-env/
-
-# IDE
-.idea/
-.vscode/
-*.swp
-*.swo
-
-# Data
-/data/
-*.tsv
-*.csv
-
-# Checkpoints & Outputs
-checkpoints/
-outputs/
-logs/
-wandb/
-
-# Test cache
-.pytest_cache/
-
-# Jupyter
-.ipynb_checkpoints/
-
-# Log files
-*.log
-
-# OS
-.DS_Store
-Thumbs.db
-
-# Models
-*.pt
-*.pth
-*.bin
-*.safetensors
-*.onnx
-
-# Tokenizer artifacts
-*.model
-*.vocab
-tokenizer.json
+# Byte-compiled / optimized / DLL files
+__pycache__/
+*.py[cod]
+*$py.class
+
+# Distribution
+dist/
+build/
+*.egg-info/
+*.egg
+
+# Virtual environments
+.venv/
+venv/
+env/
+
+# IDE
+.idea/
+.vscode/
+*.swp
+*.swo
+
+# Data
+/data/
+*.tsv
+*.csv
+
+# Checkpoints & Outputs
+checkpoints/
+outputs/
+logs/
+wandb/
+
+# Test cache
+.pytest_cache/
+
+# Jupyter
+.ipynb_checkpoints/
+
+# Log files
+*.log
+
+# OS
+.DS_Store
+Thumbs.db
+
+# Models
+*.pt
+*.pth
+*.bin
+*.safetensors
+*.onnx
+
+# Tokenizer artifacts
+*.model
+*.vocab
+tokenizer.json

README.md

@@ -1,77 +1,77 @@
----
-license: mit
----
-
-# EasyTranslate: 基于 Transformer 的英中翻译系统
-
-> 研究生 NLP 期末项目 — 使用前沿技术实现高质量英文到中文机器翻译
-
-## 亮点
-
-- **Pre-LayerNorm Transformer**: 从零构建，训练更稳定
-- **RoPE 旋转位置编码**: 替代传统正弦编码，更好捕获相对位置信息
-- **Flash Attention 2**: 利用 PyTorch 2.0+ 原生支持，训练显著加速
-- **NLLB 预训练微调**: 基于 Meta NLLB-200 模型，支持 LoRA 参数高效微调
-- **完整评估体系**: SacreBLEU + COMET + chrF++ 多维度评估
-- **6 组消融实验**: 系统验证各前沿组件的效果
-
-## 快速开始
-
-```bash
-# 安装
-conda create -n easytranslate python=3.11 && conda activate easytranslate
-pip install torch --index-url https://download.pytorch.org/whl/cu121
-pip install -r requirements.txt && pip install -e .
-
-# 从头训练 Transformer
-python scripts/train.py --config configs/default_config.yaml
-
-# 微调 NLLB + LoRA
-python scripts/train.py --config configs/default_config.yaml model.type=finetune_nllb
-
-# 评估
-python scripts/evaluate.py --checkpoint checkpoints/best_model.pt
-
-# 交互翻译
-python scripts/translate.py --checkpoint checkpoints/best_model.pt
-```
-
-## 项目结构
-
-```
-src/easytranslate/
-├── data/          # 数据加载、分词、预处理、动态批处理
-├── model/         # Transformer (Encoder-Decoder) + 注意力 + 位置编码 + LoRA
-├── training/      # 训练循环、优化器、损失函数、分布式训练
-├── evaluation/    # BLEU/COMET 指标、Beam Search 解码、评估器
-└── utils/         # 配置管理、日志、随机种子
-```
-
-## 技术架构
-
-```
-英文输入 → BPE Tokenizer → Source Embedding + RoPE
-    → Transformer Encoder (6 layers, Flash Attention, Pre-LN)
-    → Transformer Decoder (6 layers, Causal Mask, Cross-Attention)
-    → Linear Projection → Beam Search → 中文输出
-```
-
-## 团队分工
-
-详见 [TASK_ASSIGNMENT.md](TASK_ASSIGNMENT.md)
-
-| 角色 | 负责模块 | 核心任务 |
-|------|----------|----------|
-| Person A | `data/` | 数据集加载、BPE 分词器、预处理流水线 |
-| Person B | `model/` | Transformer 架构、Flash Attention、RoPE、LoRA |
-| Person C | `training/` | 训练循环、混合精度、梯度累积、早停 |
-| Person D | `evaluation/` | BLEU/COMET 评估、Beam Search、推理服务 |
-| Person E | `utils/` + 实验 | 工具模块、消融实验、可视化、报告 |
-
-## 依赖
-
-- Python >= 3.10
-- PyTorch >= 2.1 (Flash Attention 支持)
-- HuggingFace Transformers / Datasets / Tokenizers
-- PEFT (LoRA)
+---
+license: mit
+---
+
+# EasyTranslate: 基于 Transformer 的英中翻译系统
+
+> 研究生 NLP 期末项目 — 使用前沿技术实现高质量英文到中文机器翻译
+
+## 亮点
+
+- **Pre-LayerNorm Transformer**: 从零构建，训练更稳定
+- **RoPE 旋转位置编码**: 替代传统正弦编码，更好捕获相对位置信息
+- **Flash Attention 2**: 利用 PyTorch 2.0+ 原生支持，训练显著加速
+- **NLLB 预训练微调**: 基于 Meta NLLB-200 模型，支持 LoRA 参数高效微调
+- **完整评估体系**: SacreBLEU + COMET + chrF++ 多维度评估
+- **6 组消融实验**: 系统验证各前沿组件的效果
+
+## 快速开始
+
+```bash
+# 安装
+conda create -n easytranslate python=3.11 && conda activate easytranslate
+pip install torch --index-url https://download.pytorch.org/whl/cu121
+pip install -r requirements.txt && pip install -e .
+
+# 从头训练 Transformer
+python scripts/train.py --config configs/default_config.yaml
+
+# 微调 NLLB + LoRA
+python scripts/train.py --config configs/default_config.yaml model.type=finetune_nllb
+
+# 评估
+python scripts/evaluate.py --checkpoint checkpoints/best_model.pt
+
+# 交互翻译
+python scripts/translate.py --checkpoint checkpoints/best_model.pt
+```
+
+## 项目结构
+
+```
+src/easytranslate/
+├── data/          # 数据加载、分词、预处理、动态批处理
+├── model/         # Transformer (Encoder-Decoder) + 注意力 + 位置编码 + LoRA
+├── training/      # 训练循环、优化器、损失函数、分布式训练
+├── evaluation/    # BLEU/COMET 指标、Beam Search 解码、评估器
+└── utils/         # 配置管理、日志、随机种子
+```
+
+## 技术架构
+
+```
+英文输入 → BPE Tokenizer → Source Embedding + RoPE
+    → Transformer Encoder (6 layers, Flash Attention, Pre-LN)
+    → Transformer Decoder (6 layers, Causal Mask, Cross-Attention)
+    → Linear Projection → Beam Search → 中文输出
+```
+
+## 团队分工
+
+详见 [TASK_ASSIGNMENT.md](TASK_ASSIGNMENT.md)
+
+| 角色 | 负责模块 | 核心任务 |
+|------|----------|----------|
+| Person A | `data/` | 数据集加载、BPE 分词器、预处理流水线 |
+| Person B | `model/` | Transformer 架构、Flash Attention、RoPE、LoRA |
+| Person C | `training/` | 训练循环、混合精度、梯度累积、早停 |
+| Person D | `evaluation/` | BLEU/COMET 评估、Beam Search、推理服务 |
+| Person E | `utils/` + 实验 | 工具模块、消融实验、可视化、报告 |
+
+## 依赖
+
+- Python >= 3.10
+- PyTorch >= 2.1 (Flash Attention 支持)
+- HuggingFace Transformers / Datasets / Tokenizers
+- PEFT (LoRA)
 - SacreBLEU, COMET (评估)

TASK_ASSIGNMENT.md

@@ -1,414 +1,414 @@
-# EasyTranslate 任务分工文档
-
-## 项目概述
-
-**项目名称**: EasyTranslate — 基于 Transformer 的英中翻译系统  
-**课程**: 研究生 NLP 期末作业  
-**团队规模**: 5 人  
-**技术栈**: PyTorch + HuggingFace Transformers + Flash Attention 2 + RoPE + LoRA  
-
----
-
-## 📂 项目结构
-
-```
-UCAS-EasyTranslate/
-├── configs/
-│   ├── default_config.yaml          # 主配置文件
-│   └── deepspeed_config.json        # DeepSpeed 分布式训练配置
-├── src/easytranslate/
-│   ├── data/                        # [Person A] 数据模块
-│   │   ├── dataset.py               # 数据集加载
-│   │   ├── tokenizer.py             # 分词器
-│   │   ├── preprocessing.py         # 数据预处理
-│   │   └── collator.py              # 数据整理 & 动态批处理
-│   ├── model/                       # [Person B] 模型模块
-│   │   ├── transformer.py           # Transformer 主模型
-│   │   ├── encoder.py               # 编码器
-│   │   ├── decoder.py               # 解码器
-│   │   ├── attention.py             # 注意力机制 (标准 + Flash Attention)
-│   │   ├── positional.py            # 位置编码 (Sinusoidal + RoPE)
-│   │   └── finetune.py              # 预训练模型微调 (NLLB + LoRA)
-│   ├── training/                    # [Person C] 训练模块
-│   │   ├── trainer.py               # 训练器 (完整训练循环)
-│   │   ├── optimizer.py             # 优化器 & 学习率调度
-│   │   └── loss.py                  # 损失函数 (标签平滑)
-│   ├── evaluation/                  # [Person D] 评估模块
-│   │   ├── metrics.py               # 评估指标 (BLEU, COMET, chrF)
-│   │   ├── decoding.py              # 解码策略 (贪心, Beam Search, 采样)
-│   │   └── evaluator.py             # 评估器
-│   └── utils/                       # [Person E] 工具模块
-│       ├── config.py                # 配置管理
-│       ├── seed.py                  # 随机种子
-│       └── logging.py               # 日志管理
-├── scripts/
-│   ├── train.py                     # 训练入口
-│   ├── evaluate.py                  # 评估入口
-│   ├── translate.py                 # 翻译推理 (CLI + Web UI)
-│   ├── run_experiments.py           # [Person E] 实验运行器
-│   └── visualize.py                 # [Person E] 可视化分析
-├── tests/
-│   ├── test_data.py                 # 数据模块测试
-│   ├── test_model.py                # 模型模块测试
-│   ├── test_training.py             # 训练模块测试
-│   └── test_evaluation.py           # 评估模块测试
-├── requirements.txt
-├── setup.py
-└── README.md
-```
-
----
-
-## 👥 人员分工
-
-### Person A — 数据工程师
-
-**负责文件**: `src/easytranslate/data/` 目录下所有文件  
-**预计工作量**: ~800 行代码  
-**截止日期建议**: 第 1-2 周
-
-#### 核心任务
-
-| 优先级 | 文件 | 任务 | 要点 |
-|--------|------|------|------|
-| P0 | `dataset.py` | `TranslationDataset.__getitem__` | 实现 tokenize + padding + teacher forcing 输入构造 |
-| P0 | `dataset.py` | `load_wmt_dataset` | 使用 HuggingFace datasets 加载 WMT19 zh-en |
-| P0 | `tokenizer.py` | `TokenizerWrapper` | 统一分词器接口 (encode/decode/vocab_size/special_tokens) |
-| P0 | `tokenizer.py` | `train_bpe_tokenizer` | 使用 HuggingFace tokenizers 库训练 BPE |
-| P0 | `tokenizer.py` | `build_tokenizer` | 根据配置构建分词器 |
-| P1 | `preprocessing.py` | `clean_text` | Unicode 标准化 + 控制字符去除 |
-| P1 | `preprocessing.py` | `filter_by_length` | 按长度和长度比过滤 |
-| P1 | `preprocessing.py` | `preprocess_pipeline` | 完整预处理流水线 |
-| P1 | `collator.py` | `TranslationCollator` | batch padding + mask 生成 |
-| P2 | `collator.py` | `DynamicBatchSampler` | 按 token 数动态构建 batch |
-| P2 | `dataset.py` | `load_opus_dataset` | OPUS-100 数据集加载 |
-| P2 | `dataset.py` | `load_custom_dataset` | 自定义语料加载 |
-
-#### 技术参考
-- HuggingFace datasets: https://huggingface.co/docs/datasets
-- HuggingFace tokenizers: https://huggingface.co/docs/tokenizers
-- WMT19 zh-en: `datasets.load_dataset("wmt19", "zh-en")`
-
-#### 验收标准
-- [ ] `pytest tests/test_data.py` 全部通过
-- [ ] 能够成功加载 WMT19 数据集并完成预处理
-- [ ] BPE 分词器训练成功，encode/decode 往返一致
-- [ ] DataLoader 能正常迭代，batch 维度正确
-
----
-
-### Person B — 模型架构师
-
-**负责文件**: `src/easytranslate/model/` 目录下所有文件  
-**预计工作量**: ~1000 行代码  
-**截止日期建议**: 第 1-2 周
-
-#### 核心任务
-
-| 优先级 | 文件 | 任务 | 要点 |
-|--------|------|------|------|
-| P0 | `attention.py` | `MultiHeadAttention` | 标准多头注意力: QKV 投影 + scaled dot-product + mask |
-| P0 | `attention.py` | `FlashMultiHeadAttention` | 使用 F.scaled_dot_product_attention (PyTorch 2.0+) |
-| P0 | `positional.py` | `SinusoidalPositionalEncoding` | 经典正弦余弦位置编码 |
-| P0 | `positional.py` | `RotaryPositionalEmbedding` | RoPE 旋转位置编码 (前沿技术!) |
-| P0 | `encoder.py` | `TransformerEncoderLayer/Encoder` | Pre-LayerNorm Encoder |
-| P0 | `decoder.py` | `TransformerDecoderLayer/Decoder` | Pre-LayerNorm Decoder (self-attn + cross-attn + FFN) |
-| P0 | `transformer.py` | `TransformerTranslationModel` | 完整 Enc-Dec 模型: embedding + encoder + decoder + projection |
-| P1 | `transformer.py` | `encode / decode_step` | 推理用的编码和单步解码 |
-| P1 | `finetune.py` | `load_pretrained_model` | 加载 NLLB/mBART 预训练模型 |
-| P1 | `finetune.py` | `setup_lora` | 使用 PEFT 库配置 LoRA 微调 |
-
-#### 关键技术点
-
-1. **Pre-LayerNorm** (比 Post-LN 训练更稳定):
-   ```
-   x → LayerNorm → Attention → Add(x) → LayerNorm → FFN → Add
-   ```
-
-2. **RoPE 旋转位置编码** (LLaMA/GPT-NeoX 使用):
-   ```python
-   q' = q * cos(θ) + rotate_half(q) * sin(θ)
-   k' = k * cos(θ) + rotate_half(k) * sin(θ)
-   # attention(q', k') 自动编码相对位置信息
-   ```
-
-3. **Flash Attention 2** (PyTorch 2.0+ 原生支持):
-   ```python
-   F.scaled_dot_product_attention(Q, K, V, attn_mask, dropout_p, is_causal)
-   ```
-
-4. **LoRA 微调** (参数高效):
-   ```python
-   from peft import LoraConfig, get_peft_model, TaskType
-   config = LoraConfig(r=16, lora_alpha=32, target_modules=["q_proj", "v_proj"], task_type=TaskType.SEQ_2_SEQ_LM)
-   model = get_peft_model(model, config)
-   ```
-
-#### 验收标准
-- [ ] `pytest tests/test_model.py` 全部通过
-- [ ] 模型前向传播输出维度正确: `[B, T, vocab_size]`
-- [ ] Flash Attention 和标准 Attention 输出一致
-- [ ] RoPE 编码能正确应用到 Q, K
-- [ ] LoRA 模型可训练参数量远小于全量参数
-
----
-
-### Person C — 训练工程师
-
-**负责文件**: `src/easytranslate/training/` 目录下所有文件  
-**预计工作量**: ~700 行代码  
-**截止日期建议**: 第 2-3 周 (依赖 Person A, B)
-
-#### 核心任务
-
-| 优先级 | 文件 | 任务 | 要点 |
-|--------|------|------|------|
-| P0 | `loss.py` | `LabelSmoothedCrossEntropyLoss` | 带标签平滑的交叉熵损失 + 忽略 padding |
-| P0 | `optimizer.py` | `build_optimizer` | 构建 AdamW，支持参数分组 |
-| P0 | `optimizer.py` | `build_scheduler` | Cosine with warmup / Inverse sqrt 调度器 |
-| P0 | `trainer.py` | `Trainer.__init__` | 初始化训练环境 (设备, 精度, 分布式) |
-| P0 | `trainer.py` | `_train_one_epoch` | 单 epoch 训练循环 (混合精度 + 梯度累积) |
-| P0 | `trainer.py` | `_validate` | 验证循环 |
-| P1 | `trainer.py` | `train` | 主训练循环 (多 epoch) |
-| P1 | `trainer.py` | `_save/_load_checkpoint` | 检查点保存和加载 |
-| P1 | `trainer.py` | `_should_early_stop` | 早停逻辑 |
-| P2 | `trainer.py` | `_setup_distributed` | DDP/FSDP/DeepSpeed 分布式训练 |
-| P2 | `trainer.py` | `_log_metrics` | TensorBoard/WandB 日志 |
-
-#### 关键技术点
-
-1. **混合精度训练**:
-   ```python
-   scaler = torch.amp.GradScaler('cuda')
-   with torch.amp.autocast('cuda', dtype=torch.float16):
-       logits = model(src, tgt)
-       loss = criterion(logits, labels)
-   scaler.scale(loss).backward()
-   scaler.step(optimizer)
-   scaler.update()
-   ```
-
-2. **梯度累积**:
-   ```python
-   loss = loss / gradient_accumulation_steps
-   loss.backward()
-   if (step + 1) % gradient_accumulation_steps == 0:
-       optimizer.step()
-       optimizer.zero_grad()
-   ```
-
-3. **标签平滑** (smoothing=0.1):
-   ```
-   对于 target token: prob = 1 - smoothing = 0.9
-   对于其他 token: prob = smoothing / (V - 1) ≈ 0.000003
-   ```
-
-#### 验收标准
-- [ ] `pytest tests/test_training.py` 全部通过
-- [ ] 能在小数据集上完成完整训练流程
-- [ ] Loss 能正常下降
-- [ ] 检查点能正确保存和加载
-- [ ] 早停机制正常工作
-
----
-
-### Person D — 评估与推理工程师
-
-**负责文件**: `src/easytranslate/evaluation/` + `scripts/evaluate.py` + `scripts/translate.py`  
-**预计工作量**: ~800 行代码  
-**截止日期建议**: 第 2-3 周 (依赖 Person B)
-
-#### 核心任务
-
-| 优先级 | 文件 | 任务 | 要点 |
-|--------|------|------|------|
-| P0 | `metrics.py` | `compute_bleu` | SacreBLEU (tokenize="zh" 对中文分词) |
-| P0 | `metrics.py` | `compute_comet` | COMET 神经网络评估指标 |
-| P0 | `metrics.py` | `compute_chrf` / `compute_ter` | chrF++ 和 TER 指标 |
-| P0 | `decoding.py` | `greedy_decode` | 贪心解码 (逐 token argmax) |
-| P0 | `decoding.py` | `beam_search_decode` | 束搜索 (最关键的解码算法!) |
-| P1 | `decoding.py` | `sample_decode` | 采样解码 (temperature + top-k + top-p) |
-| P1 | `evaluator.py` | `Evaluator` | 统一评估接口 |
-| P1 | `scripts/evaluate.py` | 评估入口脚本 | 加载模型 → 评估 → 输出结果 |
-| P2 | `scripts/translate.py` | 交互翻译 | CLI 交互 + 文件翻译 |
-| P2 | `scripts/translate.py` | `launch_web_ui` | Gradio Web 界面 |
-
-#### 关键技术点
-
-1. **Beam Search** (翻译最核心的算法):
-   ```
-   维护 beam_size 个候选序列
-   每步扩展所有候选 → 选 top-k → 继续
-   完成后按 score / length^penalty 排序
-   ```
-
-2. **SacreBLEU** (标准化评估):
-   ```python
-   import sacrebleu
-   bleu = sacrebleu.corpus_bleu(hypotheses, [references], tokenize="zh")
-   ```
-
-3. **COMET** (最准确的评估指标):
-   ```python
-   from comet import download_model, load_from_checkpoint
-   model = load_from_checkpoint(download_model("Unbabel/wmt22-comet-da"))
-   data = [{"src": s, "mt": h, "ref": r} for s, h, r in zip(sources, hyps, refs)]
-   output = model.predict(data)
-   ```
-
-#### 验收标准
-- [ ] `pytest tests/test_evaluation.py` 全部通过
-- [ ] BLEU 计算结果与 sacrebleu CLI 一致
-- [ ] Beam search 翻译质量 ≥ Greedy
-- [ ] 能完成端到端评估流程
-- [ ] 交互翻译模式正常工作
-
----
-
-### Person E — 实验与报告
-
-**负责文件**: `src/easytranslate/utils/` + `scripts/run_experiments.py` + `scripts/visualize.py` + 实验报告  
-**预计工作量**: ~600 行代码 + 实验报告  
-**截止日期建议**: 第 1 周 (utils) + 第 3-4 周 (实验)
-
-#### 核心任务
-
-| 优先级 | 文件 | 任务 | 要点 |
-|--------|------|------|------|
-| P0 | `utils/config.py` | 配置管理 | OmegaConf 加载 + CLI 覆盖 |
-| P0 | `utils/seed.py` | 随机种子 | 全局种子设置 (可复现) |
-| P0 | `utils/logging.py` | 日志系统 | Rich 美化 + 文件日志 |
-| P1 | `scripts/run_experiments.py` | 实验运行器 | 自动化运行 6 组消融实验 |
-| P1 | `scripts/visualize.py` | 训练曲线 | Loss/BLEU/LR 曲线绘制 |
-| P1 | `scripts/visualize.py` | 实验对比 | 柱状图 + 表格对比 |
-| P2 | `scripts/visualize.py` | 注意力可视化 | Cross-attention 热力图 |
-| P2 | `scripts/visualize.py` | 翻译样例 | 好坏案例展示 |
-| P0 | — | 实验报告 | 撰写完整实验报告 |
-
-#### 实验设计 (6 组消融实验)
-
-| 实验编号 | 实验名称 | 关键变量 | 目的 |
-|----------|----------|----------|------|
-| Exp1 | Baseline Transformer | 标准 6 层, d=512, 正弦位置编码 | 基线 |
-| Exp2 | + RoPE | 替换正弦编码为 RoPE | 验证旋转位置编码效果 |
-| Exp3 | + Flash Attention | 使用 Flash Attention 2 | 验证训练加速效果 |
-| Exp4 | Full (RoPE + Flash) | RoPE + Flash Attention | 完整前沿方案 |
-| Exp5 | NLLB + LoRA | 预训练 NLLB-600M + LoRA r=16 | 预训练微调效果 |
-| Exp6 | NLLB Full FT | 预训练 NLLB-600M 全量微调 | 微调上界 |
-
-#### 报告结构建议
-1. **摘要**: 项目目标、方法、主要结果
-2. **引言**: 机器翻译背景、Transformer 发展、研究动机
-3. **相关工作**: Transformer、NLLB、LoRA、Flash Attention、RoPE
-4. **方法**:
-   - 模型架构 (附架构图)
-   - 训练策略 (混合精度、标签平滑、调度器)
-   - 预训练微调方案 (NLLB + LoRA)
-5. **实验**:
-   - 数据集 (WMT19 zh-en)
-   - 实验设置 (超参数表)
-   - 实验结果 (BLEU/COMET/chrF 表格)
-   - 消融分析 (各组件贡献)
-   - 训练效率对比 (Flash Attention 加速比)
-6. **分析**:
-   - 注意力可视化
-   - 翻译案例分析
-   - 错误分析
-7. **结论与展望**
-
-#### 验收标准
-- [ ] utils 模块功能正常
-- [ ] 6 组实验能自动化运行
-- [ ] 生成完整的可视化图表
-- [ ] 实验报告完整，图表清晰
-
----
-
-## 🗓️ 时间线
-
-```
-第 1 周: Person A (数据) + Person B (模型) + Person E (utils) 并行开发
-         ↓
-第 2 周: Person C (训练, 依赖 A+B) + Person D (评估, 依赖 B) 开始
-         Person A, B 完善和联调
-         ↓
-第 3 周: 全体联调 → scripts/train.py 整合
-         Person E 开始跑实验
-         ↓
-第 4 周: Person E 完成实验 + 报告
-         全体 Review + 优化
-```
-
-## 🔌 模块接口约定
-
-### 分词器接口 (Person A 定义, Person B/C/D 使用)
-```python
-tokenizer.encode(text: str) -> list[int]
-tokenizer.decode(ids: list[int]) -> str
-tokenizer.vocab_size -> int
-tokenizer.pad_id -> int
-tokenizer.bos_id -> int
-tokenizer.eos_id -> int
-```
-
-### 模型接口 (Person B 定义, Person C/D 使用)
-```python
-# 训练
-logits = model(src_ids, tgt_input_ids, src_padding_mask, tgt_padding_mask)
-# logits: [B, T, vocab_size]
-
-# 推理
-encoder_output = model.encode(src_ids, src_padding_mask)
-next_logits = model.decode_step(tgt_input_ids, encoder_output, src_padding_mask)
-```
-
-### 数据 Batch 格式 (Person A 定义, Person C 使用)
-```python
-batch = {
-    "src_ids": Tensor[B, S],
-    "tgt_input_ids": Tensor[B, T],
-    "labels": Tensor[B, T],
-    "src_padding_mask": BoolTensor[B, S],
-    "tgt_padding_mask": BoolTensor[B, T],
-}
-```
-
-### 评估接口 (Person D 定义, Person C/E 使用)
-```python
-evaluator = Evaluator(model, tokenizer, config)
-results = evaluator.evaluate(dataloader)
-# results: {"bleu": 25.6, "comet": 0.82, "chrf": 45.3, "ter": 55.2}
-```
-
----
-
-## ⚙️ 开发环境
-
-```bash
-# 1. 创建虚拟环境
-conda create -n easytranslate python=3.11
-conda activate easytranslate
-
-# 2. 安装 PyTorch (CUDA 12.1)
-pip install torch torchvision torchaudio --index-url https://download.pytorch.org/whl/cu121
-
-# 3. 安装项目依赖
-pip install -r requirements.txt
-
-# 4. 安装项目 (开发模式)
-pip install -e .
-
-# 5. 运行测试
-pytest tests/ -v
-```
-
----
-
-## ✅ 最终交付清单
-
-- [ ] 代码: 所有 `TODO: Person X` 均已实现
-- [ ] 测试: `pytest tests/` 全部通过
-- [ ] 训练: 至少完成 Exp1 (基线) 和 Exp5 (NLLB+LoRA) 两组实验
-- [ ] 评估: 在 WMT19 测试集上报告 BLEU / COMET 分数
-- [ ] 报告: 完整实验报告 (含图表)
-- [ ] 演示: 交互翻译 Demo 可运行
+# EasyTranslate 任务分工文档
+
+## 项目概述
+
+**项目名称**: EasyTranslate — 基于 Transformer 的英中翻译系统  
+**课程**: 研究生 NLP 期末作业  
+**团队规模**: 5 人  
+**技术栈**: PyTorch + HuggingFace Transformers + Flash Attention 2 + RoPE + LoRA  
+
+---
+
+## 📂 项目结构
+
+```
+UCAS-EasyTranslate/
+├── configs/
+│   ├── default_config.yaml          # 主配置文件
+│   └── deepspeed_config.json        # DeepSpeed 分布式训练配置
+├── src/easytranslate/
+│   ├── data/                        # [Person A] 数据模块
+│   │   ├── dataset.py               # 数据集加载
+│   │   ├── tokenizer.py             # 分词器
+│   │   ├── preprocessing.py         # 数据预处理
+│   │   └── collator.py              # 数据整理 & 动态批处理
+│   ├── model/                       # [Person B] 模型模块
+│   │   ├── transformer.py           # Transformer 主模型
+│   │   ├── encoder.py               # 编码器
+│   │   ├── decoder.py               # 解码器
+│   │   ├── attention.py             # 注意力机制 (标准 + Flash Attention)
+│   │   ├── positional.py            # 位置编码 (Sinusoidal + RoPE)
+│   │   └── finetune.py              # 预训练模型微调 (NLLB + LoRA)
+│   ├── training/                    # [Person C] 训练模块
+│   │   ├── trainer.py               # 训练器 (完整训练循环)
+│   │   ├── optimizer.py             # 优化器 & 学习率调度
+│   │   └── loss.py                  # 损失函数 (标签平滑)
+│   ├── evaluation/                  # [Person D] 评估模块
+│   │   ├── metrics.py               # 评估指标 (BLEU, COMET, chrF)
+│   │   ├── decoding.py              # 解码策略 (贪心, Beam Search, 采样)
+│   │   └── evaluator.py             # 评估器
+│   └── utils/                       # [Person E] 工具模块
+│       ├── config.py                # 配置管理
+│       ├── seed.py                  # 随机种子
+│       └── logging.py               # 日志管理
+├── scripts/
+│   ├── train.py                     # 训练入口
+│   ├── evaluate.py                  # 评估入口
+│   ├── translate.py                 # 翻译推理 (CLI + Web UI)
+│   ├── run_experiments.py           # [Person E] 实验运行器
+│   └── visualize.py                 # [Person E] 可视化分析
+├── tests/
+│   ├── test_data.py                 # 数据模块测试
+│   ├── test_model.py                # 模型模块测试
+│   ├── test_training.py             # 训练模块测试
+│   └── test_evaluation.py           # 评估模块测试
+├── requirements.txt
+├── setup.py
+└── README.md
+```
+
+---
+
+## 👥 人员分工
+
+### Person A — 数据工程师
+
+**负责文件**: `src/easytranslate/data/` 目录下所有文件  
+**预计工作量**: ~800 行代码  
+**截止日期建议**: 第 1-2 周
+
+#### 核心任务
+
+| 优先级 | 文件 | 任务 | 要点 |
+|--------|------|------|------|
+| P0 | `dataset.py` | `TranslationDataset.__getitem__` | 实现 tokenize + padding + teacher forcing 输入构造 |
+| P0 | `dataset.py` | `load_wmt_dataset` | 使用 HuggingFace datasets 加载 WMT19 zh-en |
+| P0 | `tokenizer.py` | `TokenizerWrapper` | 统一分词器接口 (encode/decode/vocab_size/special_tokens) |
+| P0 | `tokenizer.py` | `train_bpe_tokenizer` | 使用 HuggingFace tokenizers 库训练 BPE |
+| P0 | `tokenizer.py` | `build_tokenizer` | 根据配置构建分词器 |
+| P1 | `preprocessing.py` | `clean_text` | Unicode 标准化 + 控制字符去除 |
+| P1 | `preprocessing.py` | `filter_by_length` | 按长度和长度比过滤 |
+| P1 | `preprocessing.py` | `preprocess_pipeline` | 完整预处理流水线 |
+| P1 | `collator.py` | `TranslationCollator` | batch padding + mask 生成 |
+| P2 | `collator.py` | `DynamicBatchSampler` | 按 token 数动态构建 batch |
+| P2 | `dataset.py` | `load_opus_dataset` | OPUS-100 数据集加载 |
+| P2 | `dataset.py` | `load_custom_dataset` | 自定义语料加载 |
+
+#### 技术参考
+- HuggingFace datasets: https://huggingface.co/docs/datasets
+- HuggingFace tokenizers: https://huggingface.co/docs/tokenizers
+- WMT19 zh-en: `datasets.load_dataset("wmt19", "zh-en")`
+
+#### 验收标准
+- [ ] `pytest tests/test_data.py` 全部通过
+- [ ] 能够成功加载 WMT19 数据集并完成预处理
+- [ ] BPE 分词器训练成功，encode/decode 往返一致
+- [ ] DataLoader 能正常迭代，batch 维度正确
+
+---
+
+### Person B — 模型架构师
+
+**负责文件**: `src/easytranslate/model/` 目录下所有文件  
+**预计工作量**: ~1000 行代码  
+**截止日期建议**: 第 1-2 周
+
+#### 核心任务
+
+| 优先级 | 文件 | 任务 | 要点 |
+|--------|------|------|------|
+| P0 | `attention.py` | `MultiHeadAttention` | 标准多头注意力: QKV 投影 + scaled dot-product + mask |
+| P0 | `attention.py` | `FlashMultiHeadAttention` | 使用 F.scaled_dot_product_attention (PyTorch 2.0+) |
+| P0 | `positional.py` | `SinusoidalPositionalEncoding` | 经典正弦余弦位置编码 |
+| P0 | `positional.py` | `RotaryPositionalEmbedding` | RoPE 旋转位置编码 (前沿技术!) |
+| P0 | `encoder.py` | `TransformerEncoderLayer/Encoder` | Pre-LayerNorm Encoder |
+| P0 | `decoder.py` | `TransformerDecoderLayer/Decoder` | Pre-LayerNorm Decoder (self-attn + cross-attn + FFN) |
+| P0 | `transformer.py` | `TransformerTranslationModel` | 完整 Enc-Dec 模型: embedding + encoder + decoder + projection |
+| P1 | `transformer.py` | `encode / decode_step` | 推理用的编码和单步解码 |
+| P1 | `finetune.py` | `load_pretrained_model` | 加载 NLLB/mBART 预训练模型 |
+| P1 | `finetune.py` | `setup_lora` | 使用 PEFT 库配置 LoRA 微调 |
+
+#### 关键技术点
+
+1. **Pre-LayerNorm** (比 Post-LN 训练更稳定):
+   ```
+   x → LayerNorm → Attention → Add(x) → LayerNorm → FFN → Add
+   ```
+
+2. **RoPE 旋转位置编码** (LLaMA/GPT-NeoX 使用):
+   ```python
+   q' = q * cos(θ) + rotate_half(q) * sin(θ)
+   k' = k * cos(θ) + rotate_half(k) * sin(θ)
+   # attention(q', k') 自动编码相对位置信息
+   ```
+
+3. **Flash Attention 2** (PyTorch 2.0+ 原生支持):
+   ```python
+   F.scaled_dot_product_attention(Q, K, V, attn_mask, dropout_p, is_causal)
+   ```
+
+4. **LoRA 微调** (参数高效):
+   ```python
+   from peft import LoraConfig, get_peft_model, TaskType
+   config = LoraConfig(r=16, lora_alpha=32, target_modules=["q_proj", "v_proj"], task_type=TaskType.SEQ_2_SEQ_LM)
+   model = get_peft_model(model, config)
+   ```
+
+#### 验收标准
+- [ ] `pytest tests/test_model.py` 全部通过
+- [ ] 模型前向传播输出维度正确: `[B, T, vocab_size]`
+- [ ] Flash Attention 和标准 Attention 输出一致
+- [ ] RoPE 编码能正确应用到 Q, K
+- [ ] LoRA 模型可训练参数量远小于全量参数
+
+---
+
+### Person C — 训练工程师
+
+**负责文件**: `src/easytranslate/training/` 目录下所有文件  
+**预计工作量**: ~700 行代码  
+**截止日期建议**: 第 2-3 周 (依赖 Person A, B)
+
+#### 核心任务
+
+| 优先级 | 文件 | 任务 | 要点 |
+|--------|------|------|------|
+| P0 | `loss.py` | `LabelSmoothedCrossEntropyLoss` | 带标签平滑的交叉熵损失 + 忽略 padding |
+| P0 | `optimizer.py` | `build_optimizer` | 构建 AdamW，支持参数分组 |
+| P0 | `optimizer.py` | `build_scheduler` | Cosine with warmup / Inverse sqrt 调度器 |
+| P0 | `trainer.py` | `Trainer.__init__` | 初始化训练环境 (设备, 精度, 分布式) |
+| P0 | `trainer.py` | `_train_one_epoch` | 单 epoch 训练循环 (混合精度 + 梯度累积) |
+| P0 | `trainer.py` | `_validate` | 验证循环 |
+| P1 | `trainer.py` | `train` | 主训练循环 (多 epoch) |
+| P1 | `trainer.py` | `_save/_load_checkpoint` | 检查点保存和加载 |
+| P1 | `trainer.py` | `_should_early_stop` | 早停逻辑 |
+| P2 | `trainer.py` | `_setup_distributed` | DDP/FSDP/DeepSpeed 分布式训练 |
+| P2 | `trainer.py` | `_log_metrics` | TensorBoard/WandB 日志 |
+
+#### 关键技术点
+
+1. **混合精度训练**:
+   ```python
+   scaler = torch.amp.GradScaler('cuda')
+   with torch.amp.autocast('cuda', dtype=torch.float16):
+       logits = model(src, tgt)
+       loss = criterion(logits, labels)
+   scaler.scale(loss).backward()
+   scaler.step(optimizer)
+   scaler.update()
+   ```
+
+2. **梯度累积**:
+   ```python
+   loss = loss / gradient_accumulation_steps
+   loss.backward()
+   if (step + 1) % gradient_accumulation_steps == 0:
+       optimizer.step()
+       optimizer.zero_grad()
+   ```
+
+3. **标签平滑** (smoothing=0.1):
+   ```
+   对于 target token: prob = 1 - smoothing = 0.9
+   对于其他 token: prob = smoothing / (V - 1) ≈ 0.000003
+   ```
+
+#### 验收标准
+- [ ] `pytest tests/test_training.py` 全部通过
+- [ ] 能在小数据集上完成完整训练流程
+- [ ] Loss 能正常下降
+- [ ] 检查点能正确保存和加载
+- [ ] 早停机制正常工作
+
+---
+
+### Person D — 评估与推理工程师
+
+**负责文件**: `src/easytranslate/evaluation/` + `scripts/evaluate.py` + `scripts/translate.py`  
+**预计工作量**: ~800 行代码  
+**截止日期建议**: 第 2-3 周 (依赖 Person B)
+
+#### 核心任务
+
+| 优先级 | 文件 | 任务 | 要点 |
+|--------|------|------|------|
+| P0 | `metrics.py` | `compute_bleu` | SacreBLEU (tokenize="zh" 对中文分词) |
+| P0 | `metrics.py` | `compute_comet` | COMET 神经网络评估指标 |
+| P0 | `metrics.py` | `compute_chrf` / `compute_ter` | chrF++ 和 TER 指标 |
+| P0 | `decoding.py` | `greedy_decode` | 贪心解码 (逐 token argmax) |
+| P0 | `decoding.py` | `beam_search_decode` | 束搜索 (最关键的解码算法!) |
+| P1 | `decoding.py` | `sample_decode` | 采样解码 (temperature + top-k + top-p) |
+| P1 | `evaluator.py` | `Evaluator` | 统一评估接口 |
+| P1 | `scripts/evaluate.py` | 评估入口脚本 | 加载模型 → 评估 → 输出结果 |
+| P2 | `scripts/translate.py` | 交互翻译 | CLI 交互 + 文件翻译 |
+| P2 | `scripts/translate.py` | `launch_web_ui` | Gradio Web 界面 |
+
+#### 关键技术点
+
+1. **Beam Search** (翻译最核心的算法):
+   ```
+   维护 beam_size 个候选序列
+   每步扩展所有候选 → 选 top-k → 继续
+   完成后按 score / length^penalty 排序
+   ```
+
+2. **SacreBLEU** (标准化评估):
+   ```python
+   import sacrebleu
+   bleu = sacrebleu.corpus_bleu(hypotheses, [references], tokenize="zh")
+   ```
+
+3. **COMET** (最准确的评估指标):
+   ```python
+   from comet import download_model, load_from_checkpoint
+   model = load_from_checkpoint(download_model("Unbabel/wmt22-comet-da"))
+   data = [{"src": s, "mt": h, "ref": r} for s, h, r in zip(sources, hyps, refs)]
+   output = model.predict(data)
+   ```
+
+#### 验收标准
+- [ ] `pytest tests/test_evaluation.py` 全部通过
+- [ ] BLEU 计算结果与 sacrebleu CLI 一致
+- [ ] Beam search 翻译质量 ≥ Greedy
+- [ ] 能完成端到端评估流程
+- [ ] 交互翻译模式正常工作
+
+---
+
+### Person E — 实验与报告
+
+**负责文件**: `src/easytranslate/utils/` + `scripts/run_experiments.py` + `scripts/visualize.py` + 实验报告  
+**预计工作量**: ~600 行代码 + 实验报告  
+**截止日期建议**: 第 1 周 (utils) + 第 3-4 周 (实验)
+
+#### 核心任务
+
+| 优先级 | 文件 | 任务 | 要点 |
+|--------|------|------|------|
+| P0 | `utils/config.py` | 配置管理 | OmegaConf 加载 + CLI 覆盖 |
+| P0 | `utils/seed.py` | 随机种子 | 全局种子设置 (可复现) |
+| P0 | `utils/logging.py` | 日志系统 | Rich 美化 + 文件日志 |
+| P1 | `scripts/run_experiments.py` | 实验运行器 | 自动化运行 6 组消融实验 |
+| P1 | `scripts/visualize.py` | 训练曲线 | Loss/BLEU/LR 曲线绘制 |
+| P1 | `scripts/visualize.py` | 实验对比 | 柱状图 + 表格对比 |
+| P2 | `scripts/visualize.py` | 注意力可视化 | Cross-attention 热力图 |
+| P2 | `scripts/visualize.py` | 翻译样例 | 好坏案例展示 |
+| P0 | — | 实验报告 | 撰写完整实验报告 |
+
+#### 实验设计 (6 组消融实验)
+
+| 实验编号 | 实验名称 | 关键变量 | 目的 |
+|----------|----------|----------|------|
+| Exp1 | Baseline Transformer | 标准 6 层, d=512, 正弦位置编码 | 基线 |
+| Exp2 | + RoPE | 替换正弦编码为 RoPE | 验证旋转位置编码效果 |
+| Exp3 | + Flash Attention | 使用 Flash Attention 2 | 验证训练加速效果 |
+| Exp4 | Full (RoPE + Flash) | RoPE + Flash Attention | 完整前沿方案 |
+| Exp5 | NLLB + LoRA | 预训练 NLLB-600M + LoRA r=16 | 预训练微调效果 |
+| Exp6 | NLLB Full FT | 预训练 NLLB-600M 全量微调 | 微调上界 |
+
+#### 报告结构建议
+1. **摘要**: 项目目标、方法、主要结果
+2. **引言**: 机器翻译背景、Transformer 发展、研究动机
+3. **相关工作**: Transformer、NLLB、LoRA、Flash Attention、RoPE
+4. **方法**:
+   - 模型架构 (附架构图)
+   - 训练策略 (混合精度、标签平滑、调度器)
+   - 预训练微调方案 (NLLB + LoRA)
+5. **实验**:
+   - 数据集 (WMT19 zh-en)
+   - 实验设置 (超参数表)
+   - 实验结果 (BLEU/COMET/chrF 表格)
+   - 消融分析 (各组件贡献)
+   - 训练效率对比 (Flash Attention 加速比)
+6. **分析**:
+   - 注意力可视化
+   - 翻译案例分析
+   - 错误分析
+7. **结论与展望**
+
+#### 验收标准
+- [ ] utils 模块功能正常
+- [ ] 6 组实验能自动化运行
+- [ ] 生成完整的可视化图表
+- [ ] 实验报告完整，图表清晰
+
+---
+
+## 🗓️ 时间线
+
+```
+第 1 周: Person A (数据) + Person B (模型) + Person E (utils) 并行开发
+         ↓
+第 2 周: Person C (训练, 依赖 A+B) + Person D (评估, 依赖 B) 开始
+         Person A, B 完善和联调
+         ↓
+第 3 周: 全体联调 → scripts/train.py 整合
+         Person E 开始跑实验
+         ↓
+第 4 周: Person E 完成实验 + 报告
+         全体 Review + 优化
+```
+
+## 🔌 模块接口约定
+
+### 分词器接口 (Person A 定义, Person B/C/D 使用)
+```python
+tokenizer.encode(text: str) -> list[int]
+tokenizer.decode(ids: list[int]) -> str
+tokenizer.vocab_size -> int
+tokenizer.pad_id -> int
+tokenizer.bos_id -> int
+tokenizer.eos_id -> int
+```
+
+### 模型接口 (Person B 定义, Person C/D 使用)
+```python
+# 训练
+logits = model(src_ids, tgt_input_ids, src_padding_mask, tgt_padding_mask)
+# logits: [B, T, vocab_size]
+
+# 推理
+encoder_output = model.encode(src_ids, src_padding_mask)
+next_logits = model.decode_step(tgt_input_ids, encoder_output, src_padding_mask)
+```
+
+### 数据 Batch 格式 (Person A 定义, Person C 使用)
+```python
+batch = {
+    "src_ids": Tensor[B, S],
+    "tgt_input_ids": Tensor[B, T],
+    "labels": Tensor[B, T],
+    "src_padding_mask": BoolTensor[B, S],
+    "tgt_padding_mask": BoolTensor[B, T],
+}
+```
+
+### 评估接口 (Person D 定义, Person C/E 使用)
+```python
+evaluator = Evaluator(model, tokenizer, config)
+results = evaluator.evaluate(dataloader)
+# results: {"bleu": 25.6, "comet": 0.82, "chrf": 45.3, "ter": 55.2}
+```
+
+---
+
+## ⚙️ 开发环境
+
+```bash
+# 1. 创建虚拟环境
+conda create -n easytranslate python=3.11
+conda activate easytranslate
+
+# 2. 安装 PyTorch (CUDA 12.1)
+pip install torch torchvision torchaudio --index-url https://download.pytorch.org/whl/cu121
+
+# 3. 安装项目依赖
+pip install -r requirements.txt
+
+# 4. 安装项目 (开发模式)
+pip install -e .
+
+# 5. 运行测试
+pytest tests/ -v
+```
+
+---
+
+## ✅ 最终交付清单
+
+- [ ] 代码: 所有 `TODO: Person X` 均已实现
+- [ ] 测试: `pytest tests/` 全部通过
+- [ ] 训练: 至少完成 Exp1 (基线) 和 Exp5 (NLLB+LoRA) 两组实验
+- [ ] 评估: 在 WMT19 测试集上报告 BLEU / COMET 分数
+- [ ] 报告: 完整实验报告 (含图表)
+- [ ] 演示: 交互翻译 Demo 可运行

configs/deepspeed_config.json

@@ -1,28 +1,28 @@
-{
-    "fp16": {
-        "enabled": true,
-        "loss_scale": 0,
-        "loss_scale_window": 1000,
-        "initial_scale_power": 16,
-        "hysteresis": 2,
-        "min_loss_scale": 1
-    },
-    "zero_optimization": {
-        "stage": 2,
-        "offload_optimizer": {
-            "device": "cpu",
-            "pin_memory": true
-        },
-        "allgather_partitions": true,
-        "allgather_bucket_size": 2e8,
-        "overlap_comm": true,
-        "reduce_scatter": true,
-        "reduce_bucket_size": 2e8,
-        "contiguous_gradients": true
-    },
-    "gradient_accumulation_steps": 4,
-    "gradient_clipping": 1.0,
-    "train_batch_size": "auto",
-    "train_micro_batch_size_per_gpu": "auto",
-    "wall_clock_breakdown": false
-}
+{
+    "fp16": {
+        "enabled": true,
+        "loss_scale": 0,
+        "loss_scale_window": 1000,
+        "initial_scale_power": 16,
+        "hysteresis": 2,
+        "min_loss_scale": 1
+    },
+    "zero_optimization": {
+        "stage": 2,
+        "offload_optimizer": {
+            "device": "cpu",
+            "pin_memory": true
+        },
+        "allgather_partitions": true,
+        "allgather_bucket_size": 2e8,
+        "overlap_comm": true,
+        "reduce_scatter": true,
+        "reduce_bucket_size": 2e8,
+        "contiguous_gradients": true
+    },
+    "gradient_accumulation_steps": 4,
+    "gradient_clipping": 1.0,
+    "train_batch_size": "auto",
+    "train_micro_batch_size_per_gpu": "auto",
+    "wall_clock_breakdown": false
+}

configs/default_config.yaml

@@ -1,182 +1,183 @@
-# ============================================================================
-# EasyTranslate 默认配置文件
-# 基于 Transformer 架构的英中翻译模型
-# ============================================================================
-
-# ---------- 模型配置 ----------
-model:
-  # 模型类型: "transformer_scratch" | "finetune_nllb" | "finetune_mbart"
-  type: "transformer_scratch"
-
-  # === Transformer from scratch 配置 ===
-  transformer:
-    d_model: 512               # 模型维度
-    nhead: 8                   # 多头注意力头数
-    num_encoder_layers: 6      # 编码器层数
-    num_decoder_layers: 6      # 解码器层数
-    dim_feedforward: 2048      # FFN 中间维度
-    dropout: 0.1               # Dropout 概率
-    activation: "gelu"         # 激活函数: "relu" | "gelu"
-    max_seq_len: 512           # 最大序列长度
-    use_flash_attention: true  # 是否使用 Flash Attention 2
-    use_rotary_embedding: true # 是否使用旋转位置编码 (RoPE)
-    pre_norm: true             # Pre-LayerNorm (更稳定)
-
-  # === 预训练模型微调配置 ===
-  pretrained:
-    model_name: "facebook/nllb-200-distilled-600M"  # HuggingFace 模型名
-    src_lang: "eng_Latn"       # 源语言代码
-    tgt_lang: "zho_Hans"       # 目标语言代码
-    use_lora: true             # 是否使用 LoRA 微调
-    lora:
-      r: 16                    # LoRA 秩
-      alpha: 32                # LoRA alpha
-      dropout: 0.05            # LoRA dropout
-      target_modules:          # LoRA 目标模块
-        - "q_proj"
-        - "v_proj"
-        - "k_proj"
-        - "o_proj"
-
-# ---------- 分词器配置 ----------
-tokenizer:
-  # 分词器类型: "bpe" | "sentencepiece" | "pretrained"
-  type: "bpe"
-  vocab_size: 32000            # 词汇表大小
-  min_frequency: 2             # 最小词频
-  special_tokens:
-    pad: "<pad>"
-    unk: "<unk>"
-    bos: "<s>"
-    eos: "</s>"
-  max_length: 512              # 最大 token 长度
-
-# ---------- 数据配置 ----------
-data:
-  # 数据集: "wmt" | "opus" | "custom"
-  dataset_name: "wmt"
-  wmt:
-    year: "19"                 # WMT 年份
-    language_pair: "zh-en"     # 语言对
-  opus:
-    subset: "UNPC"             # OPUS 子集名
-  custom:
-    train_src: "data/train.en" # 自定义训练集源语言
-    train_tgt: "data/train.zh"
-    val_src: "data/val.en"
-    val_tgt: "data/val.zh"
-    test_src: "data/test.en"
-    test_tgt: "data/test.zh"
-
-  # 数据处理
-  preprocessing:
-    lowercase_src: false       # 源语言是否小写化
-    remove_punctuation: false  # 是否去除标点
-    max_src_len: 256           # 源语言最大长度
-    max_tgt_len: 256           # 目标语言最大长度
-    filter_by_length: true     # 是否按长度过滤
-    length_ratio_threshold: 3.0 # 长度比阈值
-
-  # DataLoader
-  dataloader:
-    batch_size: 32
-    num_workers: 4
-    pin_memory: true
-    dynamic_batching: true     # 动态 batch (按 token 数)
-    max_tokens_per_batch: 8192 # 动态 batch 最大 token 数
-
-# ---------- 训练配置 ----------
-training:
-  # 基础训练参数
-  epochs: 30
-  max_steps: -1                # -1 表示按 epoch 训练
-  gradient_accumulation_steps: 4
-  fp16: true                   # 混合精度训练
-  bf16: false                  # BF16 (A100+)
-  gradient_checkpointing: false
-
-  # 优化器
-  optimizer:
-    type: "adamw"              # "adam" | "adamw" | "adafactor"
-    lr: 3.0e-4
-    weight_decay: 0.01
-    betas: [0.9, 0.98]
-    eps: 1.0e-8
-
-  # 学习率调度
-  scheduler:
-    type: "cosine_with_warmup" # "cosine_with_warmup" | "inverse_sqrt" | "linear"
-    warmup_steps: 4000
-    min_lr: 1.0e-6
-
-  # 正则化
-  regularization:
-    label_smoothing: 0.1       # 标签平滑
-    dropout: 0.1
-
-  # 检查点
-  checkpoint:
-    save_dir: "checkpoints/"
-    save_every_n_steps: 5000
-    save_best: true            # 保存最佳模型
-    metric_for_best: "bleu"    # 选择最佳模型的指标
-    max_checkpoints: 5         # 最大保存检查点数
-
-  # 早停
-  early_stopping:
-    enabled: true
-    patience: 5
-    min_delta: 0.1
-
-  # 分布式训练
-  distributed:
-    strategy: "ddp"            # "ddp" | "fsdp" | "deepspeed"
-    deepspeed_config: "configs/deepspeed_config.json"
-
-# ---------- 评估配置 ----------
-evaluation:
-  # 评估指标
-  metrics:
-    - "bleu"                   # SacreBLEU
-    - "comet"                  # COMET (神经网络指标)
-    - "chrf"                   # chrF++
-    - "ter"                    # TER
-
-  # 解码策略
-  decoding:
-    strategy: "beam_search"    # "greedy" | "beam_search" | "sampling"
-    beam_size: 5
-    length_penalty: 1.0
-    no_repeat_ngram_size: 3
-    max_decode_len: 256
-
-    # Sampling 参数
-    sampling:
-      temperature: 0.7
-      top_k: 50
-      top_p: 0.9
-
-  # 评估频率
-  eval_every_n_steps: 1000
-  eval_on_epoch_end: true
-
-# ---------- 日志配置 ----------
-logging:
-  # 日志工具: "wandb" | "tensorboard" | "both"
-  backend: "tensorboard"
-  project_name: "EasyTranslate"
-  log_every_n_steps: 100
-  log_dir: "logs/"
-
-# ---------- 推理/部署配置 ----------
-inference:
-  model_path: "checkpoints/best_model"
-  device: "cuda"
-  batch_size: 16
-  quantization: null           # null | "int8" | "int4"
-
-# ---------- 实验配置 ----------
-experiment:
-  seed: 42
-  name: "baseline"
-  output_dir: "outputs/"
+# ============================================================================
+# EasyTranslate 默认配置文件
+# 基于 Transformer 架构的英中翻译模型
+# ============================================================================
+
+# ---------- 模型配置 ----------
+model:
+  # 模型类型: "transformer_scratch" | "finetune_nllb" | "finetune_mbart"
+  type: "transformer_scratch"
+
+  # === Transformer from scratch 配置 ===
+  transformer:
+    d_model: 512               # 模型维度
+    nhead: 8                   # 多头注意力头数
+    num_encoder_layers: 6      # 编码器层数
+    num_decoder_layers: 6      # 解码器层数
+    dim_feedforward: 2048      # FFN 中间维度
+    dropout: 0.1               # Dropout 概率
+    activation: "gelu"         # 激活函数: "relu" | "gelu"
+    max_seq_len: 512           # 最大序列长度
+    use_flash_attention: true  # 是否使用 Flash Attention 2
+    use_rotary_embedding: true # 是否使用旋转位置编码 (RoPE)
+    pre_norm: true             # Pre-LayerNorm (更稳定)
+
+  # === 预训练模型微调配置 ===
+  pretrained:
+    model_name: "facebook/nllb-200-distilled-600M"  # HuggingFace 模型名
+    src_lang: "eng_Latn"       # 源语言代码
+    tgt_lang: "zho_Hans"       # 目标语言代码
+    use_lora: true             # 是否使用 LoRA 微调
+    lora:
+      r: 16                    # LoRA 秩
+      alpha: 32                # LoRA alpha
+      dropout: 0.05            # LoRA dropout
+      target_modules:          # LoRA 目标模块
+        - "q_proj"
+        - "v_proj"
+        - "k_proj"
+        - "o_proj"
+
+# ---------- 分词器配置 ----------
+tokenizer:
+  # 分词器类型: "bpe" | "sentencepiece" | "pretrained"
+  type: "bpe"
+  vocab_size: 32000            # 词汇表大小
+  min_frequency: 2             # 最小词频
+  special_tokens:
+    pad: "<pad>"
+    unk: "<unk>"
+    bos: "<s>"
+    eos: "</s>"
+  max_length: 512              # 最大 token 长度
+  auto_train: false            # 是否允许自动从 WMT/OPUS 下载训练数据来训练 BPE tokenizer
+
+# ---------- 数据配置 ----------
+data:
+  # 数据集: "wmt" | "opus" | "custom"
+  dataset_name: "wmt"
+  wmt:
+    year: "19"                 # WMT 年份
+    language_pair: "zh-en"     # 语言对
+  opus:
+    subset: "UNPC"             # OPUS 子集名
+  custom:
+    train_src: "data/train.en" # 自定义训练集源语言
+    train_tgt: "data/train.zh"
+    val_src: "data/val.en"
+    val_tgt: "data/val.zh"
+    test_src: "data/test.en"
+    test_tgt: "data/test.zh"
+
+  # 数据处理
+  preprocessing:
+    lowercase_src: false       # 源语言是否小写化
+    remove_punctuation: false  # 是否去除标点
+    max_src_len: 256           # 源语言最大长度
+    max_tgt_len: 256           # 目标语言最大长度
+    filter_by_length: true     # 是否按长度过滤
+    length_ratio_threshold: 3.0 # 长度比阈值
+
+  # DataLoader
+  dataloader:
+    batch_size: 32
+    num_workers: 4
+    pin_memory: true
+    dynamic_batching: true     # 动态 batch (按 token 数)
+    max_tokens_per_batch: 8192 # 动态 batch 最大 token 数
+
+# ---------- 训练配置 ----------
+training:
+  # 基础训练参数
+  epochs: 30
+  max_steps: -1                # -1 表示按 epoch 训练
+  gradient_accumulation_steps: 4
+  fp16: true                   # 混合精度训练
+  bf16: false                  # BF16 (A100+)
+  gradient_checkpointing: false
+
+  # 优化器
+  optimizer:
+    type: "adamw"              # "adam" | "adamw" | "adafactor"
+    lr: 3.0e-4
+    weight_decay: 0.01
+    betas: [0.9, 0.98]
+    eps: 1.0e-8
+
+  # 学习率调度
+  scheduler:
+    type: "cosine_with_warmup" # "cosine_with_warmup" | "inverse_sqrt" | "linear"
+    warmup_steps: 4000
+    min_lr: 1.0e-6
+
+  # 正则化
+  regularization:
+    label_smoothing: 0.1       # 标签平滑
+    dropout: 0.1
+
+  # 检查点
+  checkpoint:
+    save_dir: "checkpoints/"
+    save_every_n_steps: 5000
+    save_best: true            # 保存最佳模型
+    metric_for_best: "bleu"    # 选择最佳模型的指标
+    max_checkpoints: 5         # 最大保存检查点数
+
+  # 早停
+  early_stopping:
+    enabled: true
+    patience: 5
+    min_delta: 0.1
+
+  # 分布式训练
+  distributed:
+    strategy: "ddp"            # "ddp" | "fsdp" | "deepspeed"
+    deepspeed_config: "configs/deepspeed_config.json"
+
+# ---------- 评估配置 ----------
+evaluation:
+  # 评���指标
+  metrics:
+    - "bleu"                   # SacreBLEU
+    - "comet"                  # COMET (神经网络指标)
+    - "chrf"                   # chrF++
+    - "ter"                    # TER
+
+  # 解码策略
+  decoding:
+    strategy: "beam_search"    # "greedy" | "beam_search" | "sampling"
+    beam_size: 5
+    length_penalty: 1.0
+    no_repeat_ngram_size: 3
+    max_decode_len: 256
+
+    # Sampling 参数
+    sampling:
+      temperature: 0.7
+      top_k: 50
+      top_p: 0.9
+
+  # 评估频率
+  eval_every_n_steps: 1000
+  eval_on_epoch_end: true
+
+# ---------- 日志配置 ----------
+logging:
+  # 日志工具: "wandb" | "tensorboard" | "both"
+  backend: "tensorboard"
+  project_name: "EasyTranslate"
+  log_every_n_steps: 100
+  log_dir: "logs/"
+
+# ---------- 推理/部署配置 ----------
+inference:
+  model_path: "checkpoints/best_model"
+  device: "cuda"
+  batch_size: 16
+  quantization: null           # null | "int8" | "int4"
+
+# ---------- 实验配置 ----------
+experiment:
+  seed: 42
+  name: "baseline"
+  output_dir: "outputs/"

pytest.ini

@@ -1,6 +1,6 @@
-[pytest]
-testpaths = tests
-python_files = test_*.py
-python_functions = test_*
-python_classes = Test*
-addopts = -v --tb=short
+[pytest]
+testpaths = tests
+python_files = test_*.py
+python_functions = test_*
+python_classes = Test*
+addopts = -v --tb=short

requirements.txt

@@ -1,35 +1,36 @@
-# Core
-torch>=2.1.0
-transformers>=4.36.0
-datasets>=2.16.0
-tokenizers>=0.15.0
-sentencepiece>=0.1.99
-accelerate>=0.25.0
-
-# Training
-deepspeed>=0.12.0
-bitsandbytes>=0.41.0
-peft>=0.7.0
-wandb>=0.16.0
-tensorboard>=2.15.0
-
-# Evaluation
-sacrebleu>=2.4.0
-unbabel-comet>=2.2.0
-rouge-score>=0.1.2
-nltk>=3.8.0
-
-# Utilities
-numpy>=1.24.0,<2.0
-pandas>=2.0.0
-tqdm>=4.66.0
-pyyaml>=6.0.0
-omegaconf>=2.3.0
-rich>=13.0.0
-matplotlib>=3.8.0
-seaborn>=0.13.0
-
-# Serving (optional)
-fastapi>=0.108.0
-uvicorn>=0.25.0
-gradio>=4.10.0
+# Core
+torch>=2.1.0
+transformers>=4.36.0
+datasets>=2.16.0
+tokenizers>=0.15.0
+sentencepiece>=0.1.99
+accelerate>=0.25.0
+
+# Training
+deepspeed>=0.12.0
+bitsandbytes>=0.41.0
+peft>=0.7.0
+wandb>=0.16.0
+tensorboard>=2.15.0
+
+# Evaluation
+sacrebleu>=2.4.0
+unbabel-comet>=2.2.0
+rouge-score>=0.1.2
+nltk>=3.8.0
+
+# Utilities
+numpy>=1.24.0,<2.0
+pandas>=2.0.0
+tqdm>=4.66.0
+pyyaml>=6.0.0
+omegaconf>=2.3.0
+rich>=13.0.0
+matplotlib>=3.8.0
+seaborn>=0.13.0
+
+# Serving (optional)
+fastapi>=0.108.0
+uvicorn>=0.25.0
+streamlit>=1.30.0
+gradio>=4.10.0

scripts/evaluate.py

@@ -1,50 +1,285 @@
-"""
-评估入口脚本
-
-使用方式:
-    # 在测试集上评估
-    python scripts/evaluate.py --config configs/default_config.yaml --checkpoint checkpoints/best_model.pt
-
-    # 指定解码策略
-    python scripts/evaluate.py --checkpoint checkpoints/best_model.pt evaluation.decoding.strategy=beam_search evaluation.decoding.beam_size=10
-"""
-
-import argparse
-import sys
-from pathlib import Path
-
-sys.path.insert(0, str(Path(__file__).resolve().parent.parent / "src"))
-
-
-def parse_args():
-    parser = argparse.ArgumentParser(description="EasyTranslate Evaluation")
-    parser.add_argument("--config", type=str, default="configs/default_config.yaml")
-    parser.add_argument("--checkpoint", type=str, required=True, help="模型检查点路径")
-    parser.add_argument("--output", type=str, default="outputs/evaluation_results.json", help="结果保存路径")
-    args, unknown = parser.parse_known_args()
-    return args, unknown
-
-
-def main():
-    """
-    评估主流程。
-
-    TODO [Person D]:
-    1. 加载配置和检查点
-    2. 重建模型并加载权重
-    3. 加载测试数据
-    4. 构建 Evaluator
-    5. 运行评估
-    6. 打印和保存结果
-    """
-    args, cli_overrides = parse_args()
-
-    print("=" * 60)
-    print("  EasyTranslate - Evaluation")
-    print("=" * 60)
-
-    raise NotImplementedError("TODO: Person D 实现评估主流程")
-
-
-if __name__ == "__main__":
-    main()
+"""
+评估入口脚本
+
+使用方式:
+    # 在测试集上评估
+    python scripts/evaluate.py --config configs/default_config.yaml --checkpoint checkpoints/best_model.pt
+
+    # 指定解码策略
+    python scripts/evaluate.py --checkpoint checkpoints/best_model.pt evaluation.decoding.strategy=beam_search evaluation.decoding.beam_size=10
+"""
+
+import argparse
+import json
+import sys
+from pathlib import Path
+
+try:
+    from omegaconf import OmegaConf
+except ImportError:  # pragma: no cover
+    OmegaConf = None
+    import yaml
+
+sys.path.insert(0, str(Path(__file__).resolve().parent.parent / "src"))
+
+import torch
+from torch.utils.data import DataLoader
+
+from easytranslate.data.collator import TranslationCollator
+from easytranslate.data.dataset import (
+    TranslationDataset,
+    load_custom_dataset,
+    load_opus_dataset,
+    load_wmt_dataset,
+)
+from easytranslate.data.tokenizer import TokenizerWrapper, build_tokenizer
+from easytranslate.evaluation.evaluator import Evaluator
+from easytranslate.model import TransformerTranslationModel
+from easytranslate.model.finetune import load_pretrained_model
+
+
+def parse_args():
+    parser = argparse.ArgumentParser(description="EasyTranslate Evaluation")
+    parser.add_argument("--config", type=str, default="configs/default_config.yaml")
+    parser.add_argument("--checkpoint", type=str, required=True, help="模型检查点路径")
+    parser.add_argument("--output", type=str, default="outputs/evaluation_results.json", help="结果保存路径")
+    args, unknown = parser.parse_known_args()
+    return args, unknown
+
+
+def _get_config(config, *keys, default=None):
+    value = config
+    for key in keys:
+        if isinstance(value, dict):
+            value = value.get(key, default)
+        else:
+            value = getattr(value, key, default)
+        if value is default:
+            break
+    return value
+
+
+def _load_config(path, cli_overrides=None):
+    if OmegaConf is not None:
+        config = OmegaConf.load(path)
+        if cli_overrides:
+            config = OmegaConf.merge(config, OmegaConf.from_cli(cli_overrides))
+        return config
+
+    with open(path, "r", encoding="utf-8") as fin:
+        config = yaml.safe_load(fin)
+    if cli_overrides:
+        print("Warning: OmegaConf is not installed; CLI overrides are ignored.")
+    return config
+
+
+def load_test_split(config):
+    dataset_name = _get_config(config, "data", "dataset_name")
+
+    if dataset_name == "wmt":
+        try:
+            dataset = load_wmt_dataset(
+                year=_get_config(config, "data", "wmt", "year"),
+                language_pair=_get_config(config, "data", "wmt", "language_pair"),
+                split="test",
+            )
+        except Exception:
+            dataset = load_wmt_dataset(
+                year=_get_config(config, "data", "wmt", "year"),
+                language_pair=_get_config(config, "data", "wmt", "language_pair"),
+                split="validation",
+            )
+        return list(dataset["src"]), list(dataset["tgt"])
+
+    if dataset_name == "opus":
+        try:
+            dataset = load_opus_dataset(
+                subset=_get_config(config, "data", "opus", "subset"),
+                split="test",
+            )
+        except Exception:
+            dataset = load_opus_dataset(
+                subset=_get_config(config, "data", "opus", "subset"),
+                split="validation",
+            )
+        return list(dataset["src"]), list(dataset["tgt"])
+
+    if dataset_name == "custom":
+        data = load_custom_dataset(
+            train_src=_get_config(config, "data", "custom", "train_src"),
+            train_tgt=_get_config(config, "data", "custom", "train_tgt"),
+            val_src=_get_config(config, "data", "custom", "val_src"),
+            val_tgt=_get_config(config, "data", "custom", "val_tgt"),
+            test_src=_get_config(config, "data", "custom", "test_src"),
+            test_tgt=_get_config(config, "data", "custom", "test_tgt"),
+            preprocessing_config=_get_config(config, "data", "preprocessing"),
+        )
+        if "test" not in data:
+            raise ValueError("Custom dataset missing test split")
+        return data["test"]["src"], data["test"]["tgt"]
+
+    raise ValueError(f"Unsupported dataset_name: {dataset_name}")
+
+
+def _get_tokenizer_train_texts(config, allow_auto: bool = False) -> list[str] | None:
+    dataset_name = _get_config(config, "data", "dataset_name")
+    if dataset_name == "custom":
+        custom = _get_config(config, "data", "custom") or {}
+        data = load_custom_dataset(
+            train_src=custom.get("train_src"),
+            train_tgt=custom.get("train_tgt"),
+            val_src=custom.get("val_src"),
+            val_tgt=custom.get("val_tgt"),
+            test_src=custom.get("test_src"),
+            test_tgt=custom.get("test_tgt"),
+            preprocessing_config=_get_config(config, "data", "preprocessing"),
+        )
+        return list(data["train"]["src"]) + list(data["train"]["tgt"])
+
+    if not allow_auto:
+        return None
+
+    if dataset_name == "wmt":
+        try:
+            dataset = load_wmt_dataset(
+                year=_get_config(config, "data", "wmt", "year"),
+                language_pair=_get_config(config, "data", "wmt", "language_pair"),
+                split="train",
+            )
+        except Exception:
+            dataset = load_wmt_dataset(
+                year=_get_config(config, "data", "wmt", "year"),
+                language_pair=_get_config(config, "data", "wmt", "language_pair"),
+                split="validation",
+            )
+        return list(dataset["src"]) + list(dataset["tgt"])
+
+    if dataset_name == "opus":
+        try:
+            dataset = load_opus_dataset(
+                subset=_get_config(config, "data", "opus", "subset"),
+                split="train",
+            )
+        except Exception:
+            dataset = load_opus_dataset(
+                subset=_get_config(config, "data", "opus", "subset"),
+                split="validation",
+            )
+        return list(dataset["src"]) + list(dataset["tgt"])
+
+    return None
+
+
+def build_model_and_tokenizer(config, device):
+    model_type = _get_config(config, "model", "type")
+    if model_type == "transformer_scratch":
+        tokenizer_config = _get_config(config, "tokenizer") or {}
+        tokenizer_path = tokenizer_config.get("path") or tokenizer_config.get("tokenizer_path")
+        tokenizer_type = tokenizer_config.get("type", "bpe")
+        auto_train = bool(tokenizer_config.get("auto_train", False))
+        if tokenizer_type in {"bpe", "sentencepiece"} and not tokenizer_path:
+            train_texts = _get_tokenizer_train_texts(config, allow_auto=auto_train)
+            if train_texts is None:
+                raise ValueError(
+                    "BPE tokenizer requires tokenizer.path or a local custom dataset with train texts. "
+                    "Automatic WMT/OPUS download is disabled by default. "
+                    "Set tokenizer.auto_train=true to enable it, or provide tokenizer.path/pretrained tokenizer."
+                )
+            tokenizer = build_tokenizer(tokenizer_config, train_texts=train_texts)
+        else:
+            tokenizer = build_tokenizer(tokenizer_config)
+        model = TransformerTranslationModel(
+            src_vocab_size=tokenizer.vocab_size,
+            tgt_vocab_size=tokenizer.vocab_size,
+            d_model=_get_config(config, "model", "transformer", "d_model"),
+            nhead=_get_config(config, "model", "transformer", "nhead"),
+            num_encoder_layers=_get_config(config, "model", "transformer", "num_encoder_layers"),
+            num_decoder_layers=_get_config(config, "model", "transformer", "num_decoder_layers"),
+            dim_feedforward=_get_config(config, "model", "transformer", "dim_feedforward"),
+            dropout=_get_config(config, "model", "transformer", "dropout"),
+            activation=_get_config(config, "model", "transformer", "activation"),
+            max_seq_len=_get_config(config, "model", "transformer", "max_seq_len"),
+            use_flash_attention=_get_config(config, "model", "transformer", "use_flash_attention"),
+            use_rotary_embedding=_get_config(config, "model", "transformer", "use_rotary_embedding"),
+            pre_norm=_get_config(config, "model", "transformer", "pre_norm"),
+            pad_id=tokenizer.pad_token_id,
+        )
+        return model.to(device), tokenizer
+
+    model, hf_tokenizer = load_pretrained_model(
+        config.model.pretrained.model_name,
+        config.model.pretrained.src_lang,
+        config.model.pretrained.tgt_lang,
+        device=str(device),
+    )
+    tokenizer = TokenizerWrapper(
+        hf_tokenizer,
+        pad_token=getattr(hf_tokenizer, "pad_token", "<pad>"),
+        unk_token=getattr(hf_tokenizer, "unk_token", "<unk>"),
+        bos_token=getattr(hf_tokenizer, "bos_token", "<s>"),
+        eos_token=getattr(hf_tokenizer, "eos_token", "</s>"),
+    )
+    return model, tokenizer
+
+
+def load_checkpoint(model, checkpoint_path, device):
+    checkpoint = torch.load(checkpoint_path, map_location=device)
+    if isinstance(checkpoint, dict):
+        if "model_state_dict" in checkpoint:
+            model.load_state_dict(checkpoint["model_state_dict"])
+        elif "state_dict" in checkpoint:
+            model.load_state_dict(checkpoint["state_dict"])
+        else:
+            try:
+                model.load_state_dict(checkpoint)
+            except Exception as exc:
+                raise ValueError("Checkpoint does not contain a valid model state dict") from exc
+    else:
+        raise ValueError("Unsupported checkpoint format")
+    return model
+
+
+def main():
+    args, cli_overrides = parse_args()
+
+    print("=" * 60)
+    print("  EasyTranslate - Evaluation")
+    print("=" * 60)
+
+    config = _load_config(args.config, cli_overrides)
+    device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
+    model, tokenizer = build_model_and_tokenizer(config, device)
+    model = load_checkpoint(model, args.checkpoint, device)
+
+    src_texts, tgt_texts = load_test_split(config)
+
+    dataset = TranslationDataset(
+        src_texts=src_texts,
+        tgt_texts=tgt_texts,
+        tokenizer=tokenizer,
+        max_src_len=_get_config(config, "data", "preprocessing", "max_src_len"),
+        max_tgt_len=_get_config(config, "data", "preprocessing", "max_tgt_len"),
+    )
+    collator = TranslationCollator(pad_token_id=tokenizer.pad_token_id)
+    dataloader = DataLoader(
+        dataset,
+        batch_size=config.data.dataloader.batch_size,
+        shuffle=False,
+        num_workers=int(config.data.dataloader.num_workers),
+        pin_memory=bool(config.data.dataloader.pin_memory),
+        collate_fn=collator,
+    )
+
+    evaluator = Evaluator(model, tokenizer, config)
+    results = evaluator.evaluate(dataloader, src_texts=src_texts, ref_texts=tgt_texts)
+
+    output_path = Path(args.output)
+    output_path.parent.mkdir(parents=True, exist_ok=True)
+    with output_path.open("w", encoding="utf-8") as fout:
+        json.dump(results, fout, ensure_ascii=False, indent=2)
+
+    print(json.dumps(results, ensure_ascii=False, indent=2))
+
+
+if __name__ == "__main__":
+    main()

scripts/run_experiments.py

@@ -1,113 +1,113 @@
-"""
-实验运行脚本 — Person E 负责
-
-自动化运行消融实验，比较不同配置的效果。
-
-使用方式:
-    python scripts/run_experiments.py
-"""
-
-import sys
-from pathlib import Path
-
-sys.path.insert(0, str(Path(__file__).resolve().parent.parent / "src"))
-
-
-# 实验配置列表
-EXPERIMENTS = [
-    {
-        "name": "exp1_baseline_transformer",
-        "description": "基线: 从头训练标准 Transformer (6层, d=512)",
-        "overrides": {
-            "model.type": "transformer_scratch",
-            "model.transformer.num_encoder_layers": 6,
-            "model.transformer.num_decoder_layers": 6,
-            "model.transformer.d_model": 512,
-            "model.transformer.use_flash_attention": False,
-            "model.transformer.use_rotary_embedding": False,
-        },
-    },
-    {
-        "name": "exp2_transformer_rope",
-        "description": "消融: Transformer + RoPE 位置编码",
-        "overrides": {
-            "model.type": "transformer_scratch",
-            "model.transformer.use_rotary_embedding": True,
-        },
-    },
-    {
-        "name": "exp3_transformer_flash_attn",
-        "description": "消融: Transformer + Flash Attention 2",
-        "overrides": {
-            "model.type": "transformer_scratch",
-            "model.transformer.use_flash_attention": True,
-        },
-    },
-    {
-        "name": "exp4_transformer_full",
-        "description": "完整: Transformer + RoPE + Flash Attention",
-        "overrides": {
-            "model.type": "transformer_scratch",
-            "model.transformer.use_flash_attention": True,
-            "model.transformer.use_rotary_embedding": True,
-        },
-    },
-    {
-        "name": "exp5_nllb_lora",
-        "description": "预训练微调: NLLB-600M + LoRA",
-        "overrides": {
-            "model.type": "finetune_nllb",
-            "model.pretrained.use_lora": True,
-            "model.pretrained.lora.r": 16,
-        },
-    },
-    {
-        "name": "exp6_nllb_full_finetune",
-        "description": "预训练全量微调: NLLB-600M",
-        "overrides": {
-            "model.type": "finetune_nllb",
-            "model.pretrained.use_lora": False,
-        },
-    },
-]
-
-
-def run_single_experiment(exp_config: dict):
-    """
-    运行单个实验。
-
-    TODO [Person E]:
-    1. 从 base config 创建实验配置
-    2. 应用 overrides
-    3. 设置 experiment.name 和 output_dir
-    4. 调用训练流程
-    5. 在测试集上评估
-    6. 保存实验结果
-    """
-    raise NotImplementedError("TODO: Person E 实现 run_single_experiment")
-
-
-def main():
-    """
-    运行所有实验。
-
-    TODO [Person E]:
-    1. 遍历 EXPERIMENTS 列表
-    2. 对每个实验调用 run_single_experiment
-    3. 收集所有实验结果
-    4. 生成对比表格
-    5. 保存汇总报告
-    """
-    print("=" * 60)
-    print("  EasyTranslate - Experiment Runner")
-    print("=" * 60)
-
-    for exp in EXPERIMENTS:
-        print(f"\n>>> Running: {exp['name']} - {exp['description']}")
-        # run_single_experiment(exp)
-
-    raise NotImplementedError("TODO: Person E 实现实验运行主流程")
-
-
-if __name__ == "__main__":
-    main()
+"""
+实验运行脚本 — Person E 负责
+
+自动化运行消融实验，比较不同配置的效果。
+
+使用方式:
+    python scripts/run_experiments.py
+"""
+
+import sys
+from pathlib import Path
+
+sys.path.insert(0, str(Path(__file__).resolve().parent.parent / "src"))
+
+
+# 实验配置列表
+EXPERIMENTS = [
+    {
+        "name": "exp1_baseline_transformer",
+        "description": "基线: 从头训练标准 Transformer (6层, d=512)",
+        "overrides": {
+            "model.type": "transformer_scratch",
+            "model.transformer.num_encoder_layers": 6,
+            "model.transformer.num_decoder_layers": 6,
+            "model.transformer.d_model": 512,
+            "model.transformer.use_flash_attention": False,
+            "model.transformer.use_rotary_embedding": False,
+        },
+    },
+    {
+        "name": "exp2_transformer_rope",
+        "description": "消融: Transformer + RoPE 位置编码",
+        "overrides": {
+            "model.type": "transformer_scratch",
+            "model.transformer.use_rotary_embedding": True,
+        },
+    },
+    {
+        "name": "exp3_transformer_flash_attn",
+        "description": "消融: Transformer + Flash Attention 2",
+        "overrides": {
+            "model.type": "transformer_scratch",
+            "model.transformer.use_flash_attention": True,
+        },
+    },
+    {
+        "name": "exp4_transformer_full",
+        "description": "完整: Transformer + RoPE + Flash Attention",
+        "overrides": {
+            "model.type": "transformer_scratch",
+            "model.transformer.use_flash_attention": True,
+            "model.transformer.use_rotary_embedding": True,
+        },
+    },
+    {
+        "name": "exp5_nllb_lora",
+        "description": "预训练微调: NLLB-600M + LoRA",
+        "overrides": {
+            "model.type": "finetune_nllb",
+            "model.pretrained.use_lora": True,
+            "model.pretrained.lora.r": 16,
+        },
+    },
+    {
+        "name": "exp6_nllb_full_finetune",
+        "description": "预训练全量微调: NLLB-600M",
+        "overrides": {
+            "model.type": "finetune_nllb",
+            "model.pretrained.use_lora": False,
+        },
+    },
+]
+
+
+def run_single_experiment(exp_config: dict):
+    """
+    运行单个实验。
+
+    TODO [Person E]:
+    1. 从 base config 创建实验配置
+    2. 应用 overrides
+    3. 设置 experiment.name 和 output_dir
+    4. 调用训练流程
+    5. 在测试集上评估
+    6. 保存实验结果
+    """
+    raise NotImplementedError("TODO: Person E 实现 run_single_experiment")
+
+
+def main():
+    """
+    运行所有实验。
+
+    TODO [Person E]:
+    1. 遍历 EXPERIMENTS 列表
+    2. 对每个实验调用 run_single_experiment
+    3. 收集所有实验结果
+    4. 生成对比表格
+    5. 保存汇总报告
+    """
+    print("=" * 60)
+    print("  EasyTranslate - Experiment Runner")
+    print("=" * 60)
+
+    for exp in EXPERIMENTS:
+        print(f"\n>>> Running: {exp['name']} - {exp['description']}")
+        # run_single_experiment(exp)
+
+    raise NotImplementedError("TODO: Person E 实现实验运行主流程")
+
+
+if __name__ == "__main__":
+    main()

scripts/train.py

@@ -1,76 +1,76 @@
-"""
-训练入口脚本
-
-使用方式:
-    # 从头训练 Transformer
-    python scripts/train.py --config configs/default_config.yaml
-
-    # 微调 NLLB
-    python scripts/train.py --config configs/default_config.yaml model.type=finetune_nllb
-
-    # 命令行覆盖参数
-    python scripts/train.py --config configs/default_config.yaml training.optimizer.lr=1e-4
-
-    # 分布式训练
-    torchrun --nproc_per_node=4 scripts/train.py --config configs/default_config.yaml
-"""
-
-import argparse
-import sys
-from pathlib import Path
-
-# 将 src 目录加入 Python path
-sys.path.insert(0, str(Path(__file__).resolve().parent.parent / "src"))
-
-
-def parse_args():
-    parser = argparse.ArgumentParser(description="EasyTranslate Training")
-    parser.add_argument("--config", type=str, default="configs/default_config.yaml", help="配置文件路径")
-    parser.add_argument("--resume", type=str, default=None, help="检查点路径 (断点续训)")
-    args, unknown = parser.parse_known_args()
-    return args, unknown
-
-
-def main():
-    """
-    训练主流程。
-
-    TODO [整合阶段 - 所有人协作]:
-    1. 解析命令行参数
-    2. 加载配置 (config_from_cli)
-    3. set_seed
-    4. setup_logging
-    5. 根据 model.type 选择训练模式:
-
-       A) model.type == "transformer_scratch":
-          a. 加载数据集 (load_wmt_dataset / load_opus_dataset)
-          b. 预处理 (preprocess_pipeline)
-          c. 训练分词器 (train_bpe_tokenizer) 或加载已有分词器
-          d. 构建 TranslationDataset + DataLoader
-          e. 构建 TransformerTranslationModel
-          f. 构建 Trainer 并开始训练
-
-       B) model.type == "finetune_nllb" / "finetune_mbart":
-          a. 加载预训练模型和分词器 (load_pretrained_model)
-          b. 配置 LoRA (setup_lora)
-          c. 加载数据集，使用预训练分词器处理
-          d. 构建 Trainer 并开始训练
-
-    6. 训练完成后保存最终模型
-    7. 在测试集上评估
-    """
-    args, cli_overrides = parse_args()
-
-    print("=" * 60)
-    print("  EasyTranslate - English to Chinese Translation")
-    print("=" * 60)
-
-    # TODO: 实现训练主流程
-    raise NotImplementedError(
-        "TODO: 整合阶段实现训练主流程\n"
-        "请在所有模块完成后，协作完成此脚本"
-    )
-
-
-if __name__ == "__main__":
-    main()
+"""
+训练入口脚本
+
+使用方式:
+    # 从头训练 Transformer
+    python scripts/train.py --config configs/default_config.yaml
+
+    # 微调 NLLB
+    python scripts/train.py --config configs/default_config.yaml model.type=finetune_nllb
+
+    # 命令行覆盖参数
+    python scripts/train.py --config configs/default_config.yaml training.optimizer.lr=1e-4
+
+    # 分布式训练
+    torchrun --nproc_per_node=4 scripts/train.py --config configs/default_config.yaml
+"""
+
+import argparse
+import sys
+from pathlib import Path
+
+# 将 src 目录加入 Python path
+sys.path.insert(0, str(Path(__file__).resolve().parent.parent / "src"))
+
+
+def parse_args():
+    parser = argparse.ArgumentParser(description="EasyTranslate Training")
+    parser.add_argument("--config", type=str, default="configs/default_config.yaml", help="配置文件路径")
+    parser.add_argument("--resume", type=str, default=None, help="检查点路径 (断点续训)")
+    args, unknown = parser.parse_known_args()
+    return args, unknown
+
+
+def main():
+    """
+    训练主流程。
+
+    TODO [整合阶段 - 所有人协作]:
+    1. 解析命令行参数
+    2. 加载配置 (config_from_cli)
+    3. set_seed
+    4. setup_logging
+    5. 根据 model.type 选择训练模式:
+
+       A) model.type == "transformer_scratch":
+          a. 加载数据集 (load_wmt_dataset / load_opus_dataset)
+          b. 预处理 (preprocess_pipeline)
+          c. 训练分词器 (train_bpe_tokenizer) 或加载已有分词器
+          d. 构建 TranslationDataset + DataLoader
+          e. 构建 TransformerTranslationModel
+          f. 构建 Trainer 并开始训练
+
+       B) model.type == "finetune_nllb" / "finetune_mbart":
+          a. 加载预训练模型和分词器 (load_pretrained_model)
+          b. 配置 LoRA (setup_lora)
+          c. 加载数据集，使用预训练分词器处理
+          d. 构建 Trainer 并开始训练
+
+    6. 训练完成后保存最终模型
+    7. 在测试集上评估
+    """
+    args, cli_overrides = parse_args()
+
+    print("=" * 60)
+    print("  EasyTranslate - English to Chinese Translation")
+    print("=" * 60)
+
+    # TODO: 实现训练主流程
+    raise NotImplementedError(
+        "TODO: 整合阶段实现训练主流程\n"
+        "请在所有模块完成后，协作完成此脚本"
+    )
+
+
+if __name__ == "__main__":
+    main()

scripts/translate.py

@@ -1,82 +1,316 @@
-"""
-交互式翻译推理脚本
-
-使用方式:
-    # 命令行交互翻译
-    python scripts/translate.py --checkpoint checkpoints/best_model.pt
-
-    # 翻译文件
-    python scripts/translate.py --checkpoint checkpoints/best_model.pt --input input.txt --output output.txt
-
-    # 启动 Gradio Web UI
-    python scripts/translate.py --checkpoint checkpoints/best_model.pt --web
-"""
-
-import argparse
-import sys
-from pathlib import Path
-
-sys.path.insert(0, str(Path(__file__).resolve().parent.parent / "src"))
-
-
-def parse_args():
-    parser = argparse.ArgumentParser(description="EasyTranslate Inference")
-    parser.add_argument("--config", type=str, default="configs/default_config.yaml")
-    parser.add_argument("--checkpoint", type=str, required=True)
-    parser.add_argument("--input", type=str, default=None, help="输入文件路径")
-    parser.add_argument("--output", type=str, default=None, help="输出文件路径")
-    parser.add_argument("--web", action="store_true", help="启动 Gradio Web UI")
-    return parser.parse_args()
-
-
-def interactive_translate(evaluator):
-    """
-    命令行交互翻译。
-
-    TODO [Person D]:
-    1. 循环读取用户输入
-    2. 调用 evaluator.translate_single()
-    3. 打印翻译结果
-    4. 输入 'quit' 退出
-    """
-    raise NotImplementedError("TODO: Person D 实现 interactive_translate")
-
-
-def translate_file(evaluator, input_path: str, output_path: str):
-    """
-    文件翻译。
-
-    TODO [Person D]:
-    1. 读取输入文件 (一行一句)
-    2. 批量翻译
-    3. 将结果写入输出文件
-    """
-    raise NotImplementedError("TODO: Person D 实现 translate_file")
-
-
-def launch_web_ui(evaluator):
-    """
-    启动 Gradio Web UI。
-
-    TODO [Person D]:
-    1. 创建 Gradio Interface
-    2. 输入: 英文文本框
-    3. 输出: 中文翻译结果
-    4. 调用 evaluator.translate_single()
-    """
-    raise NotImplementedError("TODO: Person D 实现 launch_web_ui")
-
-
-def main():
-    args = parse_args()
-
-    print("=" * 60)
-    print("  EasyTranslate - Translation")
-    print("=" * 60)
-
-    # TODO: 加载模型、构建 evaluator，然后根据参数选择模式
-    raise NotImplementedError("TODO: 实现推理主流程")
-
-
-if __name__ == "__main__":
-    main()
+"""
+交互式翻译推理脚本
+
+使用方式:
+    # 命令行交互翻译
+    python scripts/translate.py --checkpoint checkpoints/best_model.pt
+
+    # 翻译文件
+    python scripts/translate.py --checkpoint checkpoints/best_model.pt --input input.txt --output output.txt
+
+    # 启动 Streamlit Web UI（无需模型即可预览界面）
+    python scripts/translate.py --web
+    # 或者带模型启动翻译
+    python scripts/translate.py --checkpoint checkpoints/best_model.pt --web
+"""
+
+import argparse
+import subprocess
+import sys
+from pathlib import Path
+
+try:
+    from omegaconf import OmegaConf
+except ImportError:  # pragma: no cover
+    OmegaConf = None
+    import yaml
+
+sys.path.insert(0, str(Path(__file__).resolve().parent.parent / "src"))
+
+import torch
+
+from easytranslate.data.collator import TranslationCollator
+from easytranslate.data.dataset import (
+    TranslationDataset,
+    load_custom_dataset,
+    load_opus_dataset,
+    load_wmt_dataset,
+)
+from easytranslate.data.tokenizer import TokenizerWrapper, build_tokenizer
+from easytranslate.evaluation.evaluator import Evaluator
+from easytranslate.model import TransformerTranslationModel
+from easytranslate.model.finetune import load_pretrained_model
+
+
+def parse_args():
+    parser = argparse.ArgumentParser(description="EasyTranslate Inference")
+    parser.add_argument("--config", type=str, default="configs/default_config.yaml")
+    parser.add_argument("--checkpoint", type=str, default=None)
+    parser.add_argument("--input", type=str, default=None, help="输入文件路径")
+    parser.add_argument("--output", type=str, default=None, help="输出文件路径")
+    parser.add_argument("--web", action="store_true", help="启动 Streamlit Web UI")
+    parser.add_argument("--streamlit-app", action="store_true", help=argparse.SUPPRESS)
+    args, unknown = parser.parse_known_args()
+    return args, unknown
+
+
+def _get_config(config, *keys, default=None):
+    value = config
+    for key in keys:
+        if isinstance(value, dict):
+            value = value.get(key, default)
+        else:
+            value = getattr(value, key, default)
+        if value is default:
+            break
+    return value
+
+
+def _load_config(path, cli_overrides=None):
+    if OmegaConf is not None:
+        config = OmegaConf.load(path)
+        if cli_overrides:
+            config = OmegaConf.merge(config, OmegaConf.from_cli(cli_overrides))
+        return config
+
+    with open(path, "r", encoding="utf-8") as fin:
+        config = yaml.safe_load(fin)
+    if cli_overrides:
+        print("Warning: OmegaConf is not installed; CLI overrides are ignored.")
+    return config
+
+
+def interactive_translate(evaluator):
+    print("输入英⽂句⼦，按回车翻译；输入 'quit' 退出。")
+    while True:
+        try:
+            text = input("> ").strip()
+        except EOFError:
+            break
+        if not text:
+            continue
+        if text.lower() in {"quit", "exit"}:
+            break
+        translation = evaluator.translate_single(text)
+        print(translation)
+
+
+def translate_file(evaluator, input_path: str, output_path: str):
+    source_lines = []
+    with Path(input_path).open("r", encoding="utf-8") as fin:
+        for line in fin:
+            line = line.strip()
+            if line:
+                source_lines.append(line)
+
+    translations = evaluator.translate(source_lines)
+
+    output_file = Path(output_path)
+    output_file.parent.mkdir(parents=True, exist_ok=True)
+    with output_file.open("w", encoding="utf-8") as fout:
+        for line in translations:
+            fout.write(f"{line}\n")
+
+    print(f"Translation complete: {len(translations)} lines written to {output_path}")
+
+
+def launch_streamlit_app(args):
+    import streamlit as st
+
+    @st.cache_resource
+    def load_evaluator():
+        config = _load_config(args.config)
+        device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
+        model, tokenizer = build_model_and_tokenizer(config, device)
+        model = load_checkpoint(model, args.checkpoint, device)
+        return Evaluator(model, tokenizer, config)
+
+    st.set_page_config(page_title="EasyTranslate", layout="wide")
+    st.title("EasyTranslate")
+    st.write("English to Chinese translation powered by EasyTranslate.")
+
+    if args.checkpoint is None:
+        st.warning(
+            "当前未提供模型 checkpoint，页面仅用于预览界面效果。"
+            " 如需翻译，请传入 --checkpoint 或先训练生成模型。"
+        )
+
+    input_text = st.text_area("English Input", value="", height=200)
+    if st.button("Translate"):
+        if not input_text.strip():
+            st.warning("请输入要翻译的英文文本。")
+        elif args.checkpoint is None:
+            st.error("未提供 checkpoint，无法执行翻译。请使用 --checkpoint 参数启动。")
+        else:
+            with st.spinner("Translating..."):
+                try:
+                    evaluator = load_evaluator()
+                    translation = evaluator.translate_single(input_text.strip())
+                    st.text_area("Chinese Translation", value=translation, height=200)
+                except Exception as exc:
+                    st.error(f"模型加载或翻译失败：{exc}")
+
+    st.markdown("---")
+    st.caption("此页面用于展示前端界面；在未提供模型时，翻译功能会被禁用。")
+
+
+def launch_streamlit_process(args):
+    script_path = Path(__file__).resolve()
+    cmd = [sys.executable, "-m", "streamlit", "run", str(script_path), "--", "--streamlit-app", "--config", args.config]
+    if args.checkpoint:
+        cmd.extend(["--checkpoint", args.checkpoint])
+    if args.input:
+        cmd.extend(["--input", args.input])
+    if args.output:
+        cmd.extend(["--output", args.output])
+    subprocess.run(cmd, check=True)
+
+
+def _get_tokenizer_train_texts(config, allow_auto: bool = False) -> list[str] | None:
+    dataset_name = _get_config(config, "data", "dataset_name")
+    if dataset_name == "custom":
+        custom = _get_config(config, "data", "custom") or {}
+        data = load_custom_dataset(
+            train_src=custom.get("train_src"),
+            train_tgt=custom.get("train_tgt"),
+            val_src=custom.get("val_src"),
+            val_tgt=custom.get("val_tgt"),
+            test_src=custom.get("test_src"),
+            test_tgt=custom.get("test_tgt"),
+            preprocessing_config=_get_config(config, "data", "preprocessing"),
+        )
+        return list(data["train"]["src"]) + list(data["train"]["tgt"])
+
+    if not allow_auto:
+        return None
+
+    if dataset_name == "wmt":
+        try:
+            dataset = load_wmt_dataset(
+                year=_get_config(config, "data", "wmt", "year"),
+                language_pair=_get_config(config, "data", "wmt", "language_pair"),
+                split="train",
+            )
+        except Exception:
+            dataset = load_wmt_dataset(
+                year=_get_config(config, "data", "wmt", "year"),
+                language_pair=_get_config(config, "data", "wmt", "language_pair"),
+                split="validation",
+            )
+        return list(dataset["src"]) + list(dataset["tgt"])
+
+    if dataset_name == "opus":
+        try:
+            dataset = load_opus_dataset(
+                subset=_get_config(config, "data", "opus", "subset"),
+                split="train",
+            )
+        except Exception:
+            dataset = load_opus_dataset(
+                subset=_get_config(config, "data", "opus", "subset"),
+                split="validation",
+            )
+        return list(dataset["src"]) + list(dataset["tgt"])
+
+    return None
+
+
+def build_model_and_tokenizer(config, device):
+    model_type = _get_config(config, "model", "type")
+    if model_type == "transformer_scratch":
+        tokenizer_config = _get_config(config, "tokenizer") or {}
+        tokenizer_path = tokenizer_config.get("path") or tokenizer_config.get("tokenizer_path")
+        tokenizer_type = tokenizer_config.get("type", "bpe")
+        auto_train = bool(tokenizer_config.get("auto_train", False))
+        if tokenizer_type in {"bpe", "sentencepiece"} and not tokenizer_path:
+            train_texts = _get_tokenizer_train_texts(config, allow_auto=auto_train)
+            if train_texts is None:
+                raise ValueError(
+                    "BPE tokenizer requires tokenizer.path or a local custom dataset with train texts. "
+                    "Automatic download from WMT/OPUS is disabled by default. "
+                    "Set tokenizer.auto_train=true to enable it, or provide tokenizer.path/pretrained tokenizer."
+                )
+            tokenizer = build_tokenizer(tokenizer_config, train_texts=train_texts)
+        else:
+            tokenizer = build_tokenizer(tokenizer_config)
+
+        model = TransformerTranslationModel(
+            src_vocab_size=tokenizer.vocab_size,
+            tgt_vocab_size=tokenizer.vocab_size,
+            d_model=_get_config(config, "model", "transformer", "d_model"),
+            nhead=_get_config(config, "model", "transformer", "nhead"),
+            num_encoder_layers=_get_config(config, "model", "transformer", "num_encoder_layers"),
+            num_decoder_layers=_get_config(config, "model", "transformer", "num_decoder_layers"),
+            dim_feedforward=_get_config(config, "model", "transformer", "dim_feedforward"),
+            dropout=_get_config(config, "model", "transformer", "dropout"),
+            activation=_get_config(config, "model", "transformer", "activation"),
+            max_seq_len=_get_config(config, "model", "transformer", "max_seq_len"),
+            use_flash_attention=_get_config(config, "model", "transformer", "use_flash_attention"),
+            use_rotary_embedding=_get_config(config, "model", "transformer", "use_rotary_embedding"),
+            pre_norm=_get_config(config, "model", "transformer", "pre_norm"),
+            pad_id=tokenizer.pad_token_id,
+        )
+        return model.to(device), tokenizer
+
+    model, hf_tokenizer = load_pretrained_model(
+        config.model.pretrained.model_name,
+        config.model.pretrained.src_lang,
+        config.model.pretrained.tgt_lang,
+        device=str(device),
+    )
+    tokenizer = TokenizerWrapper(
+        hf_tokenizer,
+        pad_token=getattr(hf_tokenizer, "pad_token", "<pad>"),
+        unk_token=getattr(hf_tokenizer, "unk_token", "<unk>"),
+        bos_token=getattr(hf_tokenizer, "bos_token", "<s>"),
+        eos_token=getattr(hf_tokenizer, "eos_token", "</s>"),
+    )
+    return model, tokenizer
+
+
+def load_checkpoint(model, checkpoint_path, device):
+    checkpoint = torch.load(checkpoint_path, map_location=device)
+    if isinstance(checkpoint, dict):
+        if "model_state_dict" in checkpoint:
+            model.load_state_dict(checkpoint["model_state_dict"])
+        elif "state_dict" in checkpoint:
+            model.load_state_dict(checkpoint["state_dict"])
+        else:
+            try:
+                model.load_state_dict(checkpoint)
+            except Exception as exc:
+                raise ValueError("Checkpoint does not contain a valid model state dict") from exc
+    else:
+        raise ValueError("Unsupported checkpoint format")
+    return model
+
+
+def main():
+    args, cli_overrides = parse_args()
+
+    print("=" * 60)
+    print("  EasyTranslate - Translation")
+    print("=" * 60)
+
+    if args.web and not args.streamlit_app:
+        launch_streamlit_process(args)
+        return
+
+    if args.streamlit_app:
+        launch_streamlit_app(args)
+        return
+
+    config = _load_config(args.config, cli_overrides)
+    device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
+    model, tokenizer = build_model_and_tokenizer(config, device)
+    model = load_checkpoint(model, args.checkpoint, device)
+    evaluator = Evaluator(model, tokenizer, config)
+
+    if args.input and args.output:
+        translate_file(evaluator, args.input, args.output)
+    else:
+        interactive_translate(evaluator)
+
+
+if __name__ == "__main__":
+    main()

scripts/visualize.py

@@ -1,86 +1,86 @@
-"""
-可视化与分析脚本 — Person E 负责
-
-功能:
-1. 训练曲线绘制 (loss, BLEU, learning rate)
-2. 实验结果对比图
-3. 注意力权重可视化
-4. 翻译样例展示
-"""
-
-import sys
-from pathlib import Path
-
-sys.path.insert(0, str(Path(__file__).resolve().parent.parent / "src"))
-
-
-def plot_training_curves(log_dir: str, output_path: str = "outputs/training_curves.png"):
-    """
-    绘制训练曲线。
-
-    TODO [Person E]:
-    1. 从 TensorBoard 日志或 CSV 读取训练数据
-    2. 绘制子图:
-       - Train Loss vs Steps
-       - Val Loss vs Steps
-       - BLEU vs Epochs
-       - Learning Rate vs Steps
-    3. 保存图片
-    """
-    raise NotImplementedError("TODO: Person E 实现 plot_training_curves")
-
-
-def plot_experiment_comparison(results_dir: str, output_path: str = "outputs/experiment_comparison.png"):
-    """
-    绘制实验对比图。
-
-    TODO [Person E]:
-    1. 读取所有实验的评估结果
-    2. 绘制柱状图: BLEU / COMET / chrF 对比
-    3. 绘制表格: 所有指标汇总
-    4. 保存图片
-    """
-    raise NotImplementedError("TODO: Person E 实现 plot_experiment_comparison")
-
-
-def visualize_attention(
-    model,
-    src_text: str,
-    tgt_text: str,
-    tokenizer,
-    output_path: str = "outputs/attention_map.png",
-):
-    """
-    注意力权重可视化。
-
-    TODO [Person E]:
-    1. 获取模型的 encoder self-attention 和 cross-attention 权重
-    2. 绘制热力图 (matplotlib / seaborn)
-    3. x 轴: 源语言 tokens, y 轴: 目标语言 tokens
-    4. 支持多头注意力的分别可视化和平均可视化
-    """
-    raise NotImplementedError("TODO: Person E 实现 visualize_attention")
-
-
-def generate_translation_examples(
-    evaluator,
-    test_pairs: list[tuple[str, str]],
-    output_path: str = "outputs/translation_examples.md",
-):
-    """
-    生成翻译样例展示。
-
-    TODO [Person E]:
-    1. 翻译测试样例
-    2. 生成 Markdown 格式的对比表格:
-       | 源文 (英) | 参考翻译 (中) | 模型翻译 (中) | BLEU |
-    3. 包含好的和差的翻译案例
-    4. 保存为 Markdown 文件
-    """
-    raise NotImplementedError("TODO: Person E 实现 generate_translation_examples")
-
-
-if __name__ == "__main__":
-    print("请指定要运行的可视化任务")
-    print("  python scripts/visualize.py --task training_curves --log_dir logs/")
-    print("  python scripts/visualize.py --task comparison --results_dir outputs/")
+"""
+可视化与分析脚本 — Person E 负责
+
+功能:
+1. 训练曲线绘制 (loss, BLEU, learning rate)
+2. 实验结果对比图
+3. 注意力权重可视化
+4. 翻译样例展示
+"""
+
+import sys
+from pathlib import Path
+
+sys.path.insert(0, str(Path(__file__).resolve().parent.parent / "src"))
+
+
+def plot_training_curves(log_dir: str, output_path: str = "outputs/training_curves.png"):
+    """
+    绘制训练曲线。
+
+    TODO [Person E]:
+    1. 从 TensorBoard 日志或 CSV 读取训练数据
+    2. 绘制子图:
+       - Train Loss vs Steps
+       - Val Loss vs Steps
+       - BLEU vs Epochs
+       - Learning Rate vs Steps
+    3. 保存图片
+    """
+    raise NotImplementedError("TODO: Person E 实现 plot_training_curves")
+
+
+def plot_experiment_comparison(results_dir: str, output_path: str = "outputs/experiment_comparison.png"):
+    """
+    绘制实验对比图。
+
+    TODO [Person E]:
+    1. 读取所有实验的评估结果
+    2. 绘制柱状图: BLEU / COMET / chrF 对比
+    3. 绘制表格: 所有指标汇总
+    4. 保存图片
+    """
+    raise NotImplementedError("TODO: Person E 实现 plot_experiment_comparison")
+
+
+def visualize_attention(
+    model,
+    src_text: str,
+    tgt_text: str,
+    tokenizer,
+    output_path: str = "outputs/attention_map.png",
+):
+    """
+    注意力权重可视化。
+
+    TODO [Person E]:
+    1. 获取模型的 encoder self-attention 和 cross-attention 权重
+    2. 绘制热力图 (matplotlib / seaborn)
+    3. x 轴: 源语言 tokens, y 轴: 目标语言 tokens
+    4. 支持多头注意力的分别可视化和平均可视化
+    """
+    raise NotImplementedError("TODO: Person E 实现 visualize_attention")
+
+
+def generate_translation_examples(
+    evaluator,
+    test_pairs: list[tuple[str, str]],
+    output_path: str = "outputs/translation_examples.md",
+):
+    """
+    生成翻译样例展示。
+
+    TODO [Person E]:
+    1. 翻译测试样例
+    2. 生成 Markdown 格式的对比表格:
+       | 源文 (英) | 参考翻译 (中) | 模型翻译 (中) | BLEU |
+    3. 包含好的和差的翻译案例
+    4. 保存为 Markdown 文件
+    """
+    raise NotImplementedError("TODO: Person E 实现 generate_translation_examples")
+
+
+if __name__ == "__main__":
+    print("请指定要运行的可视化任务")
+    print("  python scripts/visualize.py --task training_curves --log_dir logs/")
+    print("  python scripts/visualize.py --task comparison --results_dir outputs/")

setup.py

@@ -1,23 +1,23 @@
-from setuptools import setup, find_packages
-
-setup(
-    name="easytranslate",
-    version="0.1.0",
-    description="Transformer-based English-to-Chinese Translation Model",
-    packages=find_packages(where="src"),
-    package_dir={"": "src"},
-    python_requires=">=3.10",
-    install_requires=[
-        "torch>=2.1.0",
-        "transformers>=4.36.0",
-        "datasets>=2.16.0",
-        "tokenizers>=0.15.0",
-        "sentencepiece>=0.1.99",
-        "accelerate>=0.25.0",
-        "sacrebleu>=2.4.0",
-        "numpy>=1.24.0,<2.0",
-        "omegaconf>=2.3.0",
-        "rich>=13.0.0",
-        "tqdm>=4.66.0",
-    ],
-)
+from setuptools import setup, find_packages
+
+setup(
+    name="easytranslate",
+    version="0.1.0",
+    description="Transformer-based English-to-Chinese Translation Model",
+    packages=find_packages(where="src"),
+    package_dir={"": "src"},
+    python_requires=">=3.10",
+    install_requires=[
+        "torch>=2.1.0",
+        "transformers>=4.36.0",
+        "datasets>=2.16.0",
+        "tokenizers>=0.15.0",
+        "sentencepiece>=0.1.99",
+        "accelerate>=0.25.0",
+        "sacrebleu>=2.4.0",
+        "numpy>=1.24.0,<2.0",
+        "omegaconf>=2.3.0",
+        "rich>=13.0.0",
+        "tqdm>=4.66.0",
+    ],
+)

src/easytranslate/__init__.py

@@ -1,3 +1,3 @@
-"""EasyTranslate: Transformer-based English-to-Chinese Translation."""
-
-__version__ = "0.1.0"
+"""EasyTranslate: Transformer-based English-to-Chinese Translation."""
+
+__version__ = "0.1.0"

src/easytranslate/data/README.md

@@ -1,126 +1,126 @@
-# EasyTranslate Data 模块使用说明
-
-本目录负责英中翻译任务的数据加载、清洗、分词、样本构造和动态批处理。推荐主数据集使用 WMT19 zh-en，调试或小样本实验可使用 OPUS-100 en-zh。
-
-## 文件职责
-
-- `dataset.py`: 加载 WMT/OPUS/custom 数据，并提供 `TranslationDataset`。
-- `tokenizer.py`: 训练 BPE tokenizer，并用 `TokenizerWrapper` 统一 tokenizer 接口。
-- `preprocessing.py`: 文本标准化、长度过滤、去重。
-- `collator.py`: batch padding、attention mask、动态 token batch。
-
-## 推荐数据集
-
-正式实验建议使用：
-
-```python
-from easytranslate.data import load_wmt_dataset
-
-raw = load_wmt_dataset(year="19", language_pair="zh-en")
-```
-
-快速调试建议使用：
-
-```python
-from easytranslate.data import load_opus_dataset
-
-raw = load_opus_dataset(subset="en-zh")
-```
-
-两个加载函数都会把样本统一成：
-
-```python
-{"src": "English sentence", "tgt": "中文句子"}
-```
-
-## 从文本到 DataLoader
-
-```python
-from torch.utils.data import DataLoader
-
-from easytranslate.data import (
-    DynamicBatchSampler,
-    TranslationCollator,
-    TranslationDataset,
-    preprocess_pipeline,
-    train_bpe_tokenizer,
-)
-
-src_texts = raw["train"]["src"]
-tgt_texts = raw["train"]["tgt"]
-
-src_texts, tgt_texts = preprocess_pipeline(
-    src_texts,
-    tgt_texts,
-    lowercase_src=False,
-    remove_punctuation=False,
-    max_src_len=256,
-    max_tgt_len=256,
-    length_ratio_threshold=3.0,
-)
-
-tokenizer = train_bpe_tokenizer(
-    list(src_texts) + list(tgt_texts),
-    vocab_size=32000,
-    min_frequency=2,
-    save_path="outputs/tokenizer/bpe.json",
-)
-
-train_dataset = TranslationDataset(
-    src_texts,
-    tgt_texts,
-    tokenizer=tokenizer,
-    max_src_len=256,
-    max_tgt_len=256,
-)
-
-lengths = [
-    (len(tokenizer.encode(src, add_special_tokens=True)), len(tokenizer.encode(tgt, add_special_tokens=True)))
-    for src, tgt in zip(src_texts, tgt_texts)
-]
-
-batch_sampler = DynamicBatchSampler(lengths, max_tokens_per_batch=8192)
-collator = TranslationCollator(pad_token_id=tokenizer.pad_token_id)
-
-loader = DataLoader(
-    train_dataset,
-    batch_sampler=batch_sampler,
-    collate_fn=collator,
-    num_workers=4,
-    pin_memory=True,
-)
-```
-
-## Batch 字段
-
-`TranslationCollator` 输出：
-
-- `src_ids`: `[B, S]`
-- `tgt_input_ids`: `[B, T]`，以 `<s>` 开头，用于 teacher forcing
-- `labels`: `[B, T]`，以 `</s>` 结尾，padding 为 `-100`
-- `src_padding_mask`: `[B, S]`，padding 位置为 `True`
-- `tgt_padding_mask`: `[B, T]`，padding 位置为 `True`
-- `src_attention_mask` / `tgt_attention_mask`: 有效 token 为 `1`
-- `src_lens` / `tgt_lens`: 原始长度
-
-## 自定义平行语料
-
-```python
-from easytranslate.data import load_custom_dataset
-
-data = load_custom_dataset(
-    train_src="data/train.en",
-    train_tgt="data/train.zh",
-    val_src="data/val.en",
-    val_tgt="data/val.zh",
-    test_src="data/test.en",
-    test_tgt="data/test.zh",
-)
-```
-
-## 处理原则
-
-- 保留英文大小写和中英文标点，默认不做 lowercase、不去标点。
-- 清洗只做 Unicode NFKC、控制字符删除、空白合并。
-- tokenizer 只用训练集训练，不使用验证集或测试集。
-- 从零训练 Transformer 时建议使用共享 bilingual BPE；微调 NLLB 时直接使用 NLLB tokenizer。
+# EasyTranslate Data 模块使用说明
+
+本目录负责英中翻译任务的数据加载、清洗、分词、样本构造和动态批处理。推荐主数据集使用 WMT19 zh-en，调试或小样本实验可使用 OPUS-100 en-zh。
+
+## 文件职责
+
+- `dataset.py`: 加载 WMT/OPUS/custom 数据，并提供 `TranslationDataset`。
+- `tokenizer.py`: 训练 BPE tokenizer，并用 `TokenizerWrapper` 统一 tokenizer 接口。
+- `preprocessing.py`: 文本标准化、长度过滤、去重。
+- `collator.py`: batch padding、attention mask、动态 token batch。
+
+## 推荐数据集
+
+正式实验建议使用：
+
+```python
+from easytranslate.data import load_wmt_dataset
+
+raw = load_wmt_dataset(year="19", language_pair="zh-en")
+```
+
+快速调试建议使用：
+
+```python
+from easytranslate.data import load_opus_dataset
+
+raw = load_opus_dataset(subset="en-zh")
+```
+
+两个加载函数都会把样本统一成：
+
+```python
+{"src": "English sentence", "tgt": "中文句子"}
+```
+
+## 从文本到 DataLoader
+
+```python
+from torch.utils.data import DataLoader
+
+from easytranslate.data import (
+    DynamicBatchSampler,
+    TranslationCollator,
+    TranslationDataset,
+    preprocess_pipeline,
+    train_bpe_tokenizer,
+)
+
+src_texts = raw["train"]["src"]
+tgt_texts = raw["train"]["tgt"]
+
+src_texts, tgt_texts = preprocess_pipeline(
+    src_texts,
+    tgt_texts,
+    lowercase_src=False,
+    remove_punctuation=False,
+    max_src_len=256,
+    max_tgt_len=256,
+    length_ratio_threshold=3.0,
+)
+
+tokenizer = train_bpe_tokenizer(
+    list(src_texts) + list(tgt_texts),
+    vocab_size=32000,
+    min_frequency=2,
+    save_path="outputs/tokenizer/bpe.json",
+)
+
+train_dataset = TranslationDataset(
+    src_texts,
+    tgt_texts,
+    tokenizer=tokenizer,
+    max_src_len=256,
+    max_tgt_len=256,
+)
+
+lengths = [
+    (len(tokenizer.encode(src, add_special_tokens=True)), len(tokenizer.encode(tgt, add_special_tokens=True)))
+    for src, tgt in zip(src_texts, tgt_texts)
+]
+
+batch_sampler = DynamicBatchSampler(lengths, max_tokens_per_batch=8192)
+collator = TranslationCollator(pad_token_id=tokenizer.pad_token_id)
+
+loader = DataLoader(
+    train_dataset,
+    batch_sampler=batch_sampler,
+    collate_fn=collator,
+    num_workers=4,
+    pin_memory=True,
+)
+```
+
+## Batch 字段
+
+`TranslationCollator` 输出：
+
+- `src_ids`: `[B, S]`
+- `tgt_input_ids`: `[B, T]`，以 `<s>` 开头，用于 teacher forcing
+- `labels`: `[B, T]`，以 `</s>` 结尾，padding 为 `-100`
+- `src_padding_mask`: `[B, S]`，padding 位置为 `True`
+- `tgt_padding_mask`: `[B, T]`，padding 位置为 `True`
+- `src_attention_mask` / `tgt_attention_mask`: 有效 token 为 `1`
+- `src_lens` / `tgt_lens`: 原始长度
+
+## 自定义平行语料
+
+```python
+from easytranslate.data import load_custom_dataset
+
+data = load_custom_dataset(
+    train_src="data/train.en",
+    train_tgt="data/train.zh",
+    val_src="data/val.en",
+    val_tgt="data/val.zh",
+    test_src="data/test.en",
+    test_tgt="data/test.zh",
+)
+```
+
+## 处理原则
+
+- 保留英文大小写和中英文标点，默认不做 lowercase、不去标点。
+- 清洗只做 Unicode NFKC、控制字符删除、空白合并。
+- tokenizer 只用训练集训练，不使用验证集或测试集。
+- 从零训练 Transformer 时建议使用共享 bilingual BPE；微调 NLLB 时直接使用 NLLB tokenizer。

src/easytranslate/data/__init__.py

@@ -1,32 +1,32 @@
-"""Data utilities for EasyTranslate."""
-
-from easytranslate.data.collator import DynamicBatchSampler, TranslationCollator
-from easytranslate.data.dataset import (
-    TranslationDataset,
-    load_custom_dataset,
-    load_opus_dataset,
-    load_wmt_dataset,
-)
-from easytranslate.data.preprocessing import (
-    clean_text,
-    deduplicate_pairs,
-    filter_by_length,
-    preprocess_pipeline,
-)
-from easytranslate.data.tokenizer import TokenizerWrapper, build_tokenizer, train_bpe_tokenizer
-
-__all__ = [
-    "DynamicBatchSampler",
-    "TokenizerWrapper",
-    "TranslationCollator",
-    "TranslationDataset",
-    "build_tokenizer",
-    "clean_text",
-    "deduplicate_pairs",
-    "filter_by_length",
-    "load_custom_dataset",
-    "load_opus_dataset",
-    "load_wmt_dataset",
-    "preprocess_pipeline",
-    "train_bpe_tokenizer",
-]
+"""Data utilities for EasyTranslate."""
+
+from easytranslate.data.collator import DynamicBatchSampler, TranslationCollator
+from easytranslate.data.dataset import (
+    TranslationDataset,
+    load_custom_dataset,
+    load_opus_dataset,
+    load_wmt_dataset,
+)
+from easytranslate.data.preprocessing import (
+    clean_text,
+    deduplicate_pairs,
+    filter_by_length,
+    preprocess_pipeline,
+)
+from easytranslate.data.tokenizer import TokenizerWrapper, build_tokenizer, train_bpe_tokenizer
+
+__all__ = [
+    "DynamicBatchSampler",
+    "TokenizerWrapper",
+    "TranslationCollator",
+    "TranslationDataset",
+    "build_tokenizer",
+    "clean_text",
+    "deduplicate_pairs",
+    "filter_by_length",
+    "load_custom_dataset",
+    "load_opus_dataset",
+    "load_wmt_dataset",
+    "preprocess_pipeline",
+    "train_bpe_tokenizer",
+]

src/easytranslate/data/collator.py

@@ -1,114 +1,114 @@
-"""Batch collation and dynamic batching for translation training."""
-
-from __future__ import annotations
-
-import random
-from typing import Iterator, Sequence
-
-import torch
-from torch.nn.utils.rnn import pad_sequence
-from torch.utils.data import Sampler
-
-
-class TranslationCollator:
-    """Pad variable-length translation examples into one batch."""
-
-    def __init__(self, pad_token_id: int = 0, label_pad_token_id: int = -100):
-        self.pad_token_id = pad_token_id
-        self.label_pad_token_id = label_pad_token_id
-
-    def __call__(self, batch: Sequence[dict[str, torch.Tensor]]) -> dict[str, torch.Tensor]:
-        src_ids = pad_sequence(
-            [item["src_ids"] for item in batch],
-            batch_first=True,
-            padding_value=self.pad_token_id,
-        )
-        tgt_input_ids = pad_sequence(
-            [item["tgt_input_ids"] for item in batch],
-            batch_first=True,
-            padding_value=self.pad_token_id,
-        )
-        labels = pad_sequence(
-            [item["labels"] for item in batch],
-            batch_first=True,
-            padding_value=self.label_pad_token_id,
-        )
-
-        src_padding_mask = src_ids.eq(self.pad_token_id)
-        tgt_padding_mask = tgt_input_ids.eq(self.pad_token_id)
-
-        return {
-            "src_ids": src_ids,
-            "tgt_input_ids": tgt_input_ids,
-            "labels": labels,
-            "src_padding_mask": src_padding_mask,
-            "tgt_padding_mask": tgt_padding_mask,
-            "src_attention_mask": (~src_padding_mask).long(),
-            "tgt_attention_mask": (~tgt_padding_mask).long(),
-            "src_lens": torch.stack([item["src_len"] for item in batch]),
-            "tgt_lens": torch.stack([item["tgt_len"] for item in batch]),
-        }
-
-
-class DynamicBatchSampler(Sampler[list[int]]):
-    """Create batches constrained by an approximate max token budget."""
-
-    def __init__(
-        self,
-        lengths: Sequence[int | tuple[int, int]],
-        max_tokens_per_batch: int = 8192,
-        shuffle: bool = True,
-        drop_last: bool = False,
-    ):
-        self.lengths = [max(length) if isinstance(length, tuple) else int(length) for length in lengths]
-        self.max_tokens_per_batch = max_tokens_per_batch
-        self.shuffle = shuffle
-        self.drop_last = drop_last
-
-    def __iter__(self) -> Iterator[list[int]]:
-        indices = list(range(len(self.lengths)))
-        if self.shuffle:
-            random.shuffle(indices)
-
-        indices.sort(key=lambda idx: self.lengths[idx])
-        batches: list[list[int]] = []
-        batch: list[int] = []
-        max_len = 0
-
-        for idx in indices:
-            candidate_max_len = max(max_len, self.lengths[idx])
-            candidate_tokens = candidate_max_len * (len(batch) + 1)
-
-            if batch and candidate_tokens > self.max_tokens_per_batch:
-                batches.append(batch)
-                batch = []
-                max_len = 0
-
-            batch.append(idx)
-            max_len = max(max_len, self.lengths[idx])
-
-        if batch and not self.drop_last:
-            batches.append(batch)
-
-        if self.shuffle:
-            random.shuffle(batches)
-
-        yield from batches
-
-    def __len__(self) -> int:
-        count = 0
-        batch_size = 0
-        max_len = 0
-
-        for length in sorted(self.lengths):
-            candidate_max_len = max(max_len, length)
-            if batch_size and candidate_max_len * (batch_size + 1) > self.max_tokens_per_batch:
-                count += 1
-                batch_size = 0
-                max_len = 0
-            batch_size += 1
-            max_len = max(max_len, length)
-
-        if batch_size and not self.drop_last:
-            count += 1
-        return count
+"""Batch collation and dynamic batching for translation training."""
+
+from __future__ import annotations
+
+import random
+from typing import Iterator, Sequence
+
+import torch
+from torch.nn.utils.rnn import pad_sequence
+from torch.utils.data import Sampler
+
+
+class TranslationCollator:
+    """Pad variable-length translation examples into one batch."""
+
+    def __init__(self, pad_token_id: int = 0, label_pad_token_id: int = -100):
+        self.pad_token_id = pad_token_id
+        self.label_pad_token_id = label_pad_token_id
+
+    def __call__(self, batch: Sequence[dict[str, torch.Tensor]]) -> dict[str, torch.Tensor]:
+        src_ids = pad_sequence(
+            [item["src_ids"] for item in batch],
+            batch_first=True,
+            padding_value=self.pad_token_id,
+        )
+        tgt_input_ids = pad_sequence(
+            [item["tgt_input_ids"] for item in batch],
+            batch_first=True,
+            padding_value=self.pad_token_id,
+        )
+        labels = pad_sequence(
+            [item["labels"] for item in batch],
+            batch_first=True,
+            padding_value=self.label_pad_token_id,
+        )
+
+        src_padding_mask = src_ids.eq(self.pad_token_id)
+        tgt_padding_mask = tgt_input_ids.eq(self.pad_token_id)
+
+        return {
+            "src_ids": src_ids,
+            "tgt_input_ids": tgt_input_ids,
+            "labels": labels,
+            "src_padding_mask": src_padding_mask,
+            "tgt_padding_mask": tgt_padding_mask,
+            "src_attention_mask": (~src_padding_mask).long(),
+            "tgt_attention_mask": (~tgt_padding_mask).long(),
+            "src_lens": torch.stack([item["src_len"] for item in batch]),
+            "tgt_lens": torch.stack([item["tgt_len"] for item in batch]),
+        }
+
+
+class DynamicBatchSampler(Sampler[list[int]]):
+    """Create batches constrained by an approximate max token budget."""
+
+    def __init__(
+        self,
+        lengths: Sequence[int | tuple[int, int]],
+        max_tokens_per_batch: int = 8192,
+        shuffle: bool = True,
+        drop_last: bool = False,
+    ):
+        self.lengths = [max(length) if isinstance(length, tuple) else int(length) for length in lengths]
+        self.max_tokens_per_batch = max_tokens_per_batch
+        self.shuffle = shuffle
+        self.drop_last = drop_last
+
+    def __iter__(self) -> Iterator[list[int]]:
+        indices = list(range(len(self.lengths)))
+        if self.shuffle:
+            random.shuffle(indices)
+
+        indices.sort(key=lambda idx: self.lengths[idx])
+        batches: list[list[int]] = []
+        batch: list[int] = []
+        max_len = 0
+
+        for idx in indices:
+            candidate_max_len = max(max_len, self.lengths[idx])
+            candidate_tokens = candidate_max_len * (len(batch) + 1)
+
+            if batch and candidate_tokens > self.max_tokens_per_batch:
+                batches.append(batch)
+                batch = []
+                max_len = 0
+
+            batch.append(idx)
+            max_len = max(max_len, self.lengths[idx])
+
+        if batch and not self.drop_last:
+            batches.append(batch)
+
+        if self.shuffle:
+            random.shuffle(batches)
+
+        yield from batches
+
+    def __len__(self) -> int:
+        count = 0
+        batch_size = 0
+        max_len = 0
+
+        for length in sorted(self.lengths):
+            candidate_max_len = max(max_len, length)
+            if batch_size and candidate_max_len * (batch_size + 1) > self.max_tokens_per_batch:
+                count += 1
+                batch_size = 0
+                max_len = 0
+            batch_size += 1
+            max_len = max(max_len, length)
+
+        if batch_size and not self.drop_last:
+            count += 1
+        return count

src/easytranslate/data/dataset.py

@@ -1,157 +1,157 @@
-"""Dataset loading and PyTorch dataset classes for English-Chinese translation."""
-
-from __future__ import annotations
-
-from pathlib import Path
-from typing import Mapping, Sequence
-
-import torch
-from torch.utils.data import Dataset
-
-from easytranslate.data.preprocessing import preprocess_pipeline
-from easytranslate.data.tokenizer import TokenizerWrapper
-
-
-class TranslationDataset(Dataset):
-    """PyTorch dataset that builds seq2seq inputs for teacher forcing."""
-
-    def __init__(
-        self,
-        src_texts: Sequence[str],
-        tgt_texts: Sequence[str],
-        tokenizer: TokenizerWrapper | None = None,
-        src_tokenizer: TokenizerWrapper | None = None,
-        tgt_tokenizer: TokenizerWrapper | None = None,
-        max_src_len: int = 256,
-        max_tgt_len: int = 256,
-    ):
-        if len(src_texts) != len(tgt_texts):
-            raise ValueError("src_texts and tgt_texts must have the same length")
-
-        if tokenizer is not None:
-            src_tokenizer = src_tokenizer or tokenizer
-            tgt_tokenizer = tgt_tokenizer or tokenizer
-        if src_tokenizer is None or tgt_tokenizer is None:
-            raise ValueError("Provide tokenizer or both src_tokenizer and tgt_tokenizer")
-
-        self.src_texts = list(src_texts)
-        self.tgt_texts = list(tgt_texts)
-        self.src_tokenizer = src_tokenizer
-        self.tgt_tokenizer = tgt_tokenizer
-        self.max_src_len = max_src_len
-        self.max_tgt_len = max_tgt_len
-
-    def __len__(self) -> int:
-        return len(self.src_texts)
-
-    def __getitem__(self, index: int) -> dict[str, torch.Tensor]:
-        src_ids = self.src_tokenizer.encode(
-            self.src_texts[index],
-            add_special_tokens=True,
-            max_length=self.max_src_len,
-        )
-
-        target_core = self.tgt_tokenizer.encode(
-            self.tgt_texts[index],
-            add_special_tokens=False,
-            max_length=max(1, self.max_tgt_len - 1),
-        )
-        tgt_input_ids = [self.tgt_tokenizer.bos_token_id] + target_core
-        labels = target_core + [self.tgt_tokenizer.eos_token_id]
-
-        return {
-            "src_ids": torch.tensor(src_ids, dtype=torch.long),
-            "tgt_input_ids": torch.tensor(tgt_input_ids, dtype=torch.long),
-            "labels": torch.tensor(labels, dtype=torch.long),
-            "src_len": torch.tensor(len(src_ids), dtype=torch.long),
-            "tgt_len": torch.tensor(len(labels), dtype=torch.long),
-        }
-
-
-def _translation_to_columns(dataset, src_lang: str, tgt_lang: str):
-    def convert(example):
-        translation = example["translation"]
-        return {"src": translation[src_lang], "tgt": translation[tgt_lang]}
-
-    return dataset.map(convert, remove_columns=dataset.column_names)
-
-
-def load_wmt_dataset(
-    year: str = "19",
-    language_pair: str = "zh-en",
-    src_lang: str = "en",
-    tgt_lang: str = "zh",
-    split: str | None = None,
-    cache_dir: str | None = None,
-):
-    """Load WMT zh-en and normalize rows to {'src', 'tgt'}."""
-    from datasets import load_dataset
-
-    dataset_name = f"wmt/wmt{year}"
-    try:
-        dataset = load_dataset(dataset_name, language_pair, split=split, cache_dir=cache_dir)
-    except Exception:
-        dataset = load_dataset(f"wmt{year}", language_pair, split=split, cache_dir=cache_dir)
-
-    if split is not None:
-        return _translation_to_columns(dataset, src_lang, tgt_lang)
-
-    return dataset.map(
-        lambda example: {"src": example["translation"][src_lang], "tgt": example["translation"][tgt_lang]},
-        remove_columns=next(iter(dataset.values())).column_names,
-    )
-
-
-def load_opus_dataset(
-    subset: str = "en-zh",
-    src_lang: str = "en",
-    tgt_lang: str = "zh",
-    split: str | None = None,
-    cache_dir: str | None = None,
-):
-    """Load OPUS-100 and normalize rows to {'src', 'tgt'}."""
-    from datasets import load_dataset
-
-    dataset = load_dataset("Helsinki-NLP/opus-100", subset, split=split, cache_dir=cache_dir)
-    if split is not None:
-        return _translation_to_columns(dataset, src_lang, tgt_lang)
-
-    return dataset.map(
-        lambda example: {"src": example["translation"][src_lang], "tgt": example["translation"][tgt_lang]},
-        remove_columns=next(iter(dataset.values())).column_names,
-    )
-
-
-def _read_lines(path: str | Path) -> list[str]:
-    with Path(path).open("r", encoding="utf-8") as f:
-        return [line.rstrip("\n") for line in f]
-
-
-def load_custom_dataset(
-    train_src: str | Path,
-    train_tgt: str | Path,
-    val_src: str | Path | None = None,
-    val_tgt: str | Path | None = None,
-    test_src: str | Path | None = None,
-    test_tgt: str | Path | None = None,
-    preprocessing_config: Mapping | None = None,
-) -> dict[str, dict[str, list[str]]]:
-    """Load parallel text files and return split dictionaries."""
-    preprocessing_config = dict(preprocessing_config or {})
-
-    def load_split(src_path: str | Path, tgt_path: str | Path) -> dict[str, list[str]]:
-        src_texts, tgt_texts = preprocess_pipeline(
-            _read_lines(src_path),
-            _read_lines(tgt_path),
-            **preprocessing_config,
-        )
-        return {"src": src_texts, "tgt": tgt_texts}
-
-    result = {"train": load_split(train_src, train_tgt)}
-
-    if val_src and val_tgt:
-        result["validation"] = load_split(val_src, val_tgt)
-    if test_src and test_tgt:
-        result["test"] = load_split(test_src, test_tgt)
-
-    return result
+"""Dataset loading and PyTorch dataset classes for English-Chinese translation."""
+
+from __future__ import annotations
+
+from pathlib import Path
+from typing import Mapping, Sequence
+
+import torch
+from torch.utils.data import Dataset
+
+from easytranslate.data.preprocessing import preprocess_pipeline
+from easytranslate.data.tokenizer import TokenizerWrapper
+
+
+class TranslationDataset(Dataset):
+    """PyTorch dataset that builds seq2seq inputs for teacher forcing."""
+
+    def __init__(
+        self,
+        src_texts: Sequence[str],
+        tgt_texts: Sequence[str],
+        tokenizer: TokenizerWrapper | None = None,
+        src_tokenizer: TokenizerWrapper | None = None,
+        tgt_tokenizer: TokenizerWrapper | None = None,
+        max_src_len: int = 256,
+        max_tgt_len: int = 256,
+    ):
+        if len(src_texts) != len(tgt_texts):
+            raise ValueError("src_texts and tgt_texts must have the same length")
+
+        if tokenizer is not None:
+            src_tokenizer = src_tokenizer or tokenizer
+            tgt_tokenizer = tgt_tokenizer or tokenizer
+        if src_tokenizer is None or tgt_tokenizer is None:
+            raise ValueError("Provide tokenizer or both src_tokenizer and tgt_tokenizer")
+
+        self.src_texts = list(src_texts)
+        self.tgt_texts = list(tgt_texts)
+        self.src_tokenizer = src_tokenizer
+        self.tgt_tokenizer = tgt_tokenizer
+        self.max_src_len = max_src_len
+        self.max_tgt_len = max_tgt_len
+
+    def __len__(self) -> int:
+        return len(self.src_texts)
+
+    def __getitem__(self, index: int) -> dict[str, torch.Tensor]:
+        src_ids = self.src_tokenizer.encode(
+            self.src_texts[index],
+            add_special_tokens=True,
+            max_length=self.max_src_len,
+        )
+
+        target_core = self.tgt_tokenizer.encode(
+            self.tgt_texts[index],
+            add_special_tokens=False,
+            max_length=max(1, self.max_tgt_len - 1),
+        )
+        tgt_input_ids = [self.tgt_tokenizer.bos_token_id] + target_core
+        labels = target_core + [self.tgt_tokenizer.eos_token_id]
+
+        return {
+            "src_ids": torch.tensor(src_ids, dtype=torch.long),
+            "tgt_input_ids": torch.tensor(tgt_input_ids, dtype=torch.long),
+            "labels": torch.tensor(labels, dtype=torch.long),
+            "src_len": torch.tensor(len(src_ids), dtype=torch.long),
+            "tgt_len": torch.tensor(len(labels), dtype=torch.long),
+        }
+
+
+def _translation_to_columns(dataset, src_lang: str, tgt_lang: str):
+    def convert(example):
+        translation = example["translation"]
+        return {"src": translation[src_lang], "tgt": translation[tgt_lang]}
+
+    return dataset.map(convert, remove_columns=dataset.column_names)
+
+
+def load_wmt_dataset(
+    year: str = "19",
+    language_pair: str = "zh-en",
+    src_lang: str = "en",
+    tgt_lang: str = "zh",
+    split: str | None = None,
+    cache_dir: str | None = None,
+):
+    """Load WMT zh-en and normalize rows to {'src', 'tgt'}."""
+    from datasets import load_dataset
+
+    dataset_name = f"wmt/wmt{year}"
+    try:
+        dataset = load_dataset(dataset_name, language_pair, split=split, cache_dir=cache_dir)
+    except Exception:
+        dataset = load_dataset(f"wmt{year}", language_pair, split=split, cache_dir=cache_dir)
+
+    if split is not None:
+        return _translation_to_columns(dataset, src_lang, tgt_lang)
+
+    return dataset.map(
+        lambda example: {"src": example["translation"][src_lang], "tgt": example["translation"][tgt_lang]},
+        remove_columns=next(iter(dataset.values())).column_names,
+    )
+
+
+def load_opus_dataset(
+    subset: str = "en-zh",
+    src_lang: str = "en",
+    tgt_lang: str = "zh",
+    split: str | None = None,
+    cache_dir: str | None = None,
+):
+    """Load OPUS-100 and normalize rows to {'src', 'tgt'}."""
+    from datasets import load_dataset
+
+    dataset = load_dataset("Helsinki-NLP/opus-100", subset, split=split, cache_dir=cache_dir)
+    if split is not None:
+        return _translation_to_columns(dataset, src_lang, tgt_lang)
+
+    return dataset.map(
+        lambda example: {"src": example["translation"][src_lang], "tgt": example["translation"][tgt_lang]},
+        remove_columns=next(iter(dataset.values())).column_names,
+    )
+
+
+def _read_lines(path: str | Path) -> list[str]:
+    with Path(path).open("r", encoding="utf-8") as f:
+        return [line.rstrip("\n") for line in f]
+
+
+def load_custom_dataset(
+    train_src: str | Path,
+    train_tgt: str | Path,
+    val_src: str | Path | None = None,
+    val_tgt: str | Path | None = None,
+    test_src: str | Path | None = None,
+    test_tgt: str | Path | None = None,
+    preprocessing_config: Mapping | None = None,
+) -> dict[str, dict[str, list[str]]]:
+    """Load parallel text files and return split dictionaries."""
+    preprocessing_config = dict(preprocessing_config or {})
+
+    def load_split(src_path: str | Path, tgt_path: str | Path) -> dict[str, list[str]]:
+        src_texts, tgt_texts = preprocess_pipeline(
+            _read_lines(src_path),
+            _read_lines(tgt_path),
+            **preprocessing_config,
+        )
+        return {"src": src_texts, "tgt": tgt_texts}
+
+    result = {"train": load_split(train_src, train_tgt)}
+
+    if val_src and val_tgt:
+        result["validation"] = load_split(val_src, val_tgt)
+    if test_src and test_tgt:
+        result["test"] = load_split(test_src, test_tgt)
+
+    return result

src/easytranslate/data/preprocessing.py

@@ -1,125 +1,125 @@
-"""Text cleaning and filtering utilities for translation corpora."""
-
-from __future__ import annotations
-
-import re
-import unicodedata
-from typing import Callable, Iterable, Sequence
-
-CONTROL_OR_ZERO_WIDTH_RE = re.compile(r"[\u0000-\u001f\u007f-\u009f\u200b\u200c\u200d\ufeff]")
-SPACE_RE = re.compile(r"\s+")
-PUNCT_RE = re.compile(r"[^\w\s\u4e00-\u9fff]", flags=re.UNICODE)
-
-
-def clean_text(text: str, lowercase: bool = False, remove_punctuation: bool = False) -> str:
-    """Normalize a single sentence without changing its meaning aggressively."""
-    if text is None:
-        return ""
-
-    text = unicodedata.normalize("NFKC", str(text))
-    text = CONTROL_OR_ZERO_WIDTH_RE.sub("", text)
-    text = SPACE_RE.sub(" ", text).strip()
-
-    if lowercase:
-        text = text.lower()
-    if remove_punctuation:
-        text = PUNCT_RE.sub("", text)
-        text = SPACE_RE.sub(" ", text).strip()
-
-    return text
-
-
-def _default_length(text: str) -> int:
-    """Use whitespace tokens for Latin text and character count for CJK-heavy text."""
-    cjk_chars = sum(1 for ch in text if "\u4e00" <= ch <= "\u9fff")
-    if cjk_chars >= max(1, len(text) // 3):
-        return len(text.replace(" ", ""))
-    return len(text.split())
-
-
-def filter_by_length(
-    src: str,
-    tgt: str,
-    min_src_len: int = 1,
-    min_tgt_len: int = 1,
-    max_src_len: int = 256,
-    max_tgt_len: int = 256,
-    length_ratio_threshold: float = 3.0,
-    length_fn: Callable[[str], int] | None = None,
-) -> bool:
-    """Return True when a sentence pair passes basic length and ratio checks."""
-    length_fn = length_fn or _default_length
-    src_len = length_fn(src)
-    tgt_len = length_fn(tgt)
-
-    if src_len < min_src_len or tgt_len < min_tgt_len:
-        return False
-    if src_len > max_src_len or tgt_len > max_tgt_len:
-        return False
-
-    shorter = max(1, min(src_len, tgt_len))
-    longer = max(src_len, tgt_len)
-    return longer / shorter <= length_ratio_threshold
-
-
-def deduplicate_pairs(pairs: Iterable[tuple[str, str]]) -> list[tuple[str, str]]:
-    """Deduplicate by exact cleaned source-target pair while preserving order."""
-    seen: set[tuple[str, str]] = set()
-    result: list[tuple[str, str]] = []
-
-    for src, tgt in pairs:
-        key = (src, tgt)
-        if key in seen:
-            continue
-        seen.add(key)
-        result.append(key)
-
-    return result
-
-
-def preprocess_pipeline(
-    src_texts: Sequence[str],
-    tgt_texts: Sequence[str],
-    lowercase_src: bool = False,
-    lowercase_tgt: bool = False,
-    remove_punctuation: bool = False,
-    max_src_len: int = 256,
-    max_tgt_len: int = 256,
-    min_src_len: int = 1,
-    min_tgt_len: int = 1,
-    filter_by_length_enabled: bool = True,
-    length_ratio_threshold: float = 3.0,
-    deduplicate: bool = True,
-) -> tuple[list[str], list[str]]:
-    """Clean, filter, and optionally deduplicate parallel source-target texts."""
-    if len(src_texts) != len(tgt_texts):
-        raise ValueError("src_texts and tgt_texts must have the same length")
-
-    pairs: list[tuple[str, str]] = []
-    for raw_src, raw_tgt in zip(src_texts, tgt_texts):
-        src = clean_text(raw_src, lowercase=lowercase_src, remove_punctuation=remove_punctuation)
-        tgt = clean_text(raw_tgt, lowercase=lowercase_tgt, remove_punctuation=remove_punctuation)
-
-        if not src or not tgt:
-            continue
-        if filter_by_length_enabled and not filter_by_length(
-            src,
-            tgt,
-            min_src_len=min_src_len,
-            min_tgt_len=min_tgt_len,
-            max_src_len=max_src_len,
-            max_tgt_len=max_tgt_len,
-            length_ratio_threshold=length_ratio_threshold,
-        ):
-            continue
-
-        pairs.append((src, tgt))
-
-    if deduplicate:
-        pairs = deduplicate_pairs(pairs)
-
-    if not pairs:
-        return [], []
-
-    src_clean, tgt_clean = zip(*pairs)
-    return list(src_clean), list(tgt_clean)
+"""Text cleaning and filtering utilities for translation corpora."""
+
+from __future__ import annotations
+
+import re
+import unicodedata
+from typing import Callable, Iterable, Sequence
+
+CONTROL_OR_ZERO_WIDTH_RE = re.compile(r"[\u0000-\u001f\u007f-\u009f\u200b\u200c\u200d\ufeff]")
+SPACE_RE = re.compile(r"\s+")
+PUNCT_RE = re.compile(r"[^\w\s\u4e00-\u9fff]", flags=re.UNICODE)
+
+
+def clean_text(text: str, lowercase: bool = False, remove_punctuation: bool = False) -> str:
+    """Normalize a single sentence without changing its meaning aggressively."""
+    if text is None:
+        return ""
+
+    text = unicodedata.normalize("NFKC", str(text))
+    text = CONTROL_OR_ZERO_WIDTH_RE.sub("", text)
+    text = SPACE_RE.sub(" ", text).strip()
+
+    if lowercase:
+        text = text.lower()
+    if remove_punctuation:
+        text = PUNCT_RE.sub("", text)
+        text = SPACE_RE.sub(" ", text).strip()
+
+    return text
+
+
+def _default_length(text: str) -> int:
+    """Use whitespace tokens for Latin text and character count for CJK-heavy text."""
+    cjk_chars = sum(1 for ch in text if "\u4e00" <= ch <= "\u9fff")
+    if cjk_chars >= max(1, len(text) // 3):
+        return len(text.replace(" ", ""))
+    return len(text.split())
+
+
+def filter_by_length(
+    src: str,
+    tgt: str,
+    min_src_len: int = 1,
+    min_tgt_len: int = 1,
+    max_src_len: int = 256,
+    max_tgt_len: int = 256,
+    length_ratio_threshold: float = 3.0,
+    length_fn: Callable[[str], int] | None = None,
+) -> bool:
+    """Return True when a sentence pair passes basic length and ratio checks."""
+    length_fn = length_fn or _default_length
+    src_len = length_fn(src)
+    tgt_len = length_fn(tgt)
+
+    if src_len < min_src_len or tgt_len < min_tgt_len:
+        return False
+    if src_len > max_src_len or tgt_len > max_tgt_len:
+        return False
+
+    shorter = max(1, min(src_len, tgt_len))
+    longer = max(src_len, tgt_len)
+    return longer / shorter <= length_ratio_threshold
+
+
+def deduplicate_pairs(pairs: Iterable[tuple[str, str]]) -> list[tuple[str, str]]:
+    """Deduplicate by exact cleaned source-target pair while preserving order."""
+    seen: set[tuple[str, str]] = set()
+    result: list[tuple[str, str]] = []
+
+    for src, tgt in pairs:
+        key = (src, tgt)
+        if key in seen:
+            continue
+        seen.add(key)
+        result.append(key)
+
+    return result
+
+
+def preprocess_pipeline(
+    src_texts: Sequence[str],
+    tgt_texts: Sequence[str],
+    lowercase_src: bool = False,
+    lowercase_tgt: bool = False,
+    remove_punctuation: bool = False,
+    max_src_len: int = 256,
+    max_tgt_len: int = 256,
+    min_src_len: int = 1,
+    min_tgt_len: int = 1,
+    filter_by_length_enabled: bool = True,
+    length_ratio_threshold: float = 3.0,
+    deduplicate: bool = True,
+) -> tuple[list[str], list[str]]:
+    """Clean, filter, and optionally deduplicate parallel source-target texts."""
+    if len(src_texts) != len(tgt_texts):
+        raise ValueError("src_texts and tgt_texts must have the same length")
+
+    pairs: list[tuple[str, str]] = []
+    for raw_src, raw_tgt in zip(src_texts, tgt_texts):
+        src = clean_text(raw_src, lowercase=lowercase_src, remove_punctuation=remove_punctuation)
+        tgt = clean_text(raw_tgt, lowercase=lowercase_tgt, remove_punctuation=remove_punctuation)
+
+        if not src or not tgt:
+            continue
+        if filter_by_length_enabled and not filter_by_length(
+            src,
+            tgt,
+            min_src_len=min_src_len,
+            min_tgt_len=min_tgt_len,
+            max_src_len=max_src_len,
+            max_tgt_len=max_tgt_len,
+            length_ratio_threshold=length_ratio_threshold,
+        ):
+            continue
+
+        pairs.append((src, tgt))
+
+    if deduplicate:
+        pairs = deduplicate_pairs(pairs)
+
+    if not pairs:
+        return [], []
+
+    src_clean, tgt_clean = zip(*pairs)
+    return list(src_clean), list(tgt_clean)

src/easytranslate/data/tokenizer.py

@@ -1,197 +1,197 @@
-"""Tokenizer wrappers and BPE training helpers."""
-
-from __future__ import annotations
-
-from pathlib import Path
-from typing import Iterable, Mapping, Sequence
-
-
-class TokenizerWrapper:
-    """Small adapter that gives HF tokenizers and tokenizers.Tokenizer one API."""
-
-    def __init__(
-        self,
-        tokenizer,
-        pad_token: str = "<pad>",
-        unk_token: str = "<unk>",
-        bos_token: str = "<s>",
-        eos_token: str = "</s>",
-    ):
-        self.tokenizer = tokenizer
-        self.pad_token = pad_token
-        self.unk_token = unk_token
-        self.bos_token = bos_token
-        self.eos_token = eos_token
-
-    @property
-    def pad_token_id(self) -> int:
-        return self.token_to_id(self.pad_token)
-
-    @property
-    def unk_token_id(self) -> int:
-        return self.token_to_id(self.unk_token)
-
-    @property
-    def bos_token_id(self) -> int:
-        return self.token_to_id(self.bos_token)
-
-    @property
-    def eos_token_id(self) -> int:
-        return self.token_to_id(self.eos_token)
-
-    @property
-    def vocab_size(self) -> int:
-        if hasattr(self.tokenizer, "get_vocab_size"):
-            return int(self.tokenizer.get_vocab_size())
-        return int(len(self.tokenizer))
-
-    def token_to_id(self, token: str) -> int:
-        if hasattr(self.tokenizer, "token_to_id"):
-            idx = self.tokenizer.token_to_id(token)
-        elif hasattr(self.tokenizer, "convert_tokens_to_ids"):
-            idx = self.tokenizer.convert_tokens_to_ids(token)
-        else:
-            raise TypeError("Unsupported tokenizer type")
-
-        if idx is None:
-            raise ValueError(f"Token {token!r} is not in the tokenizer vocabulary")
-        return int(idx)
-
-    def encode(self, text: str, add_special_tokens: bool = False, max_length: int | None = None) -> list[int]:
-        if hasattr(self.tokenizer, "encode") and self.tokenizer.__class__.__module__.startswith("tokenizers"):
-            ids = self.tokenizer.encode(text).ids
-        else:
-            ids = self.tokenizer.encode(text, add_special_tokens=add_special_tokens)
-            add_special_tokens = False
-
-        if add_special_tokens:
-            ids = [self.bos_token_id] + list(ids) + [self.eos_token_id]
-
-        if max_length is not None:
-            ids = list(ids)[:max_length]
-
-        return list(ids)
-
-    def decode(self, ids: Sequence[int], skip_special_tokens: bool = True) -> str:
-        if hasattr(self.tokenizer, "decode"):
-            try:
-                return self.tokenizer.decode(list(ids), skip_special_tokens=skip_special_tokens)
-            except TypeError:
-                return self.tokenizer.decode(list(ids))
-        raise TypeError("Unsupported tokenizer type")
-
-    def save(self, path: str | Path) -> None:
-        path = Path(path)
-        path.parent.mkdir(parents=True, exist_ok=True)
-
-        if hasattr(self.tokenizer, "save"):
-            self.tokenizer.save(str(path))
-            return
-        if hasattr(self.tokenizer, "save_pretrained"):
-            self.tokenizer.save_pretrained(str(path))
-            return
-        raise TypeError("Unsupported tokenizer type")
-
-
-def _special_tokens(config: Mapping | None = None) -> dict[str, str]:
-    tokens = {
-        "pad": "<pad>",
-        "unk": "<unk>",
-        "bos": "<s>",
-        "eos": "</s>",
-    }
-    if config:
-        tokens.update(dict(config))
-    return tokens
-
-
-def train_bpe_tokenizer(
-    texts: Iterable[str],
-    vocab_size: int = 32000,
-    min_frequency: int = 2,
-    special_tokens: Mapping[str, str] | None = None,
-    save_path: str | Path | None = None,
-) -> TokenizerWrapper:
-    """Train a byte-level BPE tokenizer on source and target training text."""
-    from tokenizers import Tokenizer
-    from tokenizers.decoders import ByteLevel as ByteLevelDecoder
-    from tokenizers.models import BPE
-    from tokenizers.normalizers import NFKC, Sequence as NormalizerSequence
-    from tokenizers.pre_tokenizers import ByteLevel
-    from tokenizers.trainers import BpeTrainer
-
-    tokens = _special_tokens(special_tokens)
-    ordered_specials = [tokens["pad"], tokens["unk"], tokens["bos"], tokens["eos"]]
-
-    tokenizer = Tokenizer(BPE(unk_token=tokens["unk"]))
-    tokenizer.normalizer = NormalizerSequence([NFKC()])
-    tokenizer.pre_tokenizer = ByteLevel(add_prefix_space=False)
-    tokenizer.decoder = ByteLevelDecoder()
-
-    trainer = BpeTrainer(
-        vocab_size=vocab_size,
-        min_frequency=min_frequency,
-        special_tokens=ordered_specials,
-        show_progress=True,
-    )
-    tokenizer.train_from_iterator((text for text in texts if text), trainer=trainer)
-
-    wrapper = TokenizerWrapper(
-        tokenizer,
-        pad_token=tokens["pad"],
-        unk_token=tokens["unk"],
-        bos_token=tokens["bos"],
-        eos_token=tokens["eos"],
-    )
-
-    if save_path is not None:
-        wrapper.save(save_path)
-
-    return wrapper
-
-
-def build_tokenizer(config: Mapping, train_texts: Iterable[str] | None = None) -> TokenizerWrapper:
-    """Build a tokenizer from project config."""
-    tokenizer_type = config.get("type", "bpe")
-    tokens = _special_tokens(config.get("special_tokens"))
-
-    if tokenizer_type == "pretrained":
-        from transformers import AutoTokenizer
-
-        model_name = config.get("model_name") or config.get("pretrained_model_name")
-        if not model_name:
-            raise ValueError("pretrained tokenizer requires config['model_name']")
-        tokenizer = AutoTokenizer.from_pretrained(model_name)
-        return TokenizerWrapper(
-            tokenizer,
-            pad_token=tokenizer.pad_token or tokens["pad"],
-            unk_token=tokenizer.unk_token or tokens["unk"],
-            bos_token=tokenizer.bos_token or tokens["bos"],
-            eos_token=tokenizer.eos_token or tokens["eos"],
-        )
-
-    if tokenizer_type in {"bpe", "sentencepiece"}:
-        tokenizer_path = config.get("path") or config.get("tokenizer_path")
-        if tokenizer_path and Path(tokenizer_path).exists():
-            from tokenizers import Tokenizer
-
-            return TokenizerWrapper(
-                Tokenizer.from_file(str(tokenizer_path)),
-                pad_token=tokens["pad"],
-                unk_token=tokens["unk"],
-                bos_token=tokens["bos"],
-                eos_token=tokens["eos"],
-            )
-
-        if train_texts is None:
-            raise ValueError("BPE tokenizer requires train_texts when no tokenizer path is provided")
-
-        return train_bpe_tokenizer(
-            train_texts,
-            vocab_size=int(config.get("vocab_size", 32000)),
-            min_frequency=int(config.get("min_frequency", 2)),
-            special_tokens=tokens,
-            save_path=tokenizer_path,
-        )
-
-    raise ValueError(f"Unsupported tokenizer type: {tokenizer_type}")
+"""Tokenizer wrappers and BPE training helpers."""
+
+from __future__ import annotations
+
+from pathlib import Path
+from typing import Iterable, Mapping, Sequence
+
+
+class TokenizerWrapper:
+    """Small adapter that gives HF tokenizers and tokenizers.Tokenizer one API."""
+
+    def __init__(
+        self,
+        tokenizer,
+        pad_token: str = "<pad>",
+        unk_token: str = "<unk>",
+        bos_token: str = "<s>",
+        eos_token: str = "</s>",
+    ):
+        self.tokenizer = tokenizer
+        self.pad_token = pad_token
+        self.unk_token = unk_token
+        self.bos_token = bos_token
+        self.eos_token = eos_token
+
+    @property
+    def pad_token_id(self) -> int:
+        return self.token_to_id(self.pad_token)
+
+    @property
+    def unk_token_id(self) -> int:
+        return self.token_to_id(self.unk_token)
+
+    @property
+    def bos_token_id(self) -> int:
+        return self.token_to_id(self.bos_token)
+
+    @property
+    def eos_token_id(self) -> int:
+        return self.token_to_id(self.eos_token)
+
+    @property
+    def vocab_size(self) -> int:
+        if hasattr(self.tokenizer, "get_vocab_size"):
+            return int(self.tokenizer.get_vocab_size())
+        return int(len(self.tokenizer))
+
+    def token_to_id(self, token: str) -> int:
+        if hasattr(self.tokenizer, "token_to_id"):
+            idx = self.tokenizer.token_to_id(token)
+        elif hasattr(self.tokenizer, "convert_tokens_to_ids"):
+            idx = self.tokenizer.convert_tokens_to_ids(token)
+        else:
+            raise TypeError("Unsupported tokenizer type")
+
+        if idx is None:
+            raise ValueError(f"Token {token!r} is not in the tokenizer vocabulary")
+        return int(idx)
+
+    def encode(self, text: str, add_special_tokens: bool = False, max_length: int | None = None) -> list[int]:
+        if hasattr(self.tokenizer, "encode") and self.tokenizer.__class__.__module__.startswith("tokenizers"):
+            ids = self.tokenizer.encode(text).ids
+        else:
+            ids = self.tokenizer.encode(text, add_special_tokens=add_special_tokens)
+            add_special_tokens = False
+
+        if add_special_tokens:
+            ids = [self.bos_token_id] + list(ids) + [self.eos_token_id]
+
+        if max_length is not None:
+            ids = list(ids)[:max_length]
+
+        return list(ids)
+
+    def decode(self, ids: Sequence[int], skip_special_tokens: bool = True) -> str:
+        if hasattr(self.tokenizer, "decode"):
+            try:
+                return self.tokenizer.decode(list(ids), skip_special_tokens=skip_special_tokens)
+            except TypeError:
+                return self.tokenizer.decode(list(ids))
+        raise TypeError("Unsupported tokenizer type")
+
+    def save(self, path: str | Path) -> None:
+        path = Path(path)
+        path.parent.mkdir(parents=True, exist_ok=True)
+
+        if hasattr(self.tokenizer, "save"):
+            self.tokenizer.save(str(path))
+            return
+        if hasattr(self.tokenizer, "save_pretrained"):
+            self.tokenizer.save_pretrained(str(path))
+            return
+        raise TypeError("Unsupported tokenizer type")
+
+
+def _special_tokens(config: Mapping | None = None) -> dict[str, str]:
+    tokens = {
+        "pad": "<pad>",
+        "unk": "<unk>",
+        "bos": "<s>",
+        "eos": "</s>",
+    }
+    if config:
+        tokens.update(dict(config))
+    return tokens
+
+
+def train_bpe_tokenizer(
+    texts: Iterable[str],
+    vocab_size: int = 32000,
+    min_frequency: int = 2,
+    special_tokens: Mapping[str, str] | None = None,
+    save_path: str | Path | None = None,
+) -> TokenizerWrapper:
+    """Train a byte-level BPE tokenizer on source and target training text."""
+    from tokenizers import Tokenizer
+    from tokenizers.decoders import ByteLevel as ByteLevelDecoder
+    from tokenizers.models import BPE
+    from tokenizers.normalizers import NFKC, Sequence as NormalizerSequence
+    from tokenizers.pre_tokenizers import ByteLevel
+    from tokenizers.trainers import BpeTrainer
+
+    tokens = _special_tokens(special_tokens)
+    ordered_specials = [tokens["pad"], tokens["unk"], tokens["bos"], tokens["eos"]]
+
+    tokenizer = Tokenizer(BPE(unk_token=tokens["unk"]))
+    tokenizer.normalizer = NormalizerSequence([NFKC()])
+    tokenizer.pre_tokenizer = ByteLevel(add_prefix_space=False)
+    tokenizer.decoder = ByteLevelDecoder()
+
+    trainer = BpeTrainer(
+        vocab_size=vocab_size,
+        min_frequency=min_frequency,
+        special_tokens=ordered_specials,
+        show_progress=True,
+    )
+    tokenizer.train_from_iterator((text for text in texts if text), trainer=trainer)
+
+    wrapper = TokenizerWrapper(
+        tokenizer,
+        pad_token=tokens["pad"],
+        unk_token=tokens["unk"],
+        bos_token=tokens["bos"],
+        eos_token=tokens["eos"],
+    )
+
+    if save_path is not None:
+        wrapper.save(save_path)
+
+    return wrapper
+
+
+def build_tokenizer(config: Mapping, train_texts: Iterable[str] | None = None) -> TokenizerWrapper:
+    """Build a tokenizer from project config."""
+    tokenizer_type = config.get("type", "bpe")
+    tokens = _special_tokens(config.get("special_tokens"))
+
+    if tokenizer_type == "pretrained":
+        from transformers import AutoTokenizer
+
+        model_name = config.get("model_name") or config.get("pretrained_model_name")
+        if not model_name:
+            raise ValueError("pretrained tokenizer requires config['model_name']")
+        tokenizer = AutoTokenizer.from_pretrained(model_name)
+        return TokenizerWrapper(
+            tokenizer,
+            pad_token=tokenizer.pad_token or tokens["pad"],
+            unk_token=tokenizer.unk_token or tokens["unk"],
+            bos_token=tokenizer.bos_token or tokens["bos"],
+            eos_token=tokenizer.eos_token or tokens["eos"],
+        )
+
+    if tokenizer_type in {"bpe", "sentencepiece"}:
+        tokenizer_path = config.get("path") or config.get("tokenizer_path")
+        if tokenizer_path and Path(tokenizer_path).exists():
+            from tokenizers import Tokenizer
+
+            return TokenizerWrapper(
+                Tokenizer.from_file(str(tokenizer_path)),
+                pad_token=tokens["pad"],
+                unk_token=tokens["unk"],
+                bos_token=tokens["bos"],
+                eos_token=tokens["eos"],
+            )
+
+        if train_texts is None:
+            raise ValueError("BPE tokenizer requires train_texts when no tokenizer path is provided")
+
+        return train_bpe_tokenizer(
+            train_texts,
+            vocab_size=int(config.get("vocab_size", 32000)),
+            min_frequency=int(config.get("min_frequency", 2)),
+            special_tokens=tokens,
+            save_path=tokenizer_path,
+        )
+
+    raise ValueError(f"Unsupported tokenizer type: {tokenizer_type}")

src/easytranslate/evaluation/__init__.py

@@ -1,16 +1,16 @@
-"""评估与推理模块 (Person D 负责)"""
-
-from easytranslate.evaluation.metrics import compute_bleu, compute_comet, compute_chrf, compute_all_metrics
-from easytranslate.evaluation.decoding import greedy_decode, beam_search_decode, sample_decode
-from easytranslate.evaluation.evaluator import Evaluator
-
-__all__ = [
-    "compute_bleu",
-    "compute_comet",
-    "compute_chrf",
-    "compute_all_metrics",
-    "greedy_decode",
-    "beam_search_decode",
-    "sample_decode",
-    "Evaluator",
-]
+"""评估与推理模块 (Person D 负责)"""
+
+from easytranslate.evaluation.metrics import compute_bleu, compute_comet, compute_chrf, compute_all_metrics
+from easytranslate.evaluation.decoding import greedy_decode, beam_search_decode, sample_decode
+from easytranslate.evaluation.evaluator import Evaluator
+
+__all__ = [
+    "compute_bleu",
+    "compute_comet",
+    "compute_chrf",
+    "compute_all_metrics",
+    "greedy_decode",
+    "beam_search_decode",
+    "sample_decode",
+    "Evaluator",
+]

src/easytranslate/evaluation/decoding.py

@@ -1,129 +1,219 @@
-"""
-解码策略模块 — Person D 负责实现
-
-功能要求:
-1. greedy_decode: 贪心解码
-2. beam_search_decode: 束搜索解码
-3. sample_decode: 采样解码 (temperature, top-k, top-p)
-
-技术要点:
-- Beam Search 是翻译任务最常用的解码策略
-- 需要高效处理批量解码
-- 支持长度惩罚 (length penalty) 和重复惩罚 (no_repeat_ngram)
-- 对于预训练模型，可以直接使用 model.generate()
-"""
-
-from __future__ import annotations
-
-import logging
-from typing import Optional
-
-import torch
-import torch.nn as nn
-import torch.nn.functional as F
-
-logger = logging.getLogger(__name__)
-
-
-@torch.no_grad()
-def greedy_decode(
-    model: nn.Module,
-    src_ids: torch.Tensor,          # [B, S]
-    src_padding_mask: torch.BoolTensor,  # [B, S]
-    bos_id: int,
-    eos_id: int,
-    max_len: int = 256,
-) -> torch.Tensor:
-    """
-    贪心解码。
-
-    TODO [Person D]: 实现以下逻辑:
-    1. encoder_output = model.encode(src_ids, src_padding_mask)
-    2. 初始化 decoder input: [B, 1] 全为 bos_id
-    3. for step in range(max_len):
-       a. logits = model.decode_step(decoder_input, encoder_output, src_padding_mask)
-       b. next_token = logits.argmax(dim=-1)
-       c. decoder_input = concat(decoder_input, next_token)
-       d. 如果所有序列都生成了 eos_id，则提前终止
-    4. 返回生成的 token ids [B, T]
-    """
-    raise NotImplementedError("TODO: Person D 实现 greedy_decode")
-
-
-@torch.no_grad()
-def beam_search_decode(
-    model: nn.Module,
-    src_ids: torch.Tensor,          # [B, S]
-    src_padding_mask: torch.BoolTensor,
-    bos_id: int,
-    eos_id: int,
-    beam_size: int = 5,
-    max_len: int = 256,
-    length_penalty: float = 1.0,
-    no_repeat_ngram_size: int = 0,
-) -> torch.Tensor:
-    """
-    束搜索解码。
-
-    TODO [Person D]: 实现以下逻辑:
-    1. encoder_output = model.encode(src_ids, src_padding_mask)
-    2. 将 encoder_output 扩展为 beam_size 份: [B*beam, S, D]
-    3. 初始化 beam:
-       - beam_scores: [B, beam_size] 初始为 0
-       - beam_tokens: [B, beam_size, 1] 初始为 bos_id
-    4. for step in range(max_len):
-       a. 对每个 beam 计算 logits
-       b. log_probs = log_softmax(logits)
-       c. (可选) 应用 no_repeat_ngram 约束
-       d. scores = beam_scores + log_probs
-       e. 选择 top-k candidates (k = beam_size)
-       f. 更新 beam_tokens 和 beam_scores
-       g. 将已完成的 beam 移到 finished pool
-    5. 对 finished beams 应用 length_penalty:
-       score = score / (length ^ length_penalty)
-    6. 选择得分最高的序列
-    7. 返回最佳翻译 [B, T]
-
-    这是翻译任务最关键的解码算法，请仔细实现。
-    """
-    raise NotImplementedError("TODO: Person D 实现 beam_search_decode")
-
-
-@torch.no_grad()
-def sample_decode(
-    model: nn.Module,
-    src_ids: torch.Tensor,
-    src_padding_mask: torch.BoolTensor,
-    bos_id: int,
-    eos_id: int,
-    max_len: int = 256,
-    temperature: float = 1.0,
-    top_k: int = 0,
-    top_p: float = 1.0,
-) -> torch.Tensor:
-    """
-    采样解码 (支持 temperature, top-k, top-p/nucleus sampling)。
-
-    TODO [Person D]: 实现以下逻辑:
-    1. 与贪心解码类似，但每步采样而非取 argmax
-    2. 应用 temperature: logits = logits / temperature
-    3. 应用 top-k: 只保留概率最高的 k 个 token
-    4. 应用 top-p (nucleus): 只保留累积概率达到 p 的 token
-    5. 从过滤后的分布中采样: torch.multinomial
-    """
-    raise NotImplementedError("TODO: Person D 实现 sample_decode")
-
-
-def _apply_no_repeat_ngram(
-    logits: torch.Tensor,
-    generated_tokens: torch.Tensor,
-    ngram_size: int,
-) -> torch.Tensor:
-    """
-    防止生成重复的 n-gram。
-
-    TODO [Person D]:
-    1. 从 generated_tokens 中提取所有已出现的 (ngram_size-1)-gram
-    2. 对于每个可能导致重复 ngram 的 next token，将其 logits 设为 -inf
-    """
-    raise NotImplementedError("TODO: Person D 实现 _apply_no_repeat_ngram")
+"""
+解码策略模块 — Person D 负责实现
+
+功能要求:
+1. greedy_decode: 贪心解码
+2. beam_search_decode: 束搜索解码
+3. sample_decode: 采样解码 (temperature, top-k, top-p)
+
+技术要点:
+- Beam Search 是翻译任务最常用的解码策略
+- 需要高效处理批量解码
+- 支持长度惩罚 (length penalty) 和重复惩罚 (no_repeat_ngram)
+- 对于预训练模型，可以直接使用 model.generate()
+"""
+
+from __future__ import annotations
+
+import logging
+from typing import Optional
+
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+
+logger = logging.getLogger(__name__)
+
+
+def _apply_no_repeat_ngram(
+    logits: torch.Tensor,
+    generated_tokens: torch.Tensor,
+    ngram_size: int,
+) -> torch.Tensor:
+    """
+    防止生成重复的 n-gram。
+    """
+    if ngram_size <= 1:
+        return logits
+
+    batch_size, vocab_size = logits.size()
+    for batch_idx in range(batch_size):
+        tokens = generated_tokens[batch_idx].tolist()
+        if len(tokens) < ngram_size - 1:
+            continue
+
+        banned_tokens: set[int] = set()
+        ngram_map: dict[tuple[int, ...], set[int]] = {}
+        for i in range(len(tokens) - ngram_size + 1):
+            prefix = tuple(tokens[i : i + ngram_size - 1])
+            next_token = tokens[i + ngram_size - 1]
+            ngram_map.setdefault(prefix, set()).add(next_token)
+
+        prefix = tuple(tokens[-(ngram_size - 1) :])
+        if prefix in ngram_map:
+            banned_tokens = ngram_map[prefix]
+            logits[batch_idx, list(banned_tokens)] = float("-inf")
+
+    return logits
+
+
+@torch.no_grad()
+def greedy_decode(
+    model: nn.Module,
+    src_ids: torch.Tensor,
+    src_padding_mask: torch.BoolTensor,
+    bos_id: int,
+    eos_id: int,
+    max_len: int = 256,
+) -> torch.Tensor:
+    """
+    贪心解码。
+    """
+    encoder_output = model.encode(src_ids, src_padding_mask)
+    batch_size = src_ids.size(0)
+    device = src_ids.device
+    generated = torch.full((batch_size, 1), bos_id, dtype=torch.long, device=device)
+    finished = torch.zeros(batch_size, dtype=torch.bool, device=device)
+
+    for _ in range(max_len):
+        logits = model.decode_step(generated, encoder_output, src_padding_mask)
+        next_token = logits.argmax(dim=-1, keepdim=True)
+        generated = torch.cat([generated, next_token], dim=1)
+        finished = finished | next_token.squeeze(-1).eq(eos_id)
+        if finished.all():
+            break
+
+    return generated
+
+
+@torch.no_grad()
+def beam_search_decode(
+    model: nn.Module,
+    src_ids: torch.Tensor,
+    src_padding_mask: torch.BoolTensor,
+    bos_id: int,
+    eos_id: int,
+    beam_size: int = 5,
+    max_len: int = 256,
+    length_penalty: float = 1.0,
+    no_repeat_ngram_size: int = 0,
+) -> torch.Tensor:
+    """
+    束搜索解码。
+    """
+    batch_size, seq_len = src_ids.size()
+    device = src_ids.device
+
+    encoder_output = model.encode(src_ids, src_padding_mask)
+    encoder_output = encoder_output.unsqueeze(1).expand(batch_size, beam_size, -1, -1)
+    encoder_output = encoder_output.reshape(batch_size * beam_size, seq_len, -1)
+    src_padding_mask = src_padding_mask.unsqueeze(1).expand(batch_size, beam_size, seq_len)
+    src_padding_mask = src_padding_mask.reshape(batch_size * beam_size, seq_len)
+
+    beam_scores = torch.full((batch_size, beam_size), float("-inf"), device=device)
+    beam_scores[:, 0] = 0.0
+    generated = torch.full((batch_size, beam_size, 1), bos_id, dtype=torch.long, device=device)
+    finished = torch.zeros((batch_size, beam_size), dtype=torch.bool, device=device)
+
+    for _ in range(max_len):
+        flat_generated = generated.view(batch_size * beam_size, -1)
+        logits = model.decode_step(flat_generated, encoder_output, src_padding_mask)
+        log_probs = F.log_softmax(logits, dim=-1)
+
+        if no_repeat_ngram_size > 0:
+            log_probs = _apply_no_repeat_ngram(log_probs, flat_generated, no_repeat_ngram_size)
+
+        finished_flat = finished.view(batch_size * beam_size)
+        if finished_flat.any():
+            log_probs[finished_flat] = float("-inf")
+            log_probs[finished_flat, eos_id] = 0.0
+
+        vocab_size = log_probs.size(-1)
+        scores = beam_scores.unsqueeze(-1) + log_probs.view(batch_size, beam_size, vocab_size)
+        scores_flat = scores.view(batch_size, -1)
+        topk_scores, topk_indices = scores_flat.topk(beam_size, dim=-1)
+
+        beam_indices = topk_indices // vocab_size
+        token_indices = topk_indices % vocab_size
+
+        next_generated = []
+        next_finished = []
+        for batch_idx in range(batch_size):
+            selected_beams = beam_indices[batch_idx]
+            selected_tokens = token_indices[batch_idx]
+            next_seq = generated[batch_idx, selected_beams]
+            next_seq = torch.cat([next_seq, selected_tokens.unsqueeze(-1)], dim=-1)
+            next_generated.append(next_seq)
+            next_finished.append(
+                finished[batch_idx, selected_beams]
+                | selected_tokens.eq(eos_id)
+            )
+
+        generated = torch.stack(next_generated, dim=0)
+        finished = torch.stack(next_finished, dim=0)
+        beam_scores = topk_scores
+
+        if finished.all():
+            break
+
+    length = generated.size(1)
+    penalty = float(length) ** float(length_penalty)
+    final_scores = beam_scores / penalty
+    best_indices = final_scores.argmax(dim=-1)
+    output = generated[torch.arange(batch_size, device=device), best_indices]
+    return output
+
+
+@torch.no_grad()
+def sample_decode(
+    model: nn.Module,
+    src_ids: torch.Tensor,
+    src_padding_mask: torch.BoolTensor,
+    bos_id: int,
+    eos_id: int,
+    max_len: int = 256,
+    temperature: float = 1.0,
+    top_k: int = 0,
+    top_p: float = 1.0,
+) -> torch.Tensor:
+    """
+    采样解码 (支持 temperature, top-k, top-p/nucleus sampling)。
+    """
+    encoder_output = model.encode(src_ids, src_padding_mask)
+    batch_size = src_ids.size(0)
+    device = src_ids.device
+    generated = torch.full((batch_size, 1), bos_id, dtype=torch.long, device=device)
+    finished = torch.zeros(batch_size, dtype=torch.bool, device=device)
+
+    for _ in range(max_len):
+        logits = model.decode_step(generated, encoder_output, src_padding_mask)
+        logits = logits / max(temperature, 1e-8)
+
+        if top_k > 0:
+            top_k = min(top_k, logits.size(-1))
+            values, indices = torch.topk(logits, top_k, dim=-1)
+            mask = torch.full_like(logits, float("-inf"))
+            logits = mask.scatter(-1, indices, values)
+
+        if 0.0 < top_p < 1.0:
+            sorted_logits, sorted_indices = torch.sort(logits, descending=True, dim=-1)
+            probs = F.softmax(sorted_logits, dim=-1)
+            cumulative_probs = torch.cumsum(probs, dim=-1)
+            cutoff = cumulative_probs > top_p
+            cutoff[:, 1:] = cutoff[:, :-1].clone()
+            cutoff[:, 0] = False
+            sorted_logits[cutoff] = float("-inf")
+            logits = torch.zeros_like(logits).scatter(-1, sorted_indices, sorted_logits)
+
+        probs = F.softmax(logits, dim=-1)
+        next_token = torch.multinomial(probs, num_samples=1)
+        next_token = next_token.clamp(min=0)
+
+        next_token = torch.where(finished.unsqueeze(-1), torch.full_like(next_token, eos_id), next_token)
+        generated = torch.cat([generated, next_token], dim=1)
+        finished = finished | next_token.squeeze(-1).eq(eos_id)
+
+        if finished.all():
+            break
+
+    return generated

src/easytranslate/evaluation/evaluator.py

@@ -1,76 +1,208 @@
-"""
-评估器模块 — Person D 负责实现
-
-功能要求:
-将解码和评估指标整合为统一的评估接口。
-
-使用方法:
-    evaluator = Evaluator(model, tokenizer, config)
-    results = evaluator.evaluate(test_loader)
-"""
-
-from __future__ import annotations
-
-import logging
-from typing import Optional
-
-import torch
-import torch.nn as nn
-from torch.utils.data import DataLoader
-from tqdm import tqdm
-
-from easytranslate.evaluation.metrics import compute_all_metrics
-from easytranslate.evaluation.decoding import greedy_decode, beam_search_decode, sample_decode
-
-logger = logging.getLogger(__name__)
-
-
-class Evaluator:
-    """
-    翻译模型评估器。
-
-    TODO [Person D]: 实现以下方法。
-    """
-
-    def __init__(self, model: nn.Module, tokenizer, config: dict):
-        """
-        TODO [Person D]:
-        1. 保存 model, tokenizer, config
-        2. 从 config 读取解码策略和评估指标配置
-        3. 根据策略选择解码函数
-        """
-        raise NotImplementedError("TODO: Person D 实现 Evaluator.__init__")
-
-    def evaluate(
-        self,
-        dataloader: DataLoader,
-        src_texts: Optional[list[str]] = None,
-        ref_texts: Optional[list[str]] = None,
-    ) -> dict:
-        """
-        在给定数据上进行评估。
-
-        TODO [Person D]: 实现以下逻辑:
-        1. model.eval()
-        2. 遍历 dataloader，使用选定的解码策略生成翻译
-        3. 将生成的 token ids 解码为文本
-        4. 调用 compute_all_metrics 计算指标
-        5. 返回评估结果 dict
-        """
-        raise NotImplementedError("TODO: Person D 实现 evaluate")
-
-    def translate(self, texts: list[str]) -> list[str]:
-        """
-        翻译一批文本。
-
-        TODO [Person D]:
-        1. tokenize 输入文本
-        2. 调用解码函数生成翻译
-        3. 解码为文本
-        4. 返回翻译结果列表
-        """
-        raise NotImplementedError("TODO: Person D 实现 translate")
-
-    def translate_single(self, text: str) -> str:
-        """翻译单条文本。"""
-        return self.translate([text])[0]
+"""
+评估器模块 — Person D 负责实现
+
+功能要求:
+将解码和评估指标整合为统一的评估接口。
+
+使用方法:
+    evaluator = Evaluator(model, tokenizer, config)
+    results = evaluator.evaluate(test_loader)
+"""
+
+from __future__ import annotations
+
+import logging
+from typing import Optional
+
+import torch
+import torch.nn as nn
+from torch.nn.utils.rnn import pad_sequence
+from torch.utils.data import DataLoader
+from tqdm import tqdm
+
+from easytranslate.evaluation.metrics import compute_all_metrics
+from easytranslate.evaluation.decoding import greedy_decode, beam_search_decode, sample_decode
+
+logger = logging.getLogger(__name__)
+
+
+def _get_config_value(config, key_path, default=None):
+    if config is None:
+        return default
+    if isinstance(config, dict):
+        value = config
+        for key in key_path:
+            value = value.get(key, default)
+            if value is default:
+                break
+        return value
+    value = config
+    for key in key_path:
+        value = getattr(value, key, default)
+        if value is default:
+            break
+    return value
+
+
+class Evaluator:
+    """
+    翻译模型评估器。
+    """
+
+    def __init__(self, model: nn.Module, tokenizer, config: dict):
+        self.model = model
+        self.tokenizer = tokenizer
+        self.config = config or {}
+
+        self.evaluation_config = _get_config_value(self.config, ["evaluation"], {})
+        self.decoding_config = _get_config_value(self.evaluation_config, ["decoding"], {})
+        self.metrics = _get_config_value(self.evaluation_config, ["metrics"], ["bleu", "comet", "chrf", "ter"])
+        self.strategy = self.decoding_config.get("strategy", "beam_search") if isinstance(self.decoding_config, dict) else getattr(self.decoding_config, "strategy", "beam_search")
+
+        self.bos_id = self.tokenizer.bos_token_id
+        self.eos_id = self.tokenizer.eos_token_id
+        self.pad_id = self.tokenizer.pad_token_id
+
+        self.max_decode_len = self.decoding_config.get("max_decode_len", 256) if isinstance(self.decoding_config, dict) else getattr(self.decoding_config, "max_decode_len", 256)
+        self.beam_size = self.decoding_config.get("beam_size", 5) if isinstance(self.decoding_config, dict) else getattr(self.decoding_config, "beam_size", 5)
+        self.length_penalty = self.decoding_config.get("length_penalty", 1.0) if isinstance(self.decoding_config, dict) else getattr(self.decoding_config, "length_penalty", 1.0)
+        self.no_repeat_ngram_size = self.decoding_config.get("no_repeat_ngram_size", 0) if isinstance(self.decoding_config, dict) else getattr(self.decoding_config, "no_repeat_ngram_size", 0)
+        self.temperature = self.decoding_config.get("sampling", {}).get("temperature", 1.0) if isinstance(self.decoding_config, dict) else getattr(getattr(self.decoding_config, "sampling", {}), "temperature", 1.0)
+        self.top_k = self.decoding_config.get("sampling", {}).get("top_k", 0) if isinstance(self.decoding_config, dict) else getattr(getattr(self.decoding_config, "sampling", {}), "top_k", 0)
+        self.top_p = self.decoding_config.get("sampling", {}).get("top_p", 1.0) if isinstance(self.decoding_config, dict) else getattr(getattr(self.decoding_config, "sampling", {}), "top_p", 1.0)
+
+        self.use_generate = hasattr(self.model, "generate") and not (hasattr(self.model, "encode") and hasattr(self.model, "decode_step"))
+
+    def _decode(self, src_ids: torch.Tensor, src_padding_mask: torch.BoolTensor) -> torch.Tensor:
+        if self.use_generate:
+            generate_kwargs = {
+                "max_length": self.max_decode_len,
+                "early_stopping": True,
+            }
+            if self.strategy == "beam_search":
+                generate_kwargs.update(
+                    {
+                        "num_beams": self.beam_size,
+                        "length_penalty": self.length_penalty,
+                        "no_repeat_ngram_size": self.no_repeat_ngram_size,
+                    }
+                )
+            elif self.strategy == "sampling":
+                generate_kwargs.update(
+                    {
+                        "do_sample": True,
+                        "temperature": self.temperature,
+                        "top_k": self.top_k,
+                        "top_p": self.top_p,
+                        "num_beams": 1,
+                    }
+                )
+            else:
+                generate_kwargs.update({"num_beams": 1})
+
+            attention_mask = (~src_padding_mask).long()
+            return self.model.generate(input_ids=src_ids, attention_mask=attention_mask, **generate_kwargs)
+
+        if self.strategy == "beam_search":
+            return beam_search_decode(
+                self.model,
+                src_ids,
+                src_padding_mask,
+                self.bos_id,
+                self.eos_id,
+                beam_size=self.beam_size,
+                max_len=self.max_decode_len,
+                length_penalty=self.length_penalty,
+                no_repeat_ngram_size=self.no_repeat_ngram_size,
+            )
+        if self.strategy == "sampling":
+            return sample_decode(
+                self.model,
+                src_ids,
+                src_padding_mask,
+                self.bos_id,
+                self.eos_id,
+                max_len=self.max_decode_len,
+                temperature=self.temperature,
+                top_k=self.top_k,
+                top_p=self.top_p,
+            )
+        return greedy_decode(
+            self.model,
+            src_ids,
+            src_padding_mask,
+            self.bos_id,
+            self.eos_id,
+            max_len=self.max_decode_len,
+        )
+
+    def evaluate(
+        self,
+        dataloader: DataLoader,
+        src_texts: Optional[list[str]] = None,
+        ref_texts: Optional[list[str]] = None,
+    ) -> dict:
+        self.model.eval()
+        device = next(self.model.parameters()).device if any(p.requires_grad or p.is_floating_point() for p in self.model.parameters()) else torch.device("cpu")
+
+        if src_texts is None and hasattr(dataloader, "dataset") and hasattr(dataloader.dataset, "src_texts"):
+            src_texts = list(dataloader.dataset.src_texts)
+        if ref_texts is None and hasattr(dataloader, "dataset") and hasattr(dataloader.dataset, "tgt_texts"):
+            ref_texts = list(dataloader.dataset.tgt_texts)
+
+        if src_texts is None or ref_texts is None:
+            raise ValueError("Source texts and reference texts must be provided for evaluation.")
+
+        hypotheses: list[str] = []
+        sources: list[str] = []
+        references: list[str] = []
+
+        for batch_idx, batch in enumerate(tqdm(dataloader, desc="Evaluating", unit="batch")):
+            src_ids = batch["src_ids"].to(device)
+            src_padding_mask = batch.get("src_padding_mask")
+            if src_padding_mask is None:
+                src_padding_mask = src_ids.eq(self.pad_id)
+            else:
+                src_padding_mask = src_padding_mask.to(device)
+
+            output_ids = self._decode(src_ids, src_padding_mask)
+            if isinstance(output_ids, torch.Tensor):
+                output_ids = output_ids.cpu()
+
+            for sample_idx in range(output_ids.size(0)):
+                decoded = self.tokenizer.decode(output_ids[sample_idx].tolist(), skip_special_tokens=True)
+                hypotheses.append(decoded)
+
+        sources = src_texts
+        references = ref_texts
+        results = compute_all_metrics(sources, hypotheses, references, metrics=self.metrics)
+        return results
+
+    def translate(self, texts: list[str]) -> list[str]:
+        self.model.eval()
+        device = next(self.model.parameters()).device if any(p.requires_grad or p.is_floating_point() for p in self.model.parameters()) else torch.device("cpu")
+
+        input_ids = []
+        for text in texts:
+            src_tokens = self.tokenizer.encode(
+                text,
+                add_special_tokens=True,
+                max_length=_get_config_value(self.config, ["data", "preprocessing", "max_src_len"], 256),
+            )
+            input_ids.append(torch.tensor(src_tokens, dtype=torch.long, device=device))
+
+        src_ids = pad_sequence(input_ids, batch_first=True, padding_value=self.pad_id)
+        src_padding_mask = src_ids.eq(self.pad_id)
+
+        output_ids = self._decode(src_ids, src_padding_mask)
+        if isinstance(output_ids, torch.Tensor):
+            output_ids = output_ids.cpu()
+
+        translations: list[str] = []
+        for sample_idx in range(output_ids.size(0)):
+            translations.append(self.tokenizer.decode(output_ids[sample_idx].tolist(), skip_special_tokens=True))
+
+        return translations
+
+    def translate_single(self, text: str) -> str:
+        return self.translate([text])[0]

src/easytranslate/evaluation/metrics.py

@@ -1,109 +1,169 @@
-"""
-评估指标模块 — Person D 负责实现
-
-功能要求:
-1. compute_bleu: 计算 SacreBLEU 分数
-2. compute_comet: 计算 COMET 分数 (神经网络评估指标)
-3. compute_chrf: 计算 chrF++ 分数
-4. compute_ter: 计算 TER (Translation Edit Rate)
-5. compute_all_metrics: 计算所有指标
-
-技术要点:
-- SacreBLEU: 标准化的 BLEU 实现，结果可复现
-- COMET: 基于预训练语言模型的评估指标，与人类评价相关性最高
-- chrF++: 基于字符 n-gram 的 F-score，对中文尤其有用
-- TER: 编辑距离，衡量翻译后编辑量
-"""
-
-from __future__ import annotations
-
-import logging
-from typing import Optional
-
-logger = logging.getLogger(__name__)
-
-
-def compute_bleu(
-    hypotheses: list[str],
-    references: list[str],
-    tokenize: str = "zh",
-) -> dict:
-    """
-    计算 SacreBLEU 分数。
-
-    TODO [Person D]: 实现以下逻辑:
-    1. 使用 sacrebleu.corpus_bleu(hypotheses, [references], tokenize=tokenize)
-    2. tokenize="zh" 对中文进行字符级别分词
-    3. 返回 {"bleu": score, "bleu_1": ..., "bleu_2": ..., "bleu_3": ..., "bleu_4": ..., "bp": ...}
-
-    注意: references 需要包装为 list of list (支持多参考)
-    """
-    raise NotImplementedError("TODO: Person D 实现 compute_bleu")
-
-
-def compute_comet(
-    sources: list[str],
-    hypotheses: list[str],
-    references: list[str],
-    model_name: str = "Unbabel/wmt22-comet-da",
-    batch_size: int = 16,
-    gpus: int = 1,
-) -> dict:
-    """
-    计算 COMET 分数。
-
-    TODO [Person D]: 实现以下逻辑:
-    1. 加载 COMET 模型: comet.download_model(model_name)
-    2. 构建输入数据: [{"src": s, "mt": h, "ref": r} for s, h, r in zip(...)]
-    3. 调用 model.predict(data, batch_size, gpus)
-    4. 返回 {"comet": system_score, "comet_scores": segment_scores}
-
-    COMET 需要源语言、翻译结果和参考翻译三者。
-    """
-    raise NotImplementedError("TODO: Person D 实现 compute_comet")
-
-
-def compute_chrf(
-    hypotheses: list[str],
-    references: list[str],
-) -> dict:
-    """
-    计算 chrF++ 分数。
-
-    TODO [Person D]:
-    1. 使用 sacrebleu.corpus_chrf(hypotheses, [references])
-    2. 返回 {"chrf": score}
-    """
-    raise NotImplementedError("TODO: Person D 实现 compute_chrf")
-
-
-def compute_ter(
-    hypotheses: list[str],
-    references: list[str],
-) -> dict:
-    """
-    计算 TER 分数。
-
-    TODO [Person D]:
-    1. 使用 sacrebleu.corpus_ter(hypotheses, [references])
-    2. 返回 {"ter": score}
-    """
-    raise NotImplementedError("TODO: Person D 实现 compute_ter")
-
-
-def compute_all_metrics(
-    sources: list[str],
-    hypotheses: list[str],
-    references: list[str],
-    metrics: list[str] = ["bleu", "comet", "chrf", "ter"],
-) -> dict:
-    """
-    计算所有指定的评估指标。
-
-    TODO [Person D]:
-    1. 遍历 metrics 列表
-    2. 调用对应的计算函数
-    3. 合并结果并返回
-    4. 记录每个指标的计算时间
-    """
-    raise NotImplementedError("TODO: Person D 实现 compute_all_metrics")
+"""
+评估指标模块 — Person D 负责实现
+
+功能要求:
+1. compute_bleu: 计算 SacreBLEU 分数
+2. compute_comet: 计算 COMET 分数 (神经网络评估指标)
+3. compute_chrf: 计算 chrF++ 分数
+4. compute_ter: 计算 TER (Translation Edit Rate)
+5. compute_all_metrics: 计算所有指标
+
+技术要点:
+- SacreBLEU: 标准化的 BLEU 实现，结果可复现
+- COMET: 基于预训练语言模型的评估指标，与人类评价相关性最高
+- chrF++: 基于字符 n-gram 的 F-score，对中文尤其有用
+- TER: 编辑距离，衡量翻译后编辑量
+"""
+
+from __future__ import annotations
+
+import logging
+import time
+from typing import Optional
+
+logger = logging.getLogger(__name__)
+
+
+def _import_sacrebleu():
+    try:
+        import sacrebleu
+    except ImportError as exc:
+        logger.error("SacreBLEU is not installed: %s", exc)
+        raise ImportError(
+            "sacrebleu is required for BLEU/chrF/TER evaluation. "
+            "Install it with `python -m pip install sacrebleu`."
+        ) from exc
+    return sacrebleu
+
+
+def compute_bleu(
+    hypotheses: list[str],
+    references: list[str],
+    tokenize: str = "zh",
+) -> dict:
+    """
+    计算 SacreBLEU 分数。
+    """
+    if not hypotheses or not references:
+        return {
+            "bleu": 0.0,
+            "bleu_1": 0.0,
+            "bleu_2": 0.0,
+            "bleu_3": 0.0,
+            "bleu_4": 0.0,
+            "bp": 0.0,
+        }
+
+    sacrebleu = _import_sacrebleu()
+    references_list = [references]
+    bleu = sacrebleu.corpus_bleu(hypotheses, references_list, tokenize=tokenize)
+    precisions = [round(float(x), 4) for x in bleu.precisions]
+    return {
+        "bleu": round(float(bleu.score), 4),
+        "bleu_1": precisions[0] if len(precisions) > 0 else 0.0,
+        "bleu_2": precisions[1] if len(precisions) > 1 else 0.0,
+        "bleu_3": precisions[2] if len(precisions) > 2 else 0.0,
+        "bleu_4": precisions[3] if len(precisions) > 3 else 0.0,
+        "bp": round(float(getattr(bleu, "bp", 0.0)), 4),
+    }
+
+
+def compute_comet(
+    sources: list[str],
+    hypotheses: list[str],
+    references: list[str],
+    model_name: str = "Unbabel/wmt22-comet-da",
+    batch_size: int = 16,
+    gpus: int = 1,
+) -> dict:
+    """
+    计算 COMET 分数。
+    """
+    if not sources or not hypotheses or not references:
+        return {"comet": 0.0, "comet_scores": []}
+
+    try:
+        from comet import download_model, load_from_checkpoint
+    except ImportError as exc:
+        logger.error("COMET library is not installed: %s", exc)
+        raise
+
+    model_path = download_model(model_name)
+    model = load_from_checkpoint(model_path)
+
+    data = [
+        {"src": src, "mt": hyp, "ref": ref}
+        for src, hyp, ref in zip(sources, hypotheses, references)
+    ]
+    prediction = model.predict(data, batch_size=batch_size, gpus=gpus)
+
+    if isinstance(prediction, dict):
+        scores = prediction.get("scores") or prediction.get("predictions") or []
+    else:
+        scores = list(prediction)
+
+    if scores is None:
+        scores = []
+
+    scores = [float(score) for score in scores]
+    system_score = float(sum(scores) / len(scores)) if scores else 0.0
+    return {"comet": round(system_score, 4), "comet_scores": scores}
+
+
+def compute_chrf(
+    hypotheses: list[str],
+    references: list[str],
+) -> dict:
+    """
+    计算 chrF++ 分数。
+    """
+    if not hypotheses or not references:
+        return {"chrf": 0.0}
+
+    sacrebleu = _import_sacrebleu()
+    chrf = sacrebleu.corpus_chrf(hypotheses, [references])
+    return {"chrf": round(float(chrf.score), 4)}
+
+
+def compute_ter(
+    hypotheses: list[str],
+    references: list[str],
+) -> dict:
+    """
+    计算 TER 分数。
+    """
+    if not hypotheses or not references:
+        return {"ter": 0.0}
+
+    sacrebleu = _import_sacrebleu()
+    ter = sacrebleu.corpus_ter(hypotheses, [references])
+    return {"ter": round(float(ter.score), 4)}
+
+
+def compute_all_metrics(
+    sources: list[str],
+    hypotheses: list[str],
+    references: list[str],
+    metrics: list[str] = ["bleu", "comet", "chrf", "ter"],
+) -> dict:
+    """
+    计算所有指定的评估指标。
+    """
+    results: dict[str, float | list[float]] = {}
+
+    for metric in metrics:
+        start_time = time.time()
+        if metric == "bleu":
+            results.update(compute_bleu(hypotheses, references))
+        elif metric == "comet":
+            results.update(compute_comet(sources, hypotheses, references))
+        elif metric == "chrf":
+            results.update(compute_chrf(hypotheses, references))
+        elif metric == "ter":
+            results.update(compute_ter(hypotheses, references))
+        else:
+            logger.warning("Unsupported evaluation metric: %s", metric)
+        elapsed = time.time() - start_time
+        logger.info("Computed %s in %.2f seconds", metric, elapsed)
+
+    return results

src/easytranslate/model/__init__.py

@@ -1,20 +1,20 @@
-"""模型架构模块 (Person B 负责)"""
-
-from easytranslate.model.transformer import TransformerTranslationModel
-from easytranslate.model.encoder import TransformerEncoder
-from easytranslate.model.decoder import TransformerDecoder
-from easytranslate.model.attention import MultiHeadAttention, FlashMultiHeadAttention
-from easytranslate.model.positional import SinusoidalPositionalEncoding, RotaryPositionalEmbedding
-from easytranslate.model.finetune import load_pretrained_model, setup_lora
-
-__all__ = [
-    "TransformerTranslationModel",
-    "TransformerEncoder",
-    "TransformerDecoder",
-    "MultiHeadAttention",
-    "FlashMultiHeadAttention",
-    "SinusoidalPositionalEncoding",
-    "RotaryPositionalEmbedding",
-    "load_pretrained_model",
-    "setup_lora",
-]
+"""模型架构模块 (Person B 负责)"""
+
+from easytranslate.model.transformer import TransformerTranslationModel
+from easytranslate.model.encoder import TransformerEncoder
+from easytranslate.model.decoder import TransformerDecoder
+from easytranslate.model.attention import MultiHeadAttention, FlashMultiHeadAttention
+from easytranslate.model.positional import SinusoidalPositionalEncoding, RotaryPositionalEmbedding
+from easytranslate.model.finetune import load_pretrained_model, setup_lora
+
+__all__ = [
+    "TransformerTranslationModel",
+    "TransformerEncoder",
+    "TransformerDecoder",
+    "MultiHeadAttention",
+    "FlashMultiHeadAttention",
+    "SinusoidalPositionalEncoding",
+    "RotaryPositionalEmbedding",
+    "load_pretrained_model",
+    "setup_lora",
+]

src/easytranslate/model/attention.py

@@ -1,231 +1,231 @@
-"""
-注意力机制模块 — Person B 负责实现
-
-包含:
-1. MultiHeadAttention: 标准多头注意力 (支持 RoPE)
-2. FlashMultiHeadAttention: Flash Attention 2 加速版本
-
-技术要点:
-- Scaled Dot-Product Attention
-- 支持 key_padding_mask 和 attn_mask
-- Flash Attention 2 使用 torch.nn.functional.scaled_dot_product_attention
-- RoPE 旋转位置编码集成
-"""
-
-from __future__ import annotations
-
-import math
-from typing import Optional
-
-import torch
-import torch.nn as nn
-import torch.nn.functional as F
-
-
-class MultiHeadAttention(nn.Module):
-    """
-    标准多头注意力机制。
-
-    TODO [Person B]: 实现以下内容:
-
-    __init__:
-    1. Q, K, V 线性投影: nn.Linear(d_model, d_model)
-    2. 输出投影: nn.Linear(d_model, d_model)
-    3. Dropout
-
-    forward(query, key, value, key_padding_mask=None, attn_mask=None):
-    1. 线性投影 Q, K, V
-    2. reshape 为 [B, nhead, L, d_k]
-    3. (可选) 应用 RoPE 旋转位置编码
-    4. 计算 attention scores: QK^T / sqrt(d_k)
-    5. 应用 masks (padding mask + causal mask)
-    6. Softmax + Dropout
-    7. 加权求和 V
-    8. reshape 回 [B, L, d_model]
-    9. 输出投影
-    """
-
-    def __init__(
-        self,
-        d_model: int = 512,
-        nhead: int = 8,
-        dropout: float = 0.1,
-        use_rotary_embedding: bool = False,
-    ):
-        super().__init__()
-        assert d_model % nhead == 0, "d_model 必须能被 nhead 整除"
-        self.d_model = d_model
-        self.nhead = nhead
-        self.d_k = d_model // nhead
-        self.use_rotary_embedding = use_rotary_embedding
-
-        self.q_proj = nn.Linear(d_model, d_model)
-        self.k_proj = nn.Linear(d_model, d_model)
-        self.v_proj = nn.Linear(d_model, d_model)
-        self.out_proj = nn.Linear(d_model, d_model)
-
-        self.dropout = nn.Dropout(p=dropout)
-        self.rope: Optional[nn.Module] = None
-
-    def forward(
-        self,
-        query: torch.Tensor,   # [B, L_q, D]
-        key: torch.Tensor,     # [B, L_k, D]
-        value: torch.Tensor,   # [B, L_v, D]
-        key_padding_mask: Optional[torch.BoolTensor] = None,  # [B, L_k]
-        attn_mask: Optional[torch.Tensor] = None,             # [L_q, L_k]
-        is_causal: bool = False,
-    ) -> torch.Tensor:
-        B, L_q, _ = query.size()
-        L_k = key.size(1)
-        L_v = value.size(1)
-
-        # 1. 线性投影
-        Q = self.q_proj(query)  # [B, L_q, D]
-        K = self.k_proj(key)    # [B, L_k, D]
-        V = self.v_proj(value)  # [B, L_v, D]
-
-        # 2. reshape 为 [B, nhead, L, d_k]
-        Q = Q.view(B, L_q, self.nhead, self.d_k).transpose(1, 2)  # [B, H, L_q, d_k]
-        K = K.view(B, L_k, self.nhead, self.d_k).transpose(1, 2)  # [B, H, L_k, d_k]
-        V = V.view(B, L_v, self.nhead, self.d_k).transpose(1, 2)  # [B, H, L_v, d_k]
-
-        # 3. (可选) 应用 RoPE
-        if self.use_rotary_embedding and self.rope is not None:
-            Q, K = self.rope.apply_rotary_pos_emb(Q, K)
-
-        # 4. 计算 attention scores: QK^T / sqrt(d_k)
-        scores = torch.matmul(Q, K.transpose(-2, -1)) / math.sqrt(self.d_k)  # [B, H, L_q, L_k]
-
-        # 5. 应用 masks
-        if key_padding_mask is not None:
-            # key_padding_mask: [B, L_k] -> [B, 1, 1, L_k]
-            scores = scores.masked_fill(
-                key_padding_mask.unsqueeze(1).unsqueeze(2), float("-inf")
-            )
-        if is_causal:
-            # 生成 causal mask
-            L_q, L_k_local = scores.size(-2), scores.size(-1)
-            causal_mask = torch.triu(
-                torch.ones(L_q, L_k_local, device=scores.device), diagonal=1
-            ).bool()
-            scores = scores.masked_fill(
-                causal_mask.unsqueeze(0).unsqueeze(0), float("-inf")
-            )
-        if attn_mask is not None:
-            # attn_mask: [L_q, L_k] -> [1, 1, L_q, L_k]
-            scores = scores.masked_fill(attn_mask.unsqueeze(0).unsqueeze(0), float("-inf"))
-
-        # 6. Softmax + Dropout
-        attn_weights = F.softmax(scores, dim=-1)
-        attn_weights = self.dropout(attn_weights)
-
-        # 7. 加权求和 V
-        attn_output = torch.matmul(attn_weights, V)  # [B, H, L_q, d_k]
-
-        # 8. reshape 回 [B, L_q, d_model]
-        attn_output = attn_output.transpose(1, 2).contiguous().view(B, L_q, self.d_model)
-
-        # 9. 输出投影
-        output = self.out_proj(attn_output)
-        return output
-
-
-class FlashMultiHeadAttention(nn.Module):
-    """
-    Flash Attention 2 加速的多头注意力。
-
-    TODO [Person B]: 使用 PyTorch 2.0+ 的 F.scaled_dot_product_attention 实现:
-    1. 与 MultiHeadAttention 结构相同
-    2. 在 forward 中使用 F.scaled_dot_product_attention(Q, K, V, attn_mask, dropout, is_causal)
-    3. 会自动选择最优的 attention kernel (Flash Attention / Memory-Efficient Attention)
-
-    注意:
-    - 需要 PyTorch >= 2.0
-    - is_causal=True 时自动生成因果掩码，不需要手动传入 attn_mask
-    """
-
-    def __init__(
-        self,
-        d_model: int = 512,
-        nhead: int = 8,
-        dropout: float = 0.1,
-        use_rotary_embedding: bool = False,
-    ):
-        super().__init__()
-        assert d_model % nhead == 0, "d_model 必须能被 nhead 整除"
-        self.d_model = d_model
-        self.nhead = nhead
-        self.d_k = d_model // nhead
-        self.use_rotary_embedding = use_rotary_embedding
-        self.dropout_p = dropout
-
-        self.q_proj = nn.Linear(d_model, d_model)
-        self.k_proj = nn.Linear(d_model, d_model)
-        self.v_proj = nn.Linear(d_model, d_model)
-        self.out_proj = nn.Linear(d_model, d_model)
-
-        self.rope: Optional[nn.Module] = None
-
-    def forward(
-        self,
-        query: torch.Tensor,
-        key: torch.Tensor,
-        value: torch.Tensor,
-        key_padding_mask: Optional[torch.BoolTensor] = None,
-        is_causal: bool = False,
-    ) -> torch.Tensor:
-        B, L_q, _ = query.size()
-        L_k = key.size(1)
-        L_v = value.size(1)
-
-        # 1. 线性投影
-        Q = self.q_proj(query)  # [B, L_q, D]
-        K = self.k_proj(key)    # [B, L_k, D]
-        V = self.v_proj(value)  # [B, L_v, D]
-
-        # 2. reshape 为 [B, nhead, L, d_k]
-        Q = Q.view(B, L_q, self.nhead, self.d_k).transpose(1, 2)  # [B, H, L_q, d_k]
-        K = K.view(B, L_k, self.nhead, self.d_k).transpose(1, 2)  # [B, H, L_k, d_k]
-        V = V.view(B, L_v, self.nhead, self.d_k).transpose(1, 2)  # [B, H, L_v, d_k]
-
-        # 3. (可选) 应用 RoPE
-        if self.use_rotary_embedding and self.rope is not None:
-            Q, K = self.rope.apply_rotary_pos_emb(Q, K)
-
-        # 4. 构建 attn_mask 以适配 scaled_dot_product_attention
-        # PyTorch >= 2.0 支持 [B, nhead, L, d_k] 的 4D 输入
-        # 注意: scaled_dot_product_attention 不允许同时设置 attn_mask 和 is_causal=True
-        attn_mask: Optional[torch.Tensor] = None
-        if is_causal or key_padding_mask is not None:
-            attn_mask = torch.zeros(
-                B, self.nhead, L_q, L_k, dtype=Q.dtype, device=Q.device
-            )
-            if is_causal:
-                # 生成 causal mask (上三角为 -inf)
-                causal_mask = torch.triu(
-                    torch.ones(L_q, L_k, device=Q.device), diagonal=1
-                ).bool()
-                attn_mask = attn_mask.masked_fill(
-                    causal_mask[None, None, :, :], float("-inf")
-                )
-            if key_padding_mask is not None:
-                # key_padding_mask: True = padding (忽略)
-                _bool_mask = key_padding_mask.unsqueeze(1).unsqueeze(2)
-                _bool_mask = _bool_mask.expand(B, self.nhead, L_q, L_k)
-                attn_mask = attn_mask.masked_fill(_bool_mask, float("-inf"))
-
-        # 5. Flash Attention (PyTorch 原生)
-        attn_output = F.scaled_dot_product_attention(
-            Q, K, V,
-            attn_mask=attn_mask,
-            dropout_p=self.dropout_p if self.training else 0.0,
-            is_causal=False,  # 已通过 attn_mask 处理
-        )  # [B, H, L_q, d_k]
-
-        # 6. reshape 回 [B, L_q, d_model]
-        attn_output = attn_output.transpose(1, 2).contiguous().view(B, L_q, self.d_model)
-
-        # 7. 输出投影
-        output = self.out_proj(attn_output)
-        return output
+"""
+注意力机制模块 — Person B 负责实现
+
+包含:
+1. MultiHeadAttention: 标准多头注意力 (支持 RoPE)
+2. FlashMultiHeadAttention: Flash Attention 2 加速版本
+
+技术要点:
+- Scaled Dot-Product Attention
+- 支持 key_padding_mask 和 attn_mask
+- Flash Attention 2 使用 torch.nn.functional.scaled_dot_product_attention
+- RoPE 旋转位置编码集成
+"""
+
+from __future__ import annotations
+
+import math
+from typing import Optional
+
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+
+
+class MultiHeadAttention(nn.Module):
+    """
+    标准多头注意力机制。
+
+    TODO [Person B]: 实现以下内容:
+
+    __init__:
+    1. Q, K, V 线性投影: nn.Linear(d_model, d_model)
+    2. 输出投影: nn.Linear(d_model, d_model)
+    3. Dropout
+
+    forward(query, key, value, key_padding_mask=None, attn_mask=None):
+    1. 线性投影 Q, K, V
+    2. reshape 为 [B, nhead, L, d_k]
+    3. (可选) 应用 RoPE 旋转位置编码
+    4. 计算 attention scores: QK^T / sqrt(d_k)
+    5. 应用 masks (padding mask + causal mask)
+    6. Softmax + Dropout
+    7. 加权求和 V
+    8. reshape 回 [B, L, d_model]
+    9. 输出投影
+    """
+
+    def __init__(
+        self,
+        d_model: int = 512,
+        nhead: int = 8,
+        dropout: float = 0.1,
+        use_rotary_embedding: bool = False,
+    ):
+        super().__init__()
+        assert d_model % nhead == 0, "d_model 必须能被 nhead 整除"
+        self.d_model = d_model
+        self.nhead = nhead
+        self.d_k = d_model // nhead
+        self.use_rotary_embedding = use_rotary_embedding
+
+        self.q_proj = nn.Linear(d_model, d_model)
+        self.k_proj = nn.Linear(d_model, d_model)
+        self.v_proj = nn.Linear(d_model, d_model)
+        self.out_proj = nn.Linear(d_model, d_model)
+
+        self.dropout = nn.Dropout(p=dropout)
+        self.rope: Optional[nn.Module] = None
+
+    def forward(
+        self,
+        query: torch.Tensor,   # [B, L_q, D]
+        key: torch.Tensor,     # [B, L_k, D]
+        value: torch.Tensor,   # [B, L_v, D]
+        key_padding_mask: Optional[torch.BoolTensor] = None,  # [B, L_k]
+        attn_mask: Optional[torch.Tensor] = None,             # [L_q, L_k]
+        is_causal: bool = False,
+    ) -> torch.Tensor:
+        B, L_q, _ = query.size()
+        L_k = key.size(1)
+        L_v = value.size(1)
+
+        # 1. 线性投影
+        Q = self.q_proj(query)  # [B, L_q, D]
+        K = self.k_proj(key)    # [B, L_k, D]
+        V = self.v_proj(value)  # [B, L_v, D]
+
+        # 2. reshape 为 [B, nhead, L, d_k]
+        Q = Q.view(B, L_q, self.nhead, self.d_k).transpose(1, 2)  # [B, H, L_q, d_k]
+        K = K.view(B, L_k, self.nhead, self.d_k).transpose(1, 2)  # [B, H, L_k, d_k]
+        V = V.view(B, L_v, self.nhead, self.d_k).transpose(1, 2)  # [B, H, L_v, d_k]
+
+        # 3. (可选) 应用 RoPE
+        if self.use_rotary_embedding and self.rope is not None:
+            Q, K = self.rope.apply_rotary_pos_emb(Q, K)
+
+        # 4. 计算 attention scores: QK^T / sqrt(d_k)
+        scores = torch.matmul(Q, K.transpose(-2, -1)) / math.sqrt(self.d_k)  # [B, H, L_q, L_k]
+
+        # 5. 应用 masks
+        if key_padding_mask is not None:
+            # key_padding_mask: [B, L_k] -> [B, 1, 1, L_k]
+            scores = scores.masked_fill(
+                key_padding_mask.unsqueeze(1).unsqueeze(2), float("-inf")
+            )
+        if is_causal:
+            # 生成 causal mask
+            L_q, L_k_local = scores.size(-2), scores.size(-1)
+            causal_mask = torch.triu(
+                torch.ones(L_q, L_k_local, device=scores.device), diagonal=1
+            ).bool()
+            scores = scores.masked_fill(
+                causal_mask.unsqueeze(0).unsqueeze(0), float("-inf")
+            )
+        if attn_mask is not None:
+            # attn_mask: [L_q, L_k] -> [1, 1, L_q, L_k]
+            scores = scores.masked_fill(attn_mask.unsqueeze(0).unsqueeze(0), float("-inf"))
+
+        # 6. Softmax + Dropout
+        attn_weights = F.softmax(scores, dim=-1)
+        attn_weights = self.dropout(attn_weights)
+
+        # 7. 加权求和 V
+        attn_output = torch.matmul(attn_weights, V)  # [B, H, L_q, d_k]
+
+        # 8. reshape 回 [B, L_q, d_model]
+        attn_output = attn_output.transpose(1, 2).contiguous().view(B, L_q, self.d_model)
+
+        # 9. 输出投影
+        output = self.out_proj(attn_output)
+        return output
+
+
+class FlashMultiHeadAttention(nn.Module):
+    """
+    Flash Attention 2 加速的多头注意力。
+
+    TODO [Person B]: 使用 PyTorch 2.0+ 的 F.scaled_dot_product_attention 实现:
+    1. 与 MultiHeadAttention 结构相同
+    2. 在 forward 中使用 F.scaled_dot_product_attention(Q, K, V, attn_mask, dropout, is_causal)
+    3. 会自动选择最优的 attention kernel (Flash Attention / Memory-Efficient Attention)
+
+    注意:
+    - 需要 PyTorch >= 2.0
+    - is_causal=True 时自动生成因果掩码，不需要手动传入 attn_mask
+    """
+
+    def __init__(
+        self,
+        d_model: int = 512,
+        nhead: int = 8,
+        dropout: float = 0.1,
+        use_rotary_embedding: bool = False,
+    ):
+        super().__init__()
+        assert d_model % nhead == 0, "d_model 必须能被 nhead 整除"
+        self.d_model = d_model
+        self.nhead = nhead
+        self.d_k = d_model // nhead
+        self.use_rotary_embedding = use_rotary_embedding
+        self.dropout_p = dropout
+
+        self.q_proj = nn.Linear(d_model, d_model)
+        self.k_proj = nn.Linear(d_model, d_model)
+        self.v_proj = nn.Linear(d_model, d_model)
+        self.out_proj = nn.Linear(d_model, d_model)
+
+        self.rope: Optional[nn.Module] = None
+
+    def forward(
+        self,
+        query: torch.Tensor,
+        key: torch.Tensor,
+        value: torch.Tensor,
+        key_padding_mask: Optional[torch.BoolTensor] = None,
+        is_causal: bool = False,
+    ) -> torch.Tensor:
+        B, L_q, _ = query.size()
+        L_k = key.size(1)
+        L_v = value.size(1)
+
+        # 1. 线性投影
+        Q = self.q_proj(query)  # [B, L_q, D]
+        K = self.k_proj(key)    # [B, L_k, D]
+        V = self.v_proj(value)  # [B, L_v, D]
+
+        # 2. reshape 为 [B, nhead, L, d_k]
+        Q = Q.view(B, L_q, self.nhead, self.d_k).transpose(1, 2)  # [B, H, L_q, d_k]
+        K = K.view(B, L_k, self.nhead, self.d_k).transpose(1, 2)  # [B, H, L_k, d_k]
+        V = V.view(B, L_v, self.nhead, self.d_k).transpose(1, 2)  # [B, H, L_v, d_k]
+
+        # 3. (可选) 应��� RoPE
+        if self.use_rotary_embedding and self.rope is not None:
+            Q, K = self.rope.apply_rotary_pos_emb(Q, K)
+
+        # 4. 构建 attn_mask 以适配 scaled_dot_product_attention
+        # PyTorch >= 2.0 支持 [B, nhead, L, d_k] 的 4D 输入
+        # 注意: scaled_dot_product_attention 不允许同时设置 attn_mask 和 is_causal=True
+        attn_mask: Optional[torch.Tensor] = None
+        if is_causal or key_padding_mask is not None:
+            attn_mask = torch.zeros(
+                B, self.nhead, L_q, L_k, dtype=Q.dtype, device=Q.device
+            )
+            if is_causal:
+                # 生成 causal mask (上三角为 -inf)
+                causal_mask = torch.triu(
+                    torch.ones(L_q, L_k, device=Q.device), diagonal=1
+                ).bool()
+                attn_mask = attn_mask.masked_fill(
+                    causal_mask[None, None, :, :], float("-inf")
+                )
+            if key_padding_mask is not None:
+                # key_padding_mask: True = padding (忽略)
+                _bool_mask = key_padding_mask.unsqueeze(1).unsqueeze(2)
+                _bool_mask = _bool_mask.expand(B, self.nhead, L_q, L_k)
+                attn_mask = attn_mask.masked_fill(_bool_mask, float("-inf"))
+
+        # 5. Flash Attention (PyTorch 原生)
+        attn_output = F.scaled_dot_product_attention(
+            Q, K, V,
+            attn_mask=attn_mask,
+            dropout_p=self.dropout_p if self.training else 0.0,
+            is_causal=False,  # 已通过 attn_mask 处理
+        )  # [B, H, L_q, d_k]
+
+        # 6. reshape 回 [B, L_q, d_model]
+        attn_output = attn_output.transpose(1, 2).contiguous().view(B, L_q, self.d_model)
+
+        # 7. 输出投影
+        output = self.out_proj(attn_output)
+        return output

src/easytranslate/model/decoder.py

@@ -1,176 +1,176 @@
-"""
-Transformer Decoder 模块 — Person B 负责实现
-
-包含:
-- TransformerDecoderLayer: 单层解码器
-- TransformerDecoder: 多层解码器堆叠
-
-架构 (Pre-LayerNorm):
-    x → LN → Masked Self-Attention → Residual
-      → LN → Cross-Attention → Residual
-      → LN → FFN → Residual
-"""
-
-from __future__ import annotations
-
-import copy
-
-import torch
-import torch.nn as nn
-from typing import Optional
-
-from easytranslate.model.attention import MultiHeadAttention, FlashMultiHeadAttention
-
-
-class TransformerDecoderLayer(nn.Module):
-    """
-    单层 Transformer Decoder。
-
-    架构 (Pre-LayerNorm):
-        x → LN → Masked Self-Attention → Residual
-          → LN → Cross-Attention → Residual
-          → LN → FFN → Residual
-    """
-
-    def __init__(
-        self,
-        d_model: int = 512,
-        nhead: int = 8,
-        dim_feedforward: int = 2048,
-        dropout: float = 0.1,
-        activation: str = "gelu",
-        use_flash_attention: bool = True,
-        use_rotary_embedding: bool = True,
-        pre_norm: bool = True,
-    ):
-        super().__init__()
-        self.pre_norm = pre_norm
-        self.d_model = d_model
-
-        attn_cls = FlashMultiHeadAttention if use_flash_attention else MultiHeadAttention
-
-        # 1. Masked Self-Attention
-        self.self_attn = attn_cls(
-            d_model=d_model,
-            nhead=nhead,
-            dropout=dropout,
-            use_rotary_embedding=use_rotary_embedding,
-        )
-
-        # 2. Cross-Attention (decoder queries encoder memory)
-        self.multihead_attn = attn_cls(
-            d_model=d_model,
-            nhead=nhead,
-            dropout=dropout,
-            use_rotary_embedding=False,  # cross-attention 不使用 RoPE
-        )
-
-        # 3. Feed-Forward Network
-        self.linear1 = nn.Linear(d_model, dim_feedforward)
-        self.activation = nn.GELU() if activation == "gelu" else nn.ReLU()
-        self.dropout = nn.Dropout(p=dropout)
-        self.linear2 = nn.Linear(dim_feedforward, d_model)
-
-        # 4. LayerNorms
-        self.norm1 = nn.LayerNorm(d_model)
-        self.norm2 = nn.LayerNorm(d_model)
-        self.norm3 = nn.LayerNorm(d_model)
-
-        # 5. Dropouts for residuals
-        self.dropout1 = nn.Dropout(p=dropout)
-        self.dropout2 = nn.Dropout(p=dropout)
-        self.dropout3 = nn.Dropout(p=dropout)
-
-    def forward(
-        self,
-        tgt: torch.Tensor,                               # [B, T, D]
-        memory: torch.Tensor,                             # [B, S, D] (encoder output)
-        tgt_mask: Optional[torch.Tensor] = None,          # [T, T] causal mask
-        memory_key_padding_mask: Optional[torch.BoolTensor] = None,  # [B, S]
-        tgt_key_padding_mask: Optional[torch.BoolTensor] = None,     # [B, T]
-    ) -> torch.Tensor:
-        """Pre-LayerNorm 前向传播。"""
-        # 适配 Flash Attention: 使用 is_causal 替代显式 causal mask
-        is_causal = tgt_mask is not None
-
-        if self.pre_norm:
-            # 1. Masked Self-Attention
-            residual = tgt
-            tgt = self.norm1(tgt)
-            tgt = self.self_attn(
-                tgt, tgt, tgt,
-                key_padding_mask=tgt_key_padding_mask,
-                is_causal=is_causal,
-            )
-            tgt = residual + self.dropout1(tgt)
-
-            # 2. Cross-Attention
-            residual = tgt
-            tgt = self.norm2(tgt)
-            tgt = self.multihead_attn(
-                tgt, memory, memory,
-                key_padding_mask=memory_key_padding_mask,
-            )
-            tgt = residual + self.dropout2(tgt)
-
-            # 3. FFN
-            residual = tgt
-            tgt = self.norm3(tgt)
-            tgt = self.linear2(self.dropout(self.activation(self.linear1(tgt))))
-            tgt = residual + self.dropout3(tgt)
-        else:
-            # Post-LayerNorm (备用)
-            residual = tgt
-            tgt = self.self_attn(
-                tgt, tgt, tgt,
-                key_padding_mask=tgt_key_padding_mask,
-                is_causal=is_causal,
-            )
-            tgt = self.norm1(residual + self.dropout1(tgt))
-
-            residual = tgt
-            tgt = self.multihead_attn(
-                tgt, memory, memory,
-                key_padding_mask=memory_key_padding_mask,
-            )
-            tgt = self.norm2(residual + self.dropout2(tgt))
-
-            residual = tgt
-            tgt = self.linear2(self.dropout(self.activation(self.linear1(tgt))))
-            tgt = self.norm3(residual + self.dropout3(tgt))
-
-        return tgt
-
-
-class TransformerDecoder(nn.Module):
-    """
-    多层 Transformer Decoder。
-    """
-
-    def __init__(self, decoder_layer: TransformerDecoderLayer, num_layers: int):
-        super().__init__()
-        self.layers = nn.ModuleList(
-            [copy.deepcopy(decoder_layer) for _ in range(num_layers)]
-        )
-        self.num_layers = num_layers
-        self.norm = nn.LayerNorm(decoder_layer.d_model)
-
-    def forward(
-        self,
-        tgt: torch.Tensor,
-        memory: torch.Tensor,
-        tgt_mask: Optional[torch.Tensor] = None,
-        memory_key_padding_mask: Optional[torch.BoolTensor] = None,
-        tgt_key_padding_mask: Optional[torch.BoolTensor] = None,
-    ) -> torch.Tensor:
-        output = tgt
-        for layer in self.layers:
-            output = layer(
-                output,
-                memory,
-                tgt_mask=tgt_mask,
-                memory_key_padding_mask=memory_key_padding_mask,
-                tgt_key_padding_mask=tgt_key_padding_mask,
-            )
-        output = self.norm(output)
-        return output
+"""
+Transformer Decoder 模块 — Person B 负责实现
+
+包含:
+- TransformerDecoderLayer: 单层解码器
+- TransformerDecoder: 多层解码器堆叠
+
+架构 (Pre-LayerNorm):
+    x → LN → Masked Self-Attention → Residual
+      → LN → Cross-Attention → Residual
+      → LN → FFN → Residual
+"""
+
+from __future__ import annotations
+
+import copy
+
+import torch
+import torch.nn as nn
+from typing import Optional
+
+from easytranslate.model.attention import MultiHeadAttention, FlashMultiHeadAttention
+
+
+class TransformerDecoderLayer(nn.Module):
+    """
+    单层 Transformer Decoder。
+
+    架构 (Pre-LayerNorm):
+        x → LN → Masked Self-Attention → Residual
+          → LN → Cross-Attention → Residual
+          → LN → FFN → Residual
+    """
+
+    def __init__(
+        self,
+        d_model: int = 512,
+        nhead: int = 8,
+        dim_feedforward: int = 2048,
+        dropout: float = 0.1,
+        activation: str = "gelu",
+        use_flash_attention: bool = True,
+        use_rotary_embedding: bool = True,
+        pre_norm: bool = True,
+    ):
+        super().__init__()
+        self.pre_norm = pre_norm
+        self.d_model = d_model
+
+        attn_cls = FlashMultiHeadAttention if use_flash_attention else MultiHeadAttention
+
+        # 1. Masked Self-Attention
+        self.self_attn = attn_cls(
+            d_model=d_model,
+            nhead=nhead,
+            dropout=dropout,
+            use_rotary_embedding=use_rotary_embedding,
+        )
+
+        # 2. Cross-Attention (decoder queries encoder memory)
+        self.multihead_attn = attn_cls(
+            d_model=d_model,
+            nhead=nhead,
+            dropout=dropout,
+            use_rotary_embedding=False,  # cross-attention 不使用 RoPE
+        )
+
+        # 3. Feed-Forward Network
+        self.linear1 = nn.Linear(d_model, dim_feedforward)
+        self.activation = nn.GELU() if activation == "gelu" else nn.ReLU()
+        self.dropout = nn.Dropout(p=dropout)
+        self.linear2 = nn.Linear(dim_feedforward, d_model)
+
+        # 4. LayerNorms
+        self.norm1 = nn.LayerNorm(d_model)
+        self.norm2 = nn.LayerNorm(d_model)
+        self.norm3 = nn.LayerNorm(d_model)
+
+        # 5. Dropouts for residuals
+        self.dropout1 = nn.Dropout(p=dropout)
+        self.dropout2 = nn.Dropout(p=dropout)
+        self.dropout3 = nn.Dropout(p=dropout)
+
+    def forward(
+        self,
+        tgt: torch.Tensor,                               # [B, T, D]
+        memory: torch.Tensor,                             # [B, S, D] (encoder output)
+        tgt_mask: Optional[torch.Tensor] = None,          # [T, T] causal mask
+        memory_key_padding_mask: Optional[torch.BoolTensor] = None,  # [B, S]
+        tgt_key_padding_mask: Optional[torch.BoolTensor] = None,     # [B, T]
+    ) -> torch.Tensor:
+        """Pre-LayerNorm 前向传播。"""
+        # 适配 Flash Attention: 使用 is_causal 替代显式 causal mask
+        is_causal = tgt_mask is not None
+
+        if self.pre_norm:
+            # 1. Masked Self-Attention
+            residual = tgt
+            tgt = self.norm1(tgt)
+            tgt = self.self_attn(
+                tgt, tgt, tgt,
+                key_padding_mask=tgt_key_padding_mask,
+                is_causal=is_causal,
+            )
+            tgt = residual + self.dropout1(tgt)
+
+            # 2. Cross-Attention
+            residual = tgt
+            tgt = self.norm2(tgt)
+            tgt = self.multihead_attn(
+                tgt, memory, memory,
+                key_padding_mask=memory_key_padding_mask,
+            )
+            tgt = residual + self.dropout2(tgt)
+
+            # 3. FFN
+            residual = tgt
+            tgt = self.norm3(tgt)
+            tgt = self.linear2(self.dropout(self.activation(self.linear1(tgt))))
+            tgt = residual + self.dropout3(tgt)
+        else:
+            # Post-LayerNorm (备用)
+            residual = tgt
+            tgt = self.self_attn(
+                tgt, tgt, tgt,
+                key_padding_mask=tgt_key_padding_mask,
+                is_causal=is_causal,
+            )
+            tgt = self.norm1(residual + self.dropout1(tgt))
+
+            residual = tgt
+            tgt = self.multihead_attn(
+                tgt, memory, memory,
+                key_padding_mask=memory_key_padding_mask,
+            )
+            tgt = self.norm2(residual + self.dropout2(tgt))
+
+            residual = tgt
+            tgt = self.linear2(self.dropout(self.activation(self.linear1(tgt))))
+            tgt = self.norm3(residual + self.dropout3(tgt))
+
+        return tgt
+
+
+class TransformerDecoder(nn.Module):
+    """
+    多层 Transformer Decoder。
+    """
+
+    def __init__(self, decoder_layer: TransformerDecoderLayer, num_layers: int):
+        super().__init__()
+        self.layers = nn.ModuleList(
+            [copy.deepcopy(decoder_layer) for _ in range(num_layers)]
+        )
+        self.num_layers = num_layers
+        self.norm = nn.LayerNorm(decoder_layer.d_model)
+
+    def forward(
+        self,
+        tgt: torch.Tensor,
+        memory: torch.Tensor,
+        tgt_mask: Optional[torch.Tensor] = None,
+        memory_key_padding_mask: Optional[torch.BoolTensor] = None,
+        tgt_key_padding_mask: Optional[torch.BoolTensor] = None,
+    ) -> torch.Tensor:
+        output = tgt
+        for layer in self.layers:
+            output = layer(
+                output,
+                memory,
+                tgt_mask=tgt_mask,
+                memory_key_padding_mask=memory_key_padding_mask,
+                tgt_key_padding_mask=tgt_key_padding_mask,
+            )
+        output = self.norm(output)
+        return output

src/easytranslate/model/encoder.py

@@ -1,127 +1,127 @@
-"""
-Transformer Encoder 模块 — Person B 负责实现
-
-包含:
-- TransformerEncoderLayer: 单层编码器
-- TransformerEncoder: 多层编码器堆叠
-
-架构 (Pre-LayerNorm):
-    x → LayerNorm → MultiHeadAttention → Residual → LayerNorm → FFN → Residual
-"""
-
-from __future__ import annotations
-
-import copy
-
-import torch
-import torch.nn as nn
-from typing import Optional
-
-from easytranslate.model.attention import MultiHeadAttention, FlashMultiHeadAttention
-
-
-class TransformerEncoderLayer(nn.Module):
-    """
-    单层 Transformer Encoder。
-
-    架构 (Pre-LayerNorm):
-        x → LayerNorm → MultiHeadAttention → Residual → LayerNorm → FFN → Residual
-    """
-
-    def __init__(
-        self,
-        d_model: int = 512,
-        nhead: int = 8,
-        dim_feedforward: int = 2048,
-        dropout: float = 0.1,
-        activation: str = "gelu",
-        use_flash_attention: bool = True,
-        use_rotary_embedding: bool = True,
-        pre_norm: bool = True,
-    ):
-        super().__init__()
-        self.pre_norm = pre_norm
-
-        # Self-Attention
-        attn_cls = FlashMultiHeadAttention if use_flash_attention else MultiHeadAttention
-        self.self_attn = attn_cls(
-            d_model=d_model,
-            nhead=nhead,
-            dropout=dropout,
-            use_rotary_embedding=use_rotary_embedding,
-        )
-
-        # Feed-Forward Network
-        self.linear1 = nn.Linear(d_model, dim_feedforward)
-        self.activation = nn.GELU() if activation == "gelu" else nn.ReLU()
-        self.dropout = nn.Dropout(p=dropout)
-        self.linear2 = nn.Linear(dim_feedforward, d_model)
-
-        # LayerNorm
-        self.norm1 = nn.LayerNorm(d_model)
-        self.norm2 = nn.LayerNorm(d_model)
-
-        # Dropout for residual
-        self.dropout1 = nn.Dropout(p=dropout)
-        self.dropout2 = nn.Dropout(p=dropout)
-
-    def forward(
-        self,
-        src: torch.Tensor,                             # [B, S, D]
-        src_key_padding_mask: Optional[torch.BoolTensor] = None,  # [B, S]
-    ) -> torch.Tensor:
-        """Pre-LayerNorm 前向传播。"""
-        if self.pre_norm:
-            # 1. Self-Attention sublayer
-            residual = src
-            src = self.norm1(src)
-            src = self.self_attn(
-                src, src, src,
-                key_padding_mask=src_key_padding_mask,
-            )
-            src = residual + self.dropout1(src)
-
-            # 2. FFN sublayer
-            residual = src
-            src = self.norm2(src)
-            src = self.linear2(self.dropout(self.activation(self.linear1(src))))
-            src = residual + self.dropout2(src)
-        else:
-            # Post-LayerNorm (备用)
-            residual = src
-            src = self.self_attn(
-                src, src, src,
-                key_padding_mask=src_key_padding_mask,
-            )
-            src = self.norm1(residual + self.dropout1(src))
-
-            residual = src
-            src = self.linear2(self.dropout(self.activation(self.linear1(src))))
-            src = self.norm2(residual + self.dropout2(src))
-
-        return src
-
-
-class TransformerEncoder(nn.Module):
-    """
-    多层 Transformer Encoder。
-    """
-
-    def __init__(self, encoder_layer: TransformerEncoderLayer, num_layers: int):
-        super().__init__()
-        self.layers = nn.ModuleList(
-            [copy.deepcopy(encoder_layer) for _ in range(num_layers)]
-        )
-        self.num_layers = num_layers
-        self.norm = nn.LayerNorm(encoder_layer.self_attn.d_model)
-
-    def forward(
-        self,
-        src: torch.Tensor,
-        src_key_padding_mask: Optional[torch.BoolTensor] = None,
-    ) -> torch.Tensor:
-        output = src
-        for layer in self.layers:
-            output = layer(output, src_key_padding_mask=src_key_padding_mask)
-        output = self.norm(output)
-        return output
+"""
+Transformer Encoder 模块 — Person B 负责实现
+
+包含:
+- TransformerEncoderLayer: 单层编码器
+- TransformerEncoder: 多层编码器堆叠
+
+架构 (Pre-LayerNorm):
+    x → LayerNorm → MultiHeadAttention → Residual → LayerNorm → FFN → Residual
+"""
+
+from __future__ import annotations
+
+import copy
+
+import torch
+import torch.nn as nn
+from typing import Optional
+
+from easytranslate.model.attention import MultiHeadAttention, FlashMultiHeadAttention
+
+
+class TransformerEncoderLayer(nn.Module):
+    """
+    单层 Transformer Encoder。
+
+    架构 (Pre-LayerNorm):
+        x → LayerNorm → MultiHeadAttention → Residual → LayerNorm → FFN → Residual
+    """
+
+    def __init__(
+        self,
+        d_model: int = 512,
+        nhead: int = 8,
+        dim_feedforward: int = 2048,
+        dropout: float = 0.1,
+        activation: str = "gelu",
+        use_flash_attention: bool = True,
+        use_rotary_embedding: bool = True,
+        pre_norm: bool = True,
+    ):
+        super().__init__()
+        self.pre_norm = pre_norm
+
+        # Self-Attention
+        attn_cls = FlashMultiHeadAttention if use_flash_attention else MultiHeadAttention
+        self.self_attn = attn_cls(
+            d_model=d_model,
+            nhead=nhead,
+            dropout=dropout,
+            use_rotary_embedding=use_rotary_embedding,
+        )
+
+        # Feed-Forward Network
+        self.linear1 = nn.Linear(d_model, dim_feedforward)
+        self.activation = nn.GELU() if activation == "gelu" else nn.ReLU()
+        self.dropout = nn.Dropout(p=dropout)
+        self.linear2 = nn.Linear(dim_feedforward, d_model)
+
+        # LayerNorm
+        self.norm1 = nn.LayerNorm(d_model)
+        self.norm2 = nn.LayerNorm(d_model)
+
+        # Dropout for residual
+        self.dropout1 = nn.Dropout(p=dropout)
+        self.dropout2 = nn.Dropout(p=dropout)
+
+    def forward(
+        self,
+        src: torch.Tensor,                             # [B, S, D]
+        src_key_padding_mask: Optional[torch.BoolTensor] = None,  # [B, S]
+    ) -> torch.Tensor:
+        """Pre-LayerNorm 前向传播。"""
+        if self.pre_norm:
+            # 1. Self-Attention sublayer
+            residual = src
+            src = self.norm1(src)
+            src = self.self_attn(
+                src, src, src,
+                key_padding_mask=src_key_padding_mask,
+            )
+            src = residual + self.dropout1(src)
+
+            # 2. FFN sublayer
+            residual = src
+            src = self.norm2(src)
+            src = self.linear2(self.dropout(self.activation(self.linear1(src))))
+            src = residual + self.dropout2(src)
+        else:
+            # Post-LayerNorm (备用)
+            residual = src
+            src = self.self_attn(
+                src, src, src,
+                key_padding_mask=src_key_padding_mask,
+            )
+            src = self.norm1(residual + self.dropout1(src))
+
+            residual = src
+            src = self.linear2(self.dropout(self.activation(self.linear1(src))))
+            src = self.norm2(residual + self.dropout2(src))
+
+        return src
+
+
+class TransformerEncoder(nn.Module):
+    """
+    多层 Transformer Encoder。
+    """
+
+    def __init__(self, encoder_layer: TransformerEncoderLayer, num_layers: int):
+        super().__init__()
+        self.layers = nn.ModuleList(
+            [copy.deepcopy(encoder_layer) for _ in range(num_layers)]
+        )
+        self.num_layers = num_layers
+        self.norm = nn.LayerNorm(encoder_layer.self_attn.d_model)
+
+    def forward(
+        self,
+        src: torch.Tensor,
+        src_key_padding_mask: Optional[torch.BoolTensor] = None,
+    ) -> torch.Tensor:
+        output = src
+        for layer in self.layers:
+            output = layer(output, src_key_padding_mask=src_key_padding_mask)
+        output = self.norm(output)
+        return output

src/easytranslate/model/finetune.py

@@ -1,116 +1,116 @@
-"""
-预训练模型微调模块 — Person B 负责实现
-
-功能要求:
-1. load_pretrained_model: 加载 NLLB / mBART 等预训练翻译模型
-2. setup_lora: 配置 LoRA 参数高效微调
-
-技术要点:
-- 使用 HuggingFace transformers 加载预训练模型
-- 使用 PEFT (Parameter-Efficient Fine-Tuning) 库配置 LoRA
-- 支持 NLLB-200 (Meta, 200种语言) 和 mBART-50 (Meta, 50种语言)
-- 冻结预训练参数，只训练 LoRA adapter
-
-这是实验对比的关键部分:
-- 从头训练 vs 预训练微调
-- Full fine-tuning vs LoRA
-"""
-
-from __future__ import annotations
-
-import logging
-from typing import Optional
-
-import torch
-import torch.nn as nn
-
-logger = logging.getLogger(__name__)
-
-
-def load_pretrained_model(
-    model_name: str = "facebook/nllb-200-distilled-600M",
-    src_lang: str = "eng_Latn",
-    tgt_lang: str = "zho_Hans",
-    device: str = "cuda",
-) -> tuple:
-    """
-    加载预训练翻译模型和分词器。
-
-    TODO [Person B]: 实现以下逻辑:
-    1. 使用 AutoModelForSeq2SeqLM.from_pretrained(model_name) 加载模型
-    2. 使用 AutoTokenizer.from_pretrained(model_name) 加载分词器
-    3. 设置源语言和目标语言
-    4. 将模型移到指定设备
-
-    支持的模型:
-    - facebook/nllb-200-distilled-600M (推荐，轻量级)
-    - facebook/nllb-200-1.3B (中等规模)
-    - facebook/mbart-large-50-many-to-many-mmt
-
-    Returns:
-        (model, tokenizer)
-    """
-    from transformers import AutoModelForSeq2SeqLM, AutoTokenizer
-
-    tokenizer = AutoTokenizer.from_pretrained(model_name)
-    model = AutoModelForSeq2SeqLM.from_pretrained(model_name)
-
-    # 设置语言对
-    if hasattr(tokenizer, "lang_code_to_id"):
-        tokenizer.src_lang = src_lang
-        tokenizer.tgt_lang = tgt_lang
-    if hasattr(model.config, "forced_bos_token_id"):
-        model.config.forced_bos_token_id = tokenizer.lang_code_to_id.get(
-            tgt_lang, tokenizer.bos_token_id
-        )
-
-    model = model.to(device)
-    logger.info(f"Loaded pretrained model: {model_name}")
-    return model, tokenizer
-
-
-def setup_lora(
-    model: nn.Module,
-    r: int = 16,
-    alpha: int = 32,
-    dropout: float = 0.05,
-    target_modules: Optional[list[str]] = None,
-) -> nn.Module:
-    """
-    为模型配置 LoRA 微调。
-
-    Returns:
-        peft_model: LoRA 包装后的模型
-    """
-    from peft import LoraConfig, get_peft_model, TaskType
-
-    if target_modules is None:
-        target_modules = ["q_proj", "v_proj"]
-
-    config = LoraConfig(
-        r=r,
-        lora_alpha=alpha,
-        lora_dropout=dropout,
-        target_modules=target_modules,
-        task_type=TaskType.SEQ_2_SEQ_LM,
-    )
-
-    model = get_peft_model(model, config)
-
-    # 打印可训练参数信息
-    trainable_params = sum(p.numel() for p in model.parameters() if p.requires_grad)
-    total_params = sum(p.numel() for p in model.parameters())
-    ratio = 100 * trainable_params / total_params if total_params > 0 else 0
-    logger.info(
-        f"LoRA setup: trainable params={trainable_params:,} "
-        f"/ total={total_params:,} ({ratio:.4f}%)"
-    )
-
-    return model
-
-
-def freeze_model_except_lora(model: nn.Module):
-    """冻结模型所有参数，只保留 LoRA 参数可训练。"""
-    for name, param in model.named_parameters():
-        if "lora" not in name.lower():
-            param.requires_grad = False
+"""
+预训练模型微调模块 — Person B 负责实现
+
+功能要求:
+1. load_pretrained_model: 加载 NLLB / mBART 等预训练翻译模型
+2. setup_lora: 配置 LoRA 参数高效微调
+
+技术要点:
+- 使用 HuggingFace transformers 加载预训练模型
+- 使用 PEFT (Parameter-Efficient Fine-Tuning) 库配置 LoRA
+- 支持 NLLB-200 (Meta, 200种语言) 和 mBART-50 (Meta, 50种语言)
+- 冻结预训练参数，只训练 LoRA adapter
+
+这是实验对比的关键部分:
+- 从头训练 vs 预训练微调
+- Full fine-tuning vs LoRA
+"""
+
+from __future__ import annotations
+
+import logging
+from typing import Optional
+
+import torch
+import torch.nn as nn
+
+logger = logging.getLogger(__name__)
+
+
+def load_pretrained_model(
+    model_name: str = "facebook/nllb-200-distilled-600M",
+    src_lang: str = "eng_Latn",
+    tgt_lang: str = "zho_Hans",
+    device: str = "cuda",
+) -> tuple:
+    """
+    加载预训练翻译模型和分词器。
+
+    TODO [Person B]: 实现以下逻辑:
+    1. 使用 AutoModelForSeq2SeqLM.from_pretrained(model_name) 加载模型
+    2. 使用 AutoTokenizer.from_pretrained(model_name) 加载分词器
+    3. 设置源语言和目标语言
+    4. 将模型移到指定设备
+
+    支持的模型:
+    - facebook/nllb-200-distilled-600M (推荐，轻量级)
+    - facebook/nllb-200-1.3B (中等规模)
+    - facebook/mbart-large-50-many-to-many-mmt
+
+    Returns:
+        (model, tokenizer)
+    """
+    from transformers import AutoModelForSeq2SeqLM, AutoTokenizer
+
+    tokenizer = AutoTokenizer.from_pretrained(model_name)
+    model = AutoModelForSeq2SeqLM.from_pretrained(model_name)
+
+    # 设置语言对
+    if hasattr(tokenizer, "lang_code_to_id"):
+        tokenizer.src_lang = src_lang
+        tokenizer.tgt_lang = tgt_lang
+    if hasattr(model.config, "forced_bos_token_id"):
+        model.config.forced_bos_token_id = tokenizer.lang_code_to_id.get(
+            tgt_lang, tokenizer.bos_token_id
+        )
+
+    model = model.to(device)
+    logger.info(f"Loaded pretrained model: {model_name}")
+    return model, tokenizer
+
+
+def setup_lora(
+    model: nn.Module,
+    r: int = 16,
+    alpha: int = 32,
+    dropout: float = 0.05,
+    target_modules: Optional[list[str]] = None,
+) -> nn.Module:
+    """
+    为模型配置 LoRA 微调。
+
+    Returns:
+        peft_model: LoRA 包装后的模型
+    """
+    from peft import LoraConfig, get_peft_model, TaskType
+
+    if target_modules is None:
+        target_modules = ["q_proj", "v_proj"]
+
+    config = LoraConfig(
+        r=r,
+        lora_alpha=alpha,
+        lora_dropout=dropout,
+        target_modules=target_modules,
+        task_type=TaskType.SEQ_2_SEQ_LM,
+    )
+
+    model = get_peft_model(model, config)
+
+    # 打印可训练参数信息
+    trainable_params = sum(p.numel() for p in model.parameters() if p.requires_grad)
+    total_params = sum(p.numel() for p in model.parameters())
+    ratio = 100 * trainable_params / total_params if total_params > 0 else 0
+    logger.info(
+        f"LoRA setup: trainable params={trainable_params:,} "
+        f"/ total={total_params:,} ({ratio:.4f}%)"
+    )
+
+    return model
+
+
+def freeze_model_except_lora(model: nn.Module):
+    """冻结模型所有参数，只保留 LoRA 参数可训练。"""
+    for name, param in model.named_parameters():
+        if "lora" not in name.lower():
+            param.requires_grad = False

src/easytranslate/model/positional.py

@@ -1,134 +1,134 @@
-"""
-位置编码模块 — Person B 负责实现
-
-包含:
-1. SinusoidalPositionalEncoding: 经典正弦余弦位置编码
-2. RotaryPositionalEmbedding (RoPE): 旋转位置编码
-
-RoPE 是当前最前沿的位置编码方案，被 LLaMA, GPT-NeoX 等大模型广泛采用。
-优势: 更好地捕获相对位置信息，支持长度外推。
-"""
-
-from __future__ import annotations
-
-import math
-
-import torch
-import torch.nn as nn
-
-
-class SinusoidalPositionalEncoding(nn.Module):
-    """
-    经典正弦余弦位置编码 (Vaswani et al., 2017)。
-
-    PE(pos, 2i)   = sin(pos / 10000^(2i/d_model))
-    PE(pos, 2i+1) = cos(pos / 10000^(2i/d_model))
-
-    TODO [Person B]: 实现以下内容:
-    1. 在 __init__ 中预计算位置编码矩阵 [max_seq_len, d_model]
-    2. 注册为 buffer (不参与梯度更新)
-    3. forward(x): 返回 x + pe[:, :x.size(1), :]
-    """
-
-    def __init__(self, d_model: int, max_seq_len: int = 5000, dropout: float = 0.1):
-        super().__init__()
-        self.dropout = nn.Dropout(p=dropout)
-
-        pe = torch.zeros(max_seq_len, d_model)
-        position = torch.arange(0, max_seq_len, dtype=torch.float).unsqueeze(1)
-        div_term = torch.exp(
-            torch.arange(0, d_model, 2).float() * (-math.log(10000.0) / d_model)
-        )
-        pe[:, 0::2] = torch.sin(position * div_term)
-        pe[:, 1::2] = torch.cos(position * div_term)
-        pe = pe.unsqueeze(0)  # [1, max_seq_len, d_model]
-        self.register_buffer("pe", pe)
-
-    def forward(self, x: torch.Tensor) -> torch.Tensor:
-        """
-        Args:
-            x: [B, L, D]
-        Returns:
-            x + positional_encoding: [B, L, D]
-        """
-        x = x + self.pe[:, : x.size(1), :]
-        return self.dropout(x)
-
-
-class RotaryPositionalEmbedding(nn.Module):
-    """
-    旋转位置编码 (RoPE) — Su et al., 2021
-
-    核心思想: 通过旋转变换将位置信息编码到 Q, K 向量中。
-    q' = q * cos(θ) + rotate_half(q) * sin(θ)
-    k' = k * cos(θ) + rotate_half(k) * sin(θ)
-
-    TODO [Person B]: 实现以下内容:
-
-    __init__:
-    1. 计算频率: inv_freq = 1.0 / (base ** (torch.arange(0, dim, 2) / dim))
-    2. 注册为 buffer
-
-    _compute_rope(seq_len):
-    1. 计算 position indices: t = [0, 1, ..., seq_len-1]
-    2. 计算 freqs = torch.outer(t, inv_freq)
-    3. 构建 cos_cached, sin_cached
-
-    apply_rotary_pos_emb(q, k):
-    1. 对 q 和 k 应用旋转变换
-    2. 返回旋转后的 q', k'
-
-    参考: https://arxiv.org/abs/2104.09864
-    """
-
-    def __init__(self, dim: int, max_seq_len: int = 2048, base: float = 10000.0):
-        super().__init__()
-        self.dim = dim
-        self.max_seq_len = max_seq_len
-        self.base = base
-
-        inv_freq = 1.0 / (
-            base ** (torch.arange(0, dim, 2).float() / dim)
-        )
-        self.register_buffer("inv_freq", inv_freq)
-
-        # 预缓存 cos/sin
-        self._cached_seq_len = 0
-        self._cached_cos: torch.Tensor | None = None
-        self._cached_sin: torch.Tensor | None = None
-
-    def _compute_rope(self, seq_len: int, device: torch.device):
-        if seq_len > self._cached_seq_len or self._cached_cos is None:
-            t = torch.arange(seq_len, device=device, dtype=self.inv_freq.dtype)
-            freqs = torch.outer(t, self.inv_freq)  # [seq_len, dim//2]
-            emb = torch.cat((freqs, freqs), dim=-1)  # [seq_len, dim]
-            self._cached_cos = emb.cos()[None, None, :, :]  # [1, 1, seq_len, dim]
-            self._cached_sin = emb.sin()[None, None, :, :]  # [1, 1, seq_len, dim]
-            self._cached_seq_len = seq_len
-        return self._cached_cos, self._cached_sin
-
-    @staticmethod
-    def _rotate_half(x: torch.Tensor) -> torch.Tensor:
-        """将 x 的后半部分取反并与前半部分交换。"""
-        x1, x2 = x.chunk(2, dim=-1)
-        return torch.cat((-x2, x1), dim=-1)
-
-    def apply_rotary_pos_emb(
-        self,
-        q: torch.Tensor,  # [B, nhead, L, d_k]
-        k: torch.Tensor,  # [B, nhead, L, d_k]
-    ) -> tuple[torch.Tensor, torch.Tensor]:
-        """
-        对 Q, K 应用 RoPE。
-
-        Returns:
-            (q_rotated, k_rotated)
-        """
-        cos, sin = self._compute_rope(q.size(2), q.device)
-        q_embed = (q * cos[:, :, : q.size(2), :]) + (
-            self._rotate_half(q) * sin[:, :, : q.size(2), :]
-        )
-        k_embed = (k * cos[:, :, : k.size(2), :]) + (
-            self._rotate_half(k) * sin[:, :, : k.size(2), :]
-        )
-        return q_embed, k_embed
+"""
+位置编码模块 — Person B 负责实现
+
+包含:
+1. SinusoidalPositionalEncoding: 经典正弦余弦位置编码
+2. RotaryPositionalEmbedding (RoPE): 旋转位置编码
+
+RoPE 是当前最前沿的位置编码方案，被 LLaMA, GPT-NeoX 等大模型广泛采用。
+优势: 更好地捕获相对位置信息，支持长度外推。
+"""
+
+from __future__ import annotations
+
+import math
+
+import torch
+import torch.nn as nn
+
+
+class SinusoidalPositionalEncoding(nn.Module):
+    """
+    经典正弦余弦位置编码 (Vaswani et al., 2017)。
+
+    PE(pos, 2i)   = sin(pos / 10000^(2i/d_model))
+    PE(pos, 2i+1) = cos(pos / 10000^(2i/d_model))
+
+    TODO [Person B]: 实现以下内容:
+    1. 在 __init__ 中预计算位置编码矩阵 [max_seq_len, d_model]
+    2. 注册为 buffer (不参与梯度更新)
+    3. forward(x): 返回 x + pe[:, :x.size(1), :]
+    """
+
+    def __init__(self, d_model: int, max_seq_len: int = 5000, dropout: float = 0.1):
+        super().__init__()
+        self.dropout = nn.Dropout(p=dropout)
+
+        pe = torch.zeros(max_seq_len, d_model)
+        position = torch.arange(0, max_seq_len, dtype=torch.float).unsqueeze(1)
+        div_term = torch.exp(
+            torch.arange(0, d_model, 2).float() * (-math.log(10000.0) / d_model)
+        )
+        pe[:, 0::2] = torch.sin(position * div_term)
+        pe[:, 1::2] = torch.cos(position * div_term)
+        pe = pe.unsqueeze(0)  # [1, max_seq_len, d_model]
+        self.register_buffer("pe", pe)
+
+    def forward(self, x: torch.Tensor) -> torch.Tensor:
+        """
+        Args:
+            x: [B, L, D]
+        Returns:
+            x + positional_encoding: [B, L, D]
+        """
+        x = x + self.pe[:, : x.size(1), :]
+        return self.dropout(x)
+
+
+class RotaryPositionalEmbedding(nn.Module):
+    """
+    旋转位置编码 (RoPE) — Su et al., 2021
+
+    核心思想: 通过旋转变换将位置信息编码到 Q, K 向量中。
+    q' = q * cos(θ) + rotate_half(q) * sin(θ)
+    k' = k * cos(θ) + rotate_half(k) * sin(θ)
+
+    TODO [Person B]: 实现以下内容:
+
+    __init__:
+    1. 计算频率: inv_freq = 1.0 / (base ** (torch.arange(0, dim, 2) / dim))
+    2. 注册为 buffer
+
+    _compute_rope(seq_len):
+    1. 计算 position indices: t = [0, 1, ..., seq_len-1]
+    2. 计算 freqs = torch.outer(t, inv_freq)
+    3. 构建 cos_cached, sin_cached
+
+    apply_rotary_pos_emb(q, k):
+    1. 对 q 和 k 应用旋转变换
+    2. 返回旋转后的 q', k'
+
+    参考: https://arxiv.org/abs/2104.09864
+    """
+
+    def __init__(self, dim: int, max_seq_len: int = 2048, base: float = 10000.0):
+        super().__init__()
+        self.dim = dim
+        self.max_seq_len = max_seq_len
+        self.base = base
+
+        inv_freq = 1.0 / (
+            base ** (torch.arange(0, dim, 2).float() / dim)
+        )
+        self.register_buffer("inv_freq", inv_freq)
+
+        # 预缓存 cos/sin
+        self._cached_seq_len = 0
+        self._cached_cos: torch.Tensor | None = None
+        self._cached_sin: torch.Tensor | None = None
+
+    def _compute_rope(self, seq_len: int, device: torch.device):
+        if seq_len > self._cached_seq_len or self._cached_cos is None:
+            t = torch.arange(seq_len, device=device, dtype=self.inv_freq.dtype)
+            freqs = torch.outer(t, self.inv_freq)  # [seq_len, dim//2]
+            emb = torch.cat((freqs, freqs), dim=-1)  # [seq_len, dim]
+            self._cached_cos = emb.cos()[None, None, :, :]  # [1, 1, seq_len, dim]
+            self._cached_sin = emb.sin()[None, None, :, :]  # [1, 1, seq_len, dim]
+            self._cached_seq_len = seq_len
+        return self._cached_cos, self._cached_sin
+
+    @staticmethod
+    def _rotate_half(x: torch.Tensor) -> torch.Tensor:
+        """将 x 的后半部分取反并与前半部分交换。"""
+        x1, x2 = x.chunk(2, dim=-1)
+        return torch.cat((-x2, x1), dim=-1)
+
+    def apply_rotary_pos_emb(
+        self,
+        q: torch.Tensor,  # [B, nhead, L, d_k]
+        k: torch.Tensor,  # [B, nhead, L, d_k]
+    ) -> tuple[torch.Tensor, torch.Tensor]:
+        """
+        对 Q, K 应用 RoPE。
+
+        Returns:
+            (q_rotated, k_rotated)
+        """
+        cos, sin = self._compute_rope(q.size(2), q.device)
+        q_embed = (q * cos[:, :, : q.size(2), :]) + (
+            self._rotate_half(q) * sin[:, :, : q.size(2), :]
+        )
+        k_embed = (k * cos[:, :, : k.size(2), :]) + (
+            self._rotate_half(k) * sin[:, :, : k.size(2), :]
+        )
+        return q_embed, k_embed

src/easytranslate/model/transformer.py

@@ -1,256 +1,256 @@
-"""
-Transformer 翻译模型主体 — Person B 负责实现
-
-这是整个模型的核心文件，定义 Encoder-Decoder 架构。
-
-架构设计:
-- Pre-LayerNorm Transformer (训练更稳定)
-- 可选 Flash Attention 2 (加速注意力计算)
-- 可选 RoPE 旋转位置编码 (替代传统正弦位置编码)
-- 共享 Embedding 权重 (可选)
-"""
-
-from __future__ import annotations
-
-import logging
-import math
-from typing import Optional
-
-import torch
-import torch.nn as nn
-import torch.nn.functional as F
-
-from easytranslate.model.encoder import TransformerEncoder, TransformerEncoderLayer
-from easytranslate.model.decoder import TransformerDecoder, TransformerDecoderLayer
-from easytranslate.model.positional import SinusoidalPositionalEncoding, RotaryPositionalEmbedding
-
-logger = logging.getLogger(__name__)
-
-
-class TransformerTranslationModel(nn.Module):
-    """
-    完整的 Transformer 英中翻译模型。
-
-    架构:
-        Source Embedding + Positional Encoding
-        → Transformer Encoder (N layers)
-        → Transformer Decoder (N layers)
-        → Linear Projection → Softmax
-
-    TODO [Person B]: 实现以下内容:
-
-    __init__:
-    1. 源语言 Embedding: nn.Embedding(src_vocab_size, d_model)
-    2. 目标语言 Embedding: nn.Embedding(tgt_vocab_size, d_model)
-    3. 位置编码: SinusoidalPositionalEncoding 或 RotaryPositionalEmbedding
-    4. Transformer Encoder
-    5. Transformer Decoder
-    6. 输出投影层: nn.Linear(d_model, tgt_vocab_size)
-    7. (可选) 共享 target embedding 和 output projection 的权重
-    8. 调用 _init_weights() 初始化参数
-
-    forward:
-    1. 源语言 embedding + 位置编码 → encoder_input
-    2. 目标语言 embedding + 位置编码 → decoder_input
-    3. 生成 masks (src_key_padding_mask, tgt_key_padding_mask, tgt_mask)
-    4. encoder_output = encoder(encoder_input, src_key_padding_mask)
-    5. decoder_output = decoder(decoder_input, encoder_output, masks...)
-    6. logits = output_projection(decoder_output)
-    7. 返回 logits [B, T, tgt_vocab_size]
-    """
-
-    def __init__(
-        self,
-        src_vocab_size: int,
-        tgt_vocab_size: int,
-        d_model: int = 512,
-        nhead: int = 8,
-        num_encoder_layers: int = 6,
-        num_decoder_layers: int = 6,
-        dim_feedforward: int = 2048,
-        dropout: float = 0.1,
-        activation: str = "gelu",
-        max_seq_len: int = 512,
-        use_flash_attention: bool = True,
-        use_rotary_embedding: bool = True,
-        pre_norm: bool = True,
-        pad_id: int = 0,
-        share_embedding: bool = False,
-    ):
-        super().__init__()
-        self.d_model = d_model
-        self.pad_id = pad_id
-        self.use_rotary_embedding = use_rotary_embedding
-        self.max_seq_len = max_seq_len
-
-        # Embeddings
-        self.src_embed = nn.Embedding(src_vocab_size, d_model)
-        self.tgt_embed = nn.Embedding(tgt_vocab_size, d_model)
-        self.embed_scale = math.sqrt(d_model)
-
-        # Positional encoding
-        if use_rotary_embedding:
-            # RoPE 在 attention 内部应用到 Q/K，不需要额外的位置编码层
-            self.pos_encoding: Optional[nn.Module] = None
-            rope = RotaryPositionalEmbedding(
-                dim=d_model // nhead,
-                max_seq_len=max_seq_len,
-            )
-        else:
-            self.pos_encoding = SinusoidalPositionalEncoding(
-                d_model=d_model,
-                max_seq_len=max_seq_len,
-                dropout=dropout,
-            )
-            rope = None
-
-        # Encoder
-        encoder_layer = TransformerEncoderLayer(
-            d_model=d_model,
-            nhead=nhead,
-            dim_feedforward=dim_feedforward,
-            dropout=dropout,
-            activation=activation,
-            use_flash_attention=use_flash_attention,
-            use_rotary_embedding=use_rotary_embedding,
-            pre_norm=pre_norm,
-        )
-        self.encoder = TransformerEncoder(encoder_layer, num_encoder_layers)
-
-        # Decoder
-        decoder_layer = TransformerDecoderLayer(
-            d_model=d_model,
-            nhead=nhead,
-            dim_feedforward=dim_feedforward,
-            dropout=dropout,
-            activation=activation,
-            use_flash_attention=use_flash_attention,
-            use_rotary_embedding=use_rotary_embedding,
-            pre_norm=pre_norm,
-        )
-        self.decoder = TransformerDecoder(decoder_layer, num_decoder_layers)
-
-        # 将 RoPE 注入到 encoder/decoder 的 attention 模块中
-        if rope is not None:
-            for layer in self.encoder.layers:
-                layer.self_attn.rope = rope
-            for layer in self.decoder.layers:
-                layer.self_attn.rope = rope
-                # cross-attention 不使用 RoPE
-                layer.multihead_attn.rope = None
-
-        # Output projection
-        self.output_projection = nn.Linear(d_model, tgt_vocab_size)
-
-        # 可选: 共享目标语言 embedding 和输出投影权重
-        self.share_embedding = share_embedding
-        if share_embedding:
-            self.output_projection.weight = self.tgt_embed.weight
-
-        self._init_weights()
-
-    def _init_weights(self):
-        """参数初始化。"""
-        for p in self.parameters():
-            if p.dim() > 1:
-                nn.init.xavier_uniform_(p)
-        for module in self.modules():
-            if isinstance(module, nn.Embedding):
-                nn.init.normal_(module.weight, mean=0, std=self.d_model ** -0.5)
-            elif isinstance(module, nn.LayerNorm):
-                nn.init.ones_(module.weight)
-                nn.init.zeros_(module.bias)
-
-    def _generate_square_subsequent_mask(self, sz: int, device: torch.device) -> torch.Tensor:
-        """生成因果注意力掩码 (causal mask)。"""
-        mask = torch.triu(torch.ones(sz, sz, device=device), diagonal=1)
-        mask = mask.masked_fill(mask == 1, float("-inf"))
-        return mask
-
-    def forward(
-        self,
-        src_ids: torch.Tensor,           # [B, S]
-        tgt_input_ids: torch.Tensor,     # [B, T]
-        src_padding_mask: Optional[torch.BoolTensor] = None,  # [B, S]
-        tgt_padding_mask: Optional[torch.BoolTensor] = None,  # [B, T]
-    ) -> torch.Tensor:
-        """
-        前向传播。
-
-        Returns:
-            logits: [B, T, tgt_vocab_size]
-        """
-        # 1. Embedding
-        src_emb = self.src_embed(src_ids) * self.embed_scale  # [B, S, D]
-        tgt_emb = self.tgt_embed(tgt_input_ids) * self.embed_scale  # [B, T, D]
-
-        # 2. Positional encoding (如果不使用 RoPE)
-        if self.pos_encoding is not None:
-            src_emb = self.pos_encoding(src_emb)
-            tgt_emb = self.pos_encoding(tgt_emb)
-
-        # 3. Masks
-        if src_padding_mask is None:
-            src_padding_mask = src_ids.eq(self.pad_id)
-        if tgt_padding_mask is None:
-            tgt_padding_mask = tgt_input_ids.eq(self.pad_id)
-
-        tgt_seq_len = tgt_input_ids.size(1)
-        tgt_mask = self._generate_square_subsequent_mask(tgt_seq_len, tgt_input_ids.device)
-
-        # 4. Encoder
-        encoder_output = self.encoder(src_emb, src_key_padding_mask=src_padding_mask)
-
-        # 5. Decoder
-        decoder_output = self.decoder(
-            tgt_emb,
-            encoder_output,
-            tgt_mask=tgt_mask,
-            memory_key_padding_mask=src_padding_mask,
-            tgt_key_padding_mask=tgt_padding_mask,
-        )
-
-        # 6. Output projection
-        logits = self.output_projection(decoder_output)
-        return logits
-
-    @torch.no_grad()
-    def encode(self, src_ids: torch.Tensor, src_padding_mask: Optional[torch.BoolTensor] = None) -> torch.Tensor:
-        """仅编码（用于推理时复用 encoder 输出）。"""
-        src_emb = self.src_embed(src_ids) * self.embed_scale
-        if self.pos_encoding is not None:
-            src_emb = self.pos_encoding(src_emb)
-        if src_padding_mask is None:
-            src_padding_mask = src_ids.eq(self.pad_id)
-        encoder_output = self.encoder(src_emb, src_key_padding_mask=src_padding_mask)
-        return encoder_output
-
-    @torch.no_grad()
-    def decode_step(
-        self,
-        tgt_input_ids: torch.Tensor,
-        encoder_output: torch.Tensor,
-        src_padding_mask: Optional[torch.BoolTensor] = None,
-    ) -> torch.Tensor:
-        """解码一步（用于自回归推理）。"""
-        tgt_emb = self.tgt_embed(tgt_input_ids) * self.embed_scale
-        if self.pos_encoding is not None:
-            tgt_emb = self.pos_encoding(tgt_emb)
-
-        tgt_seq_len = tgt_input_ids.size(1)
-        tgt_mask = self._generate_square_subsequent_mask(tgt_seq_len, tgt_input_ids.device)
-
-        decoder_output = self.decoder(
-            tgt_emb,
-            encoder_output,
-            tgt_mask=tgt_mask,
-            memory_key_padding_mask=src_padding_mask,
-        )
-
-        # 取最后一个 token 的 logits
-        logits = self.output_projection(decoder_output[:, -1, :])
-        return logits
-
-    def count_parameters(self) -> int:
-        """返回可训练参数数量。"""
-        return sum(p.numel() for p in self.parameters() if p.requires_grad)
+"""
+Transformer 翻译模型主体 — Person B 负责实现
+
+这是整个模型的核心文件，定义 Encoder-Decoder 架构。
+
+架构设计:
+- Pre-LayerNorm Transformer (训练更稳定)
+- 可选 Flash Attention 2 (加速注意力计算)
+- 可选 RoPE 旋转位置编码 (替代传统正弦位置编码)
+- 共享 Embedding 权重 (可选)
+"""
+
+from __future__ import annotations
+
+import logging
+import math
+from typing import Optional
+
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+
+from easytranslate.model.encoder import TransformerEncoder, TransformerEncoderLayer
+from easytranslate.model.decoder import TransformerDecoder, TransformerDecoderLayer
+from easytranslate.model.positional import SinusoidalPositionalEncoding, RotaryPositionalEmbedding
+
+logger = logging.getLogger(__name__)
+
+
+class TransformerTranslationModel(nn.Module):
+    """
+    完整的 Transformer 英中翻译模型。
+
+    架构:
+        Source Embedding + Positional Encoding
+        → Transformer Encoder (N layers)
+        → Transformer Decoder (N layers)
+        → Linear Projection → Softmax
+
+    TODO [Person B]: 实现以下内容:
+
+    __init__:
+    1. 源语言 Embedding: nn.Embedding(src_vocab_size, d_model)
+    2. 目标语言 Embedding: nn.Embedding(tgt_vocab_size, d_model)
+    3. 位置编码: SinusoidalPositionalEncoding 或 RotaryPositionalEmbedding
+    4. Transformer Encoder
+    5. Transformer Decoder
+    6. 输出投影层: nn.Linear(d_model, tgt_vocab_size)
+    7. (可选) 共享 target embedding 和 output projection 的权重
+    8. 调用 _init_weights() 初始化参数
+
+    forward:
+    1. 源语言 embedding + 位置编码 → encoder_input
+    2. 目标语言 embedding + 位置编码 → decoder_input
+    3. 生成 masks (src_key_padding_mask, tgt_key_padding_mask, tgt_mask)
+    4. encoder_output = encoder(encoder_input, src_key_padding_mask)
+    5. decoder_output = decoder(decoder_input, encoder_output, masks...)
+    6. logits = output_projection(decoder_output)
+    7. 返回 logits [B, T, tgt_vocab_size]
+    """
+
+    def __init__(
+        self,
+        src_vocab_size: int,
+        tgt_vocab_size: int,
+        d_model: int = 512,
+        nhead: int = 8,
+        num_encoder_layers: int = 6,
+        num_decoder_layers: int = 6,
+        dim_feedforward: int = 2048,
+        dropout: float = 0.1,
+        activation: str = "gelu",
+        max_seq_len: int = 512,
+        use_flash_attention: bool = True,
+        use_rotary_embedding: bool = True,
+        pre_norm: bool = True,
+        pad_id: int = 0,
+        share_embedding: bool = False,
+    ):
+        super().__init__()
+        self.d_model = d_model
+        self.pad_id = pad_id
+        self.use_rotary_embedding = use_rotary_embedding
+        self.max_seq_len = max_seq_len
+
+        # Embeddings
+        self.src_embed = nn.Embedding(src_vocab_size, d_model)
+        self.tgt_embed = nn.Embedding(tgt_vocab_size, d_model)
+        self.embed_scale = math.sqrt(d_model)
+
+        # Positional encoding
+        if use_rotary_embedding:
+            # RoPE 在 attention 内部应用到 Q/K，不需要额外的位置编码层
+            self.pos_encoding: Optional[nn.Module] = None
+            rope = RotaryPositionalEmbedding(
+                dim=d_model // nhead,
+                max_seq_len=max_seq_len,
+            )
+        else:
+            self.pos_encoding = SinusoidalPositionalEncoding(
+                d_model=d_model,
+                max_seq_len=max_seq_len,
+                dropout=dropout,
+            )
+            rope = None
+
+        # Encoder
+        encoder_layer = TransformerEncoderLayer(
+            d_model=d_model,
+            nhead=nhead,
+            dim_feedforward=dim_feedforward,
+            dropout=dropout,
+            activation=activation,
+            use_flash_attention=use_flash_attention,
+            use_rotary_embedding=use_rotary_embedding,
+            pre_norm=pre_norm,
+        )
+        self.encoder = TransformerEncoder(encoder_layer, num_encoder_layers)
+
+        # Decoder
+        decoder_layer = TransformerDecoderLayer(
+            d_model=d_model,
+            nhead=nhead,
+            dim_feedforward=dim_feedforward,
+            dropout=dropout,
+            activation=activation,
+            use_flash_attention=use_flash_attention,
+            use_rotary_embedding=use_rotary_embedding,
+            pre_norm=pre_norm,
+        )
+        self.decoder = TransformerDecoder(decoder_layer, num_decoder_layers)
+
+        # 将 RoPE 注入到 encoder/decoder 的 attention 模块中
+        if rope is not None:
+            for layer in self.encoder.layers:
+                layer.self_attn.rope = rope
+            for layer in self.decoder.layers:
+                layer.self_attn.rope = rope
+                # cross-attention 不使用 RoPE
+                layer.multihead_attn.rope = None
+
+        # Output projection
+        self.output_projection = nn.Linear(d_model, tgt_vocab_size)
+
+        # 可选: 共享目标语言 embedding 和输出投影权重
+        self.share_embedding = share_embedding
+        if share_embedding:
+            self.output_projection.weight = self.tgt_embed.weight
+
+        self._init_weights()
+
+    def _init_weights(self):
+        """参数初始化。"""
+        for p in self.parameters():
+            if p.dim() > 1:
+                nn.init.xavier_uniform_(p)
+        for module in self.modules():
+            if isinstance(module, nn.Embedding):
+                nn.init.normal_(module.weight, mean=0, std=self.d_model ** -0.5)
+            elif isinstance(module, nn.LayerNorm):
+                nn.init.ones_(module.weight)
+                nn.init.zeros_(module.bias)
+
+    def _generate_square_subsequent_mask(self, sz: int, device: torch.device) -> torch.Tensor:
+        """生成因果注意力掩码 (causal mask)。"""
+        mask = torch.triu(torch.ones(sz, sz, device=device), diagonal=1)
+        mask = mask.masked_fill(mask == 1, float("-inf"))
+        return mask
+
+    def forward(
+        self,
+        src_ids: torch.Tensor,           # [B, S]
+        tgt_input_ids: torch.Tensor,     # [B, T]
+        src_padding_mask: Optional[torch.BoolTensor] = None,  # [B, S]
+        tgt_padding_mask: Optional[torch.BoolTensor] = None,  # [B, T]
+    ) -> torch.Tensor:
+        """
+        前向传播。
+
+        Returns:
+            logits: [B, T, tgt_vocab_size]
+        """
+        # 1. Embedding
+        src_emb = self.src_embed(src_ids) * self.embed_scale  # [B, S, D]
+        tgt_emb = self.tgt_embed(tgt_input_ids) * self.embed_scale  # [B, T, D]
+
+        # 2. Positional encoding (如果不使用 RoPE)
+        if self.pos_encoding is not None:
+            src_emb = self.pos_encoding(src_emb)
+            tgt_emb = self.pos_encoding(tgt_emb)
+
+        # 3. Masks
+        if src_padding_mask is None:
+            src_padding_mask = src_ids.eq(self.pad_id)
+        if tgt_padding_mask is None:
+            tgt_padding_mask = tgt_input_ids.eq(self.pad_id)
+
+        tgt_seq_len = tgt_input_ids.size(1)
+        tgt_mask = self._generate_square_subsequent_mask(tgt_seq_len, tgt_input_ids.device)
+
+        # 4. Encoder
+        encoder_output = self.encoder(src_emb, src_key_padding_mask=src_padding_mask)
+
+        # 5. Decoder
+        decoder_output = self.decoder(
+            tgt_emb,
+            encoder_output,
+            tgt_mask=tgt_mask,
+            memory_key_padding_mask=src_padding_mask,
+            tgt_key_padding_mask=tgt_padding_mask,
+        )
+
+        # 6. Output projection
+        logits = self.output_projection(decoder_output)
+        return logits
+
+    @torch.no_grad()
+    def encode(self, src_ids: torch.Tensor, src_padding_mask: Optional[torch.BoolTensor] = None) -> torch.Tensor:
+        """仅编码（用于推理时复用 encoder 输出）。"""
+        src_emb = self.src_embed(src_ids) * self.embed_scale
+        if self.pos_encoding is not None:
+            src_emb = self.pos_encoding(src_emb)
+        if src_padding_mask is None:
+            src_padding_mask = src_ids.eq(self.pad_id)
+        encoder_output = self.encoder(src_emb, src_key_padding_mask=src_padding_mask)
+        return encoder_output
+
+    @torch.no_grad()
+    def decode_step(
+        self,
+        tgt_input_ids: torch.Tensor,
+        encoder_output: torch.Tensor,
+        src_padding_mask: Optional[torch.BoolTensor] = None,
+    ) -> torch.Tensor:
+        """解码一步（用于自回归推理）。"""
+        tgt_emb = self.tgt_embed(tgt_input_ids) * self.embed_scale
+        if self.pos_encoding is not None:
+            tgt_emb = self.pos_encoding(tgt_emb)
+
+        tgt_seq_len = tgt_input_ids.size(1)
+        tgt_mask = self._generate_square_subsequent_mask(tgt_seq_len, tgt_input_ids.device)
+
+        decoder_output = self.decoder(
+            tgt_emb,
+            encoder_output,
+            tgt_mask=tgt_mask,
+            memory_key_padding_mask=src_padding_mask,
+        )
+
+        # 取最后一个 token 的 logits
+        logits = self.output_projection(decoder_output[:, -1, :])
+        return logits
+
+    def count_parameters(self) -> int:
+        """返回可训练参数数量。"""
+        return sum(p.numel() for p in self.parameters() if p.requires_grad)

src/easytranslate/training/__init__.py

@@ -1,12 +1,12 @@
-"""训练模块 (Person C 负责)"""
-
-from easytranslate.training.trainer import Trainer
-from easytranslate.training.optimizer import build_optimizer, build_scheduler
-from easytranslate.training.loss import LabelSmoothedCrossEntropyLoss
-
-__all__ = [
-    "Trainer",
-    "build_optimizer",
-    "build_scheduler",
-    "LabelSmoothedCrossEntropyLoss",
-]
+"""训练模块 (Person C 负责)"""
+
+from easytranslate.training.trainer import Trainer
+from easytranslate.training.optimizer import build_optimizer, build_scheduler
+from easytranslate.training.loss import LabelSmoothedCrossEntropyLoss
+
+__all__ = [
+    "Trainer",
+    "build_optimizer",
+    "build_scheduler",
+    "LabelSmoothedCrossEntropyLoss",
+]

src/easytranslate/training/loss.py

@@ -1,57 +1,57 @@
-"""
-损失函数模块 — Person C 负责实现
-
-功能要求:
-1. LabelSmoothedCrossEntropyLoss: 带标签平滑的交叉熵损失
-
-技术要点:
-- 标签平滑 (Label Smoothing) 是 Transformer 训练的标准技巧
-- smoothing=0.1 意味着将 10% 的概率均匀分配给非目标类
-- 需要忽略 padding token 的损失
-"""
-
-from __future__ import annotations
-
-import torch
-import torch.nn as nn
-import torch.nn.functional as F
-
-
-class LabelSmoothedCrossEntropyLoss(nn.Module):
-    """
-    带标签平滑的交叉熵损失。
-
-    TODO [Person C]: 实现以下逻辑:
-
-    __init__:
-    1. 保存 smoothing 系数和 pad_id
-
-    forward(logits, targets):
-    1. logits: [B, T, V] — 模型输出
-    2. targets: [B, T] — 目标 token ids
-    3. 展平为 [B*T, V] 和 [B*T]
-    4. 创建 smooth label distribution:
-       - target token 概率 = 1 - smoothing
-       - 其他 token 概率 = smoothing / (V - 1)
-    5. 计算 KL 散度作为损失
-    6. 忽略 pad_id 位置的损失
-    7. 返回平均损失
-
-    参考: "Rethinking the Inception Architecture for Computer Vision" (Szegedy et al.)
-    """
-
-    def __init__(self, smoothing: float = 0.1, pad_id: int = 0):
-        super().__init__()
-        self.smoothing = smoothing
-        self.pad_id = pad_id
-        raise NotImplementedError("TODO: Person C 实现 LabelSmoothedCrossEntropyLoss.__init__")
-
-    def forward(self, logits: torch.Tensor, targets: torch.Tensor) -> torch.Tensor:
-        """
-        Args:
-            logits: [B, T, V]
-            targets: [B, T]
-        Returns:
-            loss: scalar
-        """
-        raise NotImplementedError("TODO: Person C 实现 LabelSmoothedCrossEntropyLoss.forward")
+"""
+损失函数模块 — Person C 负责实现
+
+功能要求:
+1. LabelSmoothedCrossEntropyLoss: 带标签平滑的交叉熵损失
+
+技术要点:
+- 标签平滑 (Label Smoothing) 是 Transformer 训练的标准技巧
+- smoothing=0.1 意味着将 10% 的概率均匀分配给非目标类
+- 需要忽略 padding token 的损失
+"""
+
+from __future__ import annotations
+
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+
+
+class LabelSmoothedCrossEntropyLoss(nn.Module):
+    """
+    带标签平滑的交叉熵损失。
+
+    TODO [Person C]: 实现以下逻辑:
+
+    __init__:
+    1. 保存 smoothing 系数和 pad_id
+
+    forward(logits, targets):
+    1. logits: [B, T, V] — 模型输出
+    2. targets: [B, T] — 目标 token ids
+    3. 展平为 [B*T, V] 和 [B*T]
+    4. 创建 smooth label distribution:
+       - target token 概率 = 1 - smoothing
+       - 其他 token 概率 = smoothing / (V - 1)
+    5. 计算 KL 散度作为损失
+    6. 忽略 pad_id 位置的损失
+    7. 返回平均损失
+
+    参考: "Rethinking the Inception Architecture for Computer Vision" (Szegedy et al.)
+    """
+
+    def __init__(self, smoothing: float = 0.1, pad_id: int = 0):
+        super().__init__()
+        self.smoothing = smoothing
+        self.pad_id = pad_id
+        raise NotImplementedError("TODO: Person C 实现 LabelSmoothedCrossEntropyLoss.__init__")
+
+    def forward(self, logits: torch.Tensor, targets: torch.Tensor) -> torch.Tensor:
+        """
+        Args:
+            logits: [B, T, V]
+            targets: [B, T]
+        Returns:
+            loss: scalar
+        """
+        raise NotImplementedError("TODO: Person C 实现 LabelSmoothedCrossEntropyLoss.forward")

src/easytranslate/training/optimizer.py

@@ -1,76 +1,76 @@
-"""
-优化器与学习率调度器 — Person C 负责实现
-
-功能要求:
-1. build_optimizer: 根据配置创建优化器
-2. build_scheduler: 根据配置创建学习率调度器
-3. InverseSqrtScheduler: 自定义 inverse square root 调度器
-
-技术要点:
-- AdamW 是 Transformer 训练的标准优化器
-- Cosine with warmup 是目前最流行的调度策略
-- Inverse sqrt 是经典 Transformer 论文使用的调度策略
-"""
-
-from __future__ import annotations
-
-import math
-from typing import Optional
-
-import torch
-from torch.optim import Adam, AdamW
-from torch.optim.lr_scheduler import LambdaLR
-
-
-def build_optimizer(model, config: dict) -> torch.optim.Optimizer:
-    """
-    根据配置创建优化器。
-
-    TODO [Person C]: 实现以下逻辑:
-    1. 从 config 中读取 optimizer type, lr, weight_decay, betas, eps
-    2. 根据 type 创建 Adam / AdamW / Adafactor
-    3. (可选) 对不同参数组设置不同学习率:
-       - embedding 层可以用较小的 lr
-       - LayerNorm 的 bias 不加 weight_decay
-    """
-    raise NotImplementedError("TODO: Person C 实现 build_optimizer")
-
-
-def build_scheduler(
-    optimizer: torch.optim.Optimizer,
-    config: dict,
-    num_training_steps: Optional[int] = None,
-) -> torch.optim.lr_scheduler._LRScheduler:
-    """
-    根据配置创建学习率调度器。
-
-    TODO [Person C]: 实现以下逻辑:
-    1. 从 config 中读取 scheduler type, warmup_steps, min_lr
-    2. type == "cosine_with_warmup":
-       使用 get_cosine_schedule_with_warmup (transformers 库)
-    3. type == "inverse_sqrt":
-       实现经典的 lr = d_model^(-0.5) * min(step^(-0.5), step * warmup^(-1.5))
-    4. type == "linear":
-       使用 get_linear_schedule_with_warmup
-
-    参考: Attention Is All You Need, Section 5.3
-    """
-    raise NotImplementedError("TODO: Person C 实现 build_scheduler")
-
-
-class InverseSqrtScheduler(LambdaLR):
-    """
-    Inverse Square Root 学习率调度器。
-
-    lr = base_lr * min(step^{-0.5}, step * warmup_steps^{-1.5})
-
-    这是原始 Transformer 论文使用的调度策略。
-
-    TODO [Person C]:
-    1. 实现 lr_lambda 函数
-    2. warmup 阶段线性增长
-    3. warmup 后按 step^{-0.5} 衰减
-    """
-
-    def __init__(self, optimizer, warmup_steps: int = 4000):
-        raise NotImplementedError("TODO: Person C 实现 InverseSqrtScheduler.__init__")
+"""
+优化器与学习率调度器 — Person C 负责实现
+
+功能要求:
+1. build_optimizer: 根据配置创建优化器
+2. build_scheduler: 根据配置创建学习率调度器
+3. InverseSqrtScheduler: 自定义 inverse square root 调度器
+
+技术要点:
+- AdamW 是 Transformer 训练的标准优化器
+- Cosine with warmup 是目前最流行的调度策略
+- Inverse sqrt 是经典 Transformer 论文使用的调度策略
+"""
+
+from __future__ import annotations
+
+import math
+from typing import Optional
+
+import torch
+from torch.optim import Adam, AdamW
+from torch.optim.lr_scheduler import LambdaLR
+
+
+def build_optimizer(model, config: dict) -> torch.optim.Optimizer:
+    """
+    根据配置创建优化器。
+
+    TODO [Person C]: 实现以下逻辑:
+    1. 从 config 中读取 optimizer type, lr, weight_decay, betas, eps
+    2. 根据 type 创建 Adam / AdamW / Adafactor
+    3. (可选) 对不同参数组设置不同学习率:
+       - embedding 层可以用较小的 lr
+       - LayerNorm 的 bias 不加 weight_decay
+    """
+    raise NotImplementedError("TODO: Person C 实现 build_optimizer")
+
+
+def build_scheduler(
+    optimizer: torch.optim.Optimizer,
+    config: dict,
+    num_training_steps: Optional[int] = None,
+) -> torch.optim.lr_scheduler._LRScheduler:
+    """
+    根据配置创建学习率调度器。
+
+    TODO [Person C]: 实现以下逻辑:
+    1. 从 config 中读取 scheduler type, warmup_steps, min_lr
+    2. type == "cosine_with_warmup":
+       使用 get_cosine_schedule_with_warmup (transformers 库)
+    3. type == "inverse_sqrt":
+       实现经典的 lr = d_model^(-0.5) * min(step^(-0.5), step * warmup^(-1.5))
+    4. type == "linear":
+       使用 get_linear_schedule_with_warmup
+
+    参考: Attention Is All You Need, Section 5.3
+    """
+    raise NotImplementedError("TODO: Person C 实现 build_scheduler")
+
+
+class InverseSqrtScheduler(LambdaLR):
+    """
+    Inverse Square Root 学习率调度器。
+
+    lr = base_lr * min(step^{-0.5}, step * warmup_steps^{-1.5})
+
+    这是原始 Transformer 论文使用的调度策略。
+
+    TODO [Person C]:
+    1. 实现 lr_lambda 函数
+    2. warmup 阶段线性增长
+    3. warmup 后按 step^{-0.5} 衰减
+    """
+
+    def __init__(self, optimizer, warmup_steps: int = 4000):
+        raise NotImplementedError("TODO: Person C 实现 InverseSqrtScheduler.__init__")

src/easytranslate/training/trainer.py

@@ -1,186 +1,186 @@
-"""
-训练器模块 — Person C 负责实现
-
-这是训练的核心控制器，负责:
-1. 训练循环 (training loop)
-2. 验证循环 (validation loop)
-3. 检查点保存/加载
-4. 日志记录
-5. 分布式训练支持
-6. 混合精度训练
-7. 梯度累积
-8. 早停
-"""
-
-from __future__ import annotations
-
-import logging
-import os
-import time
-from pathlib import Path
-from typing import Optional
-
-import torch
-import torch.nn as nn
-from torch.utils.data import DataLoader
-from tqdm import tqdm
-
-logger = logging.getLogger(__name__)
-
-
-class Trainer:
-    """
-    翻译模型训练器。
-
-    TODO [Person C]: 实现以下所有方法。
-
-    使用方法:
-        trainer = Trainer(model, train_loader, val_loader, config)
-        trainer.train()
-    """
-
-    def __init__(
-        self,
-        model: nn.Module,
-        train_loader: DataLoader,
-        val_loader: DataLoader,
-        config: dict,
-        optimizer=None,
-        scheduler=None,
-        criterion=None,
-        evaluator=None,
-    ):
-        """
-        TODO [Person C]: 初始化训练器:
-        1. 保存模型、数据加载器、配置
-        2. 如果 optimizer 为 None，调用 build_optimizer 创建
-        3. 如果 scheduler 为 None，调用 build_scheduler 创建
-        4. 如果 criterion 为 None，创建 LabelSmoothedCrossEntropyLoss
-        5. 设置混合精度: torch.amp.GradScaler (如果配置了 fp16/bf16)
-        6. 设置分布式训练: 根据配置初始化 DDP/FSDP/DeepSpeed
-        7. 初始化日志记录器 (TensorBoard / WandB)
-        8. 初始化训练状态 (epoch, step, best_metric)
-        """
-        self.model = model
-        self.train_loader = train_loader
-        self.val_loader = val_loader
-        self.config = config
-        raise NotImplementedError("TODO: Person C 实现 Trainer.__init__")
-
-    def train(self):
-        """
-        主训练循环。
-
-        TODO [Person C]: 实现以下逻辑:
-        1. for epoch in range(start_epoch, num_epochs):
-        2.   train_loss = self._train_one_epoch(epoch)
-        3.   val_metrics = self._validate(epoch)
-        4.   self._log_metrics(epoch, train_loss, val_metrics)
-        5.   self._save_checkpoint(epoch, val_metrics)
-        6.   if self._should_early_stop(val_metrics): break
-        7. self._log_final_results()
-        """
-        raise NotImplementedError("TODO: Person C 实现 train")
-
-    def _train_one_epoch(self, epoch: int) -> float:
-        """
-        训练一个 epoch。
-
-        TODO [Person C]: 实现以下逻辑:
-        1. model.train()
-        2. 遍历 train_loader:
-           a. 将 batch 移到设备
-           b. 混合精度上下文: with torch.amp.autocast('cuda'):
-           c. logits = model(src_ids, tgt_input_ids, masks...)
-           d. loss = criterion(logits, labels)
-           e. loss = loss / gradient_accumulation_steps
-           f. scaler.scale(loss).backward()
-           g. 每 gradient_accumulation_steps 步:
-              - scaler.unscale_(optimizer)
-              - torch.nn.utils.clip_grad_norm_(model.parameters(), max_norm=1.0)
-              - scaler.step(optimizer)
-              - scaler.update()
-              - scheduler.step()
-              - optimizer.zero_grad()
-           h. 记录 loss, lr 等指标
-        3. 返回平均训练 loss
-        """
-        raise NotImplementedError("TODO: Person C 实现 _train_one_epoch")
-
-    def _validate(self, epoch: int) -> dict:
-        """
-        验证。
-
-        TODO [Person C]: 实现以下逻辑:
-        1. model.eval()
-        2. with torch.no_grad():
-        3.   遍历 val_loader，计算 loss
-        4.   每 N 步调用 evaluator 计算 BLEU 等指标
-        5. 返回 {"val_loss": ..., "bleu": ..., "comet": ...}
-        """
-        raise NotImplementedError("TODO: Person C 实现 _validate")
-
-    def _save_checkpoint(self, epoch: int, metrics: dict):
-        """
-        保存检查点。
-
-        TODO [Person C]: 实现以下逻辑:
-        1. 构建 checkpoint dict:
-           {
-               "epoch": epoch,
-               "step": self.global_step,
-               "model_state_dict": model.state_dict(),
-               "optimizer_state_dict": optimizer.state_dict(),
-               "scheduler_state_dict": scheduler.state_dict(),
-               "metrics": metrics,
-               "config": config,
-           }
-        2. 保存到 checkpoint_dir/checkpoint_epoch_{epoch}.pt
-        3. 如果是最佳模型，额外保存为 best_model.pt
-        4. 清理旧的检查点 (保留最近 max_checkpoints 个)
-        """
-        raise NotImplementedError("TODO: Person C 实现 _save_checkpoint")
-
-    def _load_checkpoint(self, checkpoint_path: str):
-        """
-        加载检查点继续训练。
-
-        TODO [Person C]:
-        1. torch.load(checkpoint_path)
-        2. 恢复 model, optimizer, scheduler 状态
-        3. 恢复 epoch, step 计数器
-        """
-        raise NotImplementedError("TODO: Person C 实现 _load_checkpoint")
-
-    def _should_early_stop(self, metrics: dict) -> bool:
-        """
-        判断是否应该早停。
-
-        TODO [Person C]:
-        1. 比较当前指标与最佳指标
-        2. 如果连续 patience 个 epoch 没有改善，返回 True
-        """
-        raise NotImplementedError("TODO: Person C 实现 _should_early_stop")
-
-    def _log_metrics(self, epoch: int, train_loss: float, val_metrics: dict):
-        """
-        记录训练指标到 TensorBoard / WandB。
-
-        TODO [Person C]:
-        1. 使用 self.logger 记录 train_loss, val_loss, bleu, lr 等
-        2. 打印到控制台 (使用 rich 库的 table 格式)
-        """
-        raise NotImplementedError("TODO: Person C 实现 _log_metrics")
-
-    def _setup_distributed(self):
-        """
-        设置分布式训练。
-
-        TODO [Person C]: 根据 config["training"]["distributed"]["strategy"]:
-        1. "ddp": 使用 torch.nn.parallel.DistributedDataParallel
-        2. "fsdp": 使用 torch.distributed.fsdp.FullyShardedDataParallel
-        3. "deepspeed": 使用 deepspeed.initialize()
-
-        也可以使用 HuggingFace Accelerate 统一处理。
-        """
-        raise NotImplementedError("TODO: Person C 实现 _setup_distributed")
+"""
+训练器模块 — Person C 负责实现
+
+这是训练的核心控制器，负责:
+1. 训练循环 (training loop)
+2. 验证循环 (validation loop)
+3. 检查点保存/加载
+4. 日志记录
+5. 分布式训练支持
+6. 混合精度训练
+7. 梯度累积
+8. 早停
+"""
+
+from __future__ import annotations
+
+import logging
+import os
+import time
+from pathlib import Path
+from typing import Optional
+
+import torch
+import torch.nn as nn
+from torch.utils.data import DataLoader
+from tqdm import tqdm
+
+logger = logging.getLogger(__name__)
+
+
+class Trainer:
+    """
+    翻译模型训练器。
+
+    TODO [Person C]: 实现以下所有方法。
+
+    使用方法:
+        trainer = Trainer(model, train_loader, val_loader, config)
+        trainer.train()
+    """
+
+    def __init__(
+        self,
+        model: nn.Module,
+        train_loader: DataLoader,
+        val_loader: DataLoader,
+        config: dict,
+        optimizer=None,
+        scheduler=None,
+        criterion=None,
+        evaluator=None,
+    ):
+        """
+        TODO [Person C]: 初始化训练器:
+        1. 保存模型、数据加载器、配置
+        2. 如果 optimizer 为 None，调用 build_optimizer 创建
+        3. 如果 scheduler 为 None，调用 build_scheduler 创建
+        4. 如果 criterion 为 None，创建 LabelSmoothedCrossEntropyLoss
+        5. 设置混合精度: torch.amp.GradScaler (如果配置了 fp16/bf16)
+        6. 设置分布式训练: 根据配置初始化 DDP/FSDP/DeepSpeed
+        7. 初始化日志记录器 (TensorBoard / WandB)
+        8. 初始化训练状态 (epoch, step, best_metric)
+        """
+        self.model = model
+        self.train_loader = train_loader
+        self.val_loader = val_loader
+        self.config = config
+        raise NotImplementedError("TODO: Person C 实现 Trainer.__init__")
+
+    def train(self):
+        """
+        主训练循环。
+
+        TODO [Person C]: 实现以下逻辑:
+        1. for epoch in range(start_epoch, num_epochs):
+        2.   train_loss = self._train_one_epoch(epoch)
+        3.   val_metrics = self._validate(epoch)
+        4.   self._log_metrics(epoch, train_loss, val_metrics)
+        5.   self._save_checkpoint(epoch, val_metrics)
+        6.   if self._should_early_stop(val_metrics): break
+        7. self._log_final_results()
+        """
+        raise NotImplementedError("TODO: Person C 实现 train")
+
+    def _train_one_epoch(self, epoch: int) -> float:
+        """
+        训练一个 epoch。
+
+        TODO [Person C]: 实现以下逻辑:
+        1. model.train()
+        2. 遍历 train_loader:
+           a. 将 batch 移到设备
+           b. 混合精度上下文: with torch.amp.autocast('cuda'):
+           c. logits = model(src_ids, tgt_input_ids, masks...)
+           d. loss = criterion(logits, labels)
+           e. loss = loss / gradient_accumulation_steps
+           f. scaler.scale(loss).backward()
+           g. 每 gradient_accumulation_steps 步:
+              - scaler.unscale_(optimizer)
+              - torch.nn.utils.clip_grad_norm_(model.parameters(), max_norm=1.0)
+              - scaler.step(optimizer)
+              - scaler.update()
+              - scheduler.step()
+              - optimizer.zero_grad()
+           h. 记录 loss, lr 等指标
+        3. 返回平均训练 loss
+        """
+        raise NotImplementedError("TODO: Person C 实现 _train_one_epoch")
+
+    def _validate(self, epoch: int) -> dict:
+        """
+        验证。
+
+        TODO [Person C]: 实现以下逻辑:
+        1. model.eval()
+        2. with torch.no_grad():
+        3.   遍历 val_loader，计算 loss
+        4.   每 N 步调用 evaluator 计算 BLEU 等指标
+        5. 返回 {"val_loss": ..., "bleu": ..., "comet": ...}
+        """
+        raise NotImplementedError("TODO: Person C 实现 _validate")
+
+    def _save_checkpoint(self, epoch: int, metrics: dict):
+        """
+        保存检查点。
+
+        TODO [Person C]: 实现以下逻辑:
+        1. 构建 checkpoint dict:
+           {
+               "epoch": epoch,
+               "step": self.global_step,
+               "model_state_dict": model.state_dict(),
+               "optimizer_state_dict": optimizer.state_dict(),
+               "scheduler_state_dict": scheduler.state_dict(),
+               "metrics": metrics,
+               "config": config,
+           }
+        2. 保存到 checkpoint_dir/checkpoint_epoch_{epoch}.pt
+        3. 如果是最佳模型，额外保存为 best_model.pt
+        4. 清理旧的检查点 (保留最近 max_checkpoints 个)
+        """
+        raise NotImplementedError("TODO: Person C 实现 _save_checkpoint")
+
+    def _load_checkpoint(self, checkpoint_path: str):
+        """
+        加载检查点继续训练。
+
+        TODO [Person C]:
+        1. torch.load(checkpoint_path)
+        2. 恢复 model, optimizer, scheduler 状态
+        3. 恢复 epoch, step 计数器
+        """
+        raise NotImplementedError("TODO: Person C 实现 _load_checkpoint")
+
+    def _should_early_stop(self, metrics: dict) -> bool:
+        """
+        判断是否应该早停。
+
+        TODO [Person C]:
+        1. 比较当前指标与最佳指标
+        2. 如果连续 patience 个 epoch 没有改善，返回 True
+        """
+        raise NotImplementedError("TODO: Person C 实现 _should_early_stop")
+
+    def _log_metrics(self, epoch: int, train_loss: float, val_metrics: dict):
+        """
+        记录训练指标到 TensorBoard / WandB。
+
+        TODO [Person C]:
+        1. 使用 self.logger 记录 train_loss, val_loss, bleu, lr 等
+        2. 打印到控制台 (使用 rich 库的 table 格式)
+        """
+        raise NotImplementedError("TODO: Person C 实现 _log_metrics")
+
+    def _setup_distributed(self):
+        """
+        设置分布式训练。
+
+        TODO [Person C]: 根据 config["training"]["distributed"]["strategy"]:
+        1. "ddp": 使用 torch.nn.parallel.DistributedDataParallel
+        2. "fsdp": 使用 torch.distributed.fsdp.FullyShardedDataParallel
+        3. "deepspeed": 使用 deepspeed.initialize()
+
+        也可以使用 HuggingFace Accelerate 统一处理。
+        """
+        raise NotImplementedError("TODO: Person C 实现 _setup_distributed")

src/easytranslate/utils/__init__.py

@@ -1,7 +1,7 @@
-"""工具模块"""
-
-from easytranslate.utils.config import load_config, merge_configs
-from easytranslate.utils.seed import set_seed
-from easytranslate.utils.logging import setup_logging
-
-__all__ = ["load_config", "merge_configs", "set_seed", "setup_logging"]
+"""工具模块"""
+
+from easytranslate.utils.config import load_config, merge_configs
+from easytranslate.utils.seed import set_seed
+from easytranslate.utils.logging import setup_logging
+
+__all__ = ["load_config", "merge_configs", "set_seed", "setup_logging"]

src/easytranslate/utils/config.py

@@ -1,51 +1,51 @@
-"""
-配置管理模块 — 公共模块 (Person E 可以先完成)
-
-功能: 加载 YAML 配置，支持命令行覆盖和多配置合并。
-"""
-
-from __future__ import annotations
-
-import logging
-from pathlib import Path
-from typing import Any, Optional
-
-from omegaconf import OmegaConf, DictConfig
-
-logger = logging.getLogger(__name__)
-
-
-def load_config(config_path: str | Path) -> DictConfig:
-    """
-    加载 YAML 配置文件。
-
-    TODO [Person E]: 实现以下逻辑:
-    1. 使用 OmegaConf.load(config_path) 加载配置
-    2. 验证配置完整性
-    3. 返回 DictConfig 对象
-    """
-    raise NotImplementedError("TODO: Person E 实现 load_config")
-
-
-def merge_configs(base_config: DictConfig, override_config: DictConfig) -> DictConfig:
-    """
-    合并配置 (override 覆盖 base)。
-
-    TODO [Person E]:
-    使用 OmegaConf.merge(base_config, override_config)
-    """
-    raise NotImplementedError("TODO: Person E 实现 merge_configs")
-
-
-def config_from_cli(config_path: str, cli_args: list[str]) -> DictConfig:
-    """
-    从配置文件 + 命令行参数构建最终配置。
-
-    TODO [Person E]:
-    1. 加载 config_path
-    2. 使用 OmegaConf.from_cli(cli_args) 解析命令行参数
-    3. 合并并返回
-
-    使用方式: python train.py --config configs/default.yaml training.lr=1e-4 model.nhead=16
-    """
-    raise NotImplementedError("TODO: Person E 实现 config_from_cli")
+"""
+配置管理模块 — 公共模块 (Person E 可以先完成)
+
+功能: 加载 YAML 配置，支持命令行覆盖和多配置合并。
+"""
+
+from __future__ import annotations
+
+import logging
+from pathlib import Path
+from typing import Any, Optional
+
+from omegaconf import OmegaConf, DictConfig
+
+logger = logging.getLogger(__name__)
+
+
+def load_config(config_path: str | Path) -> DictConfig:
+    """
+    加载 YAML 配置文件。
+
+    TODO [Person E]: 实现以下逻辑:
+    1. 使用 OmegaConf.load(config_path) 加载配置
+    2. 验证配置完整性
+    3. 返回 DictConfig 对象
+    """
+    raise NotImplementedError("TODO: Person E 实现 load_config")
+
+
+def merge_configs(base_config: DictConfig, override_config: DictConfig) -> DictConfig:
+    """
+    合并配置 (override 覆盖 base)。
+
+    TODO [Person E]:
+    使用 OmegaConf.merge(base_config, override_config)
+    """
+    raise NotImplementedError("TODO: Person E 实现 merge_configs")
+
+
+def config_from_cli(config_path: str, cli_args: list[str]) -> DictConfig:
+    """
+    从配置文件 + 命令行参数构建最终配置。
+
+    TODO [Person E]:
+    1. 加载 config_path
+    2. 使用 OmegaConf.from_cli(cli_args) 解析命令行参数
+    3. 合并并返回
+
+    使用方式: python train.py --config configs/default.yaml training.lr=1e-4 model.nhead=16
+    """
+    raise NotImplementedError("TODO: Person E 实现 config_from_cli")

src/easytranslate/utils/logging.py

@@ -1,26 +1,26 @@
-"""
-日志管理模块 — 公共模块
-"""
-
-import logging
-import sys
-from pathlib import Path
-from typing import Optional
-
-
-def setup_logging(
-    log_dir: Optional[str] = None,
-    level: int = logging.INFO,
-    log_file: Optional[str] = None,
-):
-    """
-    配置日志系统。
-
-    TODO [Person E]: 实现以下逻辑:
-    1. 创建 root logger
-    2. 设置 StreamHandler (控制台输出)
-    3. (可选) 设置 FileHandler (文件输出)
-    4. 使用 rich.logging.RichHandler 美化控制台输出
-    5. 设置日志格式: [时间] [级别] [模块名] 消息
-    """
-    raise NotImplementedError("TODO: Person E 实现 setup_logging")
+"""
+日志管理模块 — 公共模块
+"""
+
+import logging
+import sys
+from pathlib import Path
+from typing import Optional
+
+
+def setup_logging(
+    log_dir: Optional[str] = None,
+    level: int = logging.INFO,
+    log_file: Optional[str] = None,
+):
+    """
+    配置日志系统。
+
+    TODO [Person E]: 实现以下逻辑:
+    1. 创建 root logger
+    2. 设置 StreamHandler (控制台输出)
+    3. (可选) 设置 FileHandler (文件输出)
+    4. 使用 rich.logging.RichHandler 美化控制台输出
+    5. 设置日志格式: [时间] [级别] [模块名] 消息
+    """
+    raise NotImplementedError("TODO: Person E 实现 setup_logging")

src/easytranslate/utils/seed.py

@@ -1,23 +1,23 @@
-"""
-随机种子管理模块 — 公共模块
-"""
-
-import random
-
-import numpy as np
-import torch
-
-
-def set_seed(seed: int = 42):
-    """
-    设置全局随机种子以确保实验可复现。
-
-    TODO [Person E]:
-    1. random.seed(seed)
-    2. np.random.seed(seed)
-    3. torch.manual_seed(seed)
-    4. torch.cuda.manual_seed_all(seed)
-    5. torch.backends.cudnn.deterministic = True
-    6. torch.backends.cudnn.benchmark = False
-    """
-    raise NotImplementedError("TODO: Person E 实现 set_seed")
+"""
+随机种子管理模块 — 公共模块
+"""
+
+import random
+
+import numpy as np
+import torch
+
+
+def set_seed(seed: int = 42):
+    """
+    设置全局随机种子以确保实验可复现。
+
+    TODO [Person E]:
+    1. random.seed(seed)
+    2. np.random.seed(seed)
+    3. torch.manual_seed(seed)
+    4. torch.cuda.manual_seed_all(seed)
+    5. torch.backends.cudnn.deterministic = True
+    6. torch.backends.cudnn.benchmark = False
+    """
+    raise NotImplementedError("TODO: Person E 实现 set_seed")

tests/test_data.py

@@ -1,143 +1,143 @@
-"""
-单元测试 — 数据模块 (Person A 实现后需通过)
-"""
-
-import pytest
-import torch
-
-from easytranslate.data import (
-    DynamicBatchSampler,
-    TokenizerWrapper,
-    TranslationCollator,
-    TranslationDataset,
-    clean_text,
-    deduplicate_pairs,
-    filter_by_length,
-    preprocess_pipeline,
-    train_bpe_tokenizer,
-)
-
-
-@pytest.fixture()
-def tiny_tokenizer() -> TokenizerWrapper:
-    texts = [
-        "Hello world",
-        "Machine translation is useful",
-        "I love natural language processing",
-        "你好 世界",
-        "机器翻译 很 有用",
-        "我 喜欢 自然语言处理",
-    ]
-    return train_bpe_tokenizer(texts, vocab_size=80, min_frequency=1)
-
-
-class TestTranslationDataset:
-    """测试 TranslationDataset 类。"""
-
-    def test_dataset_length(self, tiny_tokenizer):
-        """测试数据集长度。"""
-        dataset = TranslationDataset(["Hello world"], ["你好 世界"], tokenizer=tiny_tokenizer)
-        assert len(dataset) == 1
-
-    def test_getitem_returns_correct_keys(self, tiny_tokenizer):
-        """测试 __getitem__ 返回正确的字段。"""
-        dataset = TranslationDataset(["Hello world"], ["你好 世界"], tokenizer=tiny_tokenizer)
-        item = dataset[0]
-        assert set(item) == {"src_ids", "tgt_input_ids", "labels", "src_len", "tgt_len"}
-
-    def test_getitem_tensor_types(self, tiny_tokenizer):
-        """测试返回的 tensor 类型正确。"""
-        dataset = TranslationDataset(["Hello world"], ["你好 世界"], tokenizer=tiny_tokenizer)
-        item = dataset[0]
-        assert item["src_ids"].dtype == torch.long
-        assert item["tgt_input_ids"].dtype == torch.long
-        assert item["labels"].dtype == torch.long
-
-    def test_src_tgt_mismatch_raises(self):
-        """测试源目标数量不匹配时抛出异常。"""
-        with pytest.raises(ValueError):
-            TranslationDataset(["a", "b"], ["甲"], tokenizer=object())
-
-
-class TestTokenizer:
-    """测试分词器。"""
-
-    def test_bpe_train_and_encode(self, tiny_tokenizer):
-        """测试 BPE 训练和编码。"""
-        ids = tiny_tokenizer.encode("Hello world", add_special_tokens=True)
-        assert len(ids) >= 3
-        assert ids[0] == tiny_tokenizer.bos_token_id
-        assert ids[-1] == tiny_tokenizer.eos_token_id
-
-    def test_encode_decode_roundtrip(self, tiny_tokenizer):
-        """测试编码-解码往返一致性。"""
-        ids = tiny_tokenizer.encode("Hello world")
-        decoded = tiny_tokenizer.decode(ids)
-        assert "Hello" in decoded
-        assert "world" in decoded
-
-    def test_special_tokens(self, tiny_tokenizer):
-        """测试特殊 token 正确。"""
-        assert tiny_tokenizer.pad_token_id == 0
-        assert tiny_tokenizer.unk_token_id == 1
-        assert tiny_tokenizer.bos_token_id == 2
-        assert tiny_tokenizer.eos_token_id == 3
-
-
-class TestPreprocessing:
-    """测试预处理。"""
-
-    def test_clean_text_unicode(self):
-        """测试 Unicode 标准化。"""
-        assert clean_text("ＡＢＣ\u200b  １２３") == "ABC 123"
-
-    def test_filter_by_length(self):
-        """测试按长度过滤。"""
-        assert filter_by_length("hello world", "你好世界", max_src_len=10, max_tgt_len=10)
-        assert not filter_by_length("hello " * 300, "你好", max_src_len=256)
-
-    def test_deduplicate(self):
-        """测试去重。"""
-        pairs = deduplicate_pairs([("a", "甲"), ("a", "甲"), ("b", "乙")])
-        assert pairs == [("a", "甲"), ("b", "乙")]
-
-    def test_preprocess_pipeline(self):
-        """测试完整预处理流水线。"""
-        src, tgt = preprocess_pipeline([" Hello  world ", "", "Hello  world"], [" 你好 世界 ", "空", "你好 世界"])
-        assert src == ["Hello world"]
-        assert tgt == ["你好 世界"]
-
-
-class TestCollator:
-    """测试数据整理器。"""
-
-    def test_padding(self, tiny_tokenizer):
-        """测试 padding 正确。"""
-        dataset = TranslationDataset(
-            ["Hello world", "Machine translation is useful"],
-            ["你好 世界", "机器翻译 很 有用"],
-            tokenizer=tiny_tokenizer,
-        )
-        batch = TranslationCollator(pad_token_id=tiny_tokenizer.pad_token_id)([dataset[0], dataset[1]])
-        assert batch["src_ids"].ndim == 2
-        assert batch["tgt_input_ids"].ndim == 2
-        assert batch["labels"].shape == batch["tgt_input_ids"].shape
-
-    def test_attention_mask(self, tiny_tokenizer):
-        """测试 attention mask 正确。"""
-        dataset = TranslationDataset(
-            ["Hello", "Machine translation is useful"],
-            ["你好", "机器翻译 很 有用"],
-            tokenizer=tiny_tokenizer,
-        )
-        batch = TranslationCollator(pad_token_id=tiny_tokenizer.pad_token_id)([dataset[0], dataset[1]])
-        assert batch["src_padding_mask"].dtype == torch.bool
-        assert torch.equal(batch["src_attention_mask"], (~batch["src_padding_mask"]).long())
-
-    def test_dynamic_batch_sampler(self):
-        """测试动态 batch 不超过 token 预算。"""
-        sampler = DynamicBatchSampler([5, 6, 20, 21], max_tokens_per_batch=24, shuffle=False)
-        batches = list(sampler)
-        assert batches
-        for batch in batches:
-            assert max([5, 6, 20, 21][idx] for idx in batch) * len(batch) <= 24
+"""
+单元测试 — 数据模块 (Person A 实现后需通过)
+"""
+
+import pytest
+import torch
+
+from easytranslate.data import (
+    DynamicBatchSampler,
+    TokenizerWrapper,
+    TranslationCollator,
+    TranslationDataset,
+    clean_text,
+    deduplicate_pairs,
+    filter_by_length,
+    preprocess_pipeline,
+    train_bpe_tokenizer,
+)
+
+
+@pytest.fixture()
+def tiny_tokenizer() -> TokenizerWrapper:
+    texts = [
+        "Hello world",
+        "Machine translation is useful",
+        "I love natural language processing",
+        "你好 世界",
+        "机器翻译 很 有用",
+        "我 喜欢 自然语言处理",
+    ]
+    return train_bpe_tokenizer(texts, vocab_size=80, min_frequency=1)
+
+
+class TestTranslationDataset:
+    """测试 TranslationDataset 类。"""
+
+    def test_dataset_length(self, tiny_tokenizer):
+        """测试数据集长度。"""
+        dataset = TranslationDataset(["Hello world"], ["你好 世界"], tokenizer=tiny_tokenizer)
+        assert len(dataset) == 1
+
+    def test_getitem_returns_correct_keys(self, tiny_tokenizer):
+        """测试 __getitem__ 返回正确的字段。"""
+        dataset = TranslationDataset(["Hello world"], ["你好 世界"], tokenizer=tiny_tokenizer)
+        item = dataset[0]
+        assert set(item) == {"src_ids", "tgt_input_ids", "labels", "src_len", "tgt_len"}
+
+    def test_getitem_tensor_types(self, tiny_tokenizer):
+        """测试返回的 tensor 类型正确。"""
+        dataset = TranslationDataset(["Hello world"], ["你好 世界"], tokenizer=tiny_tokenizer)
+        item = dataset[0]
+        assert item["src_ids"].dtype == torch.long
+        assert item["tgt_input_ids"].dtype == torch.long
+        assert item["labels"].dtype == torch.long
+
+    def test_src_tgt_mismatch_raises(self):
+        """测试源目标数量不匹配时抛出异常。"""
+        with pytest.raises(ValueError):
+            TranslationDataset(["a", "b"], ["甲"], tokenizer=object())
+
+
+class TestTokenizer:
+    """测试分词器。"""
+
+    def test_bpe_train_and_encode(self, tiny_tokenizer):
+        """测试 BPE 训练和编码。"""
+        ids = tiny_tokenizer.encode("Hello world", add_special_tokens=True)
+        assert len(ids) >= 3
+        assert ids[0] == tiny_tokenizer.bos_token_id
+        assert ids[-1] == tiny_tokenizer.eos_token_id
+
+    def test_encode_decode_roundtrip(self, tiny_tokenizer):
+        """测试编码-解码往返一致性。"""
+        ids = tiny_tokenizer.encode("Hello world")
+        decoded = tiny_tokenizer.decode(ids)
+        assert "Hello" in decoded
+        assert "world" in decoded
+
+    def test_special_tokens(self, tiny_tokenizer):
+        """测试特殊 token 正确。"""
+        assert tiny_tokenizer.pad_token_id == 0
+        assert tiny_tokenizer.unk_token_id == 1
+        assert tiny_tokenizer.bos_token_id == 2
+        assert tiny_tokenizer.eos_token_id == 3
+
+
+class TestPreprocessing:
+    """测试预处理。"""
+
+    def test_clean_text_unicode(self):
+        """测试 Unicode 标准化。"""
+        assert clean_text("ＡＢＣ\u200b  １２３") == "ABC 123"
+
+    def test_filter_by_length(self):
+        """测试按长度过滤。"""
+        assert filter_by_length("hello world", "你好世界", max_src_len=10, max_tgt_len=10)
+        assert not filter_by_length("hello " * 300, "你好", max_src_len=256)
+
+    def test_deduplicate(self):
+        """测试去重。"""
+        pairs = deduplicate_pairs([("a", "甲"), ("a", "甲"), ("b", "乙")])
+        assert pairs == [("a", "甲"), ("b", "乙")]
+
+    def test_preprocess_pipeline(self):
+        """测试完整预处理流水线。"""
+        src, tgt = preprocess_pipeline([" Hello  world ", "", "Hello  world"], [" 你好 世界 ", "空", "你好 世界"])
+        assert src == ["Hello world"]
+        assert tgt == ["你好 世界"]
+
+
+class TestCollator:
+    """测试数据整理器。"""
+
+    def test_padding(self, tiny_tokenizer):
+        """测试 padding 正确。"""
+        dataset = TranslationDataset(
+            ["Hello world", "Machine translation is useful"],
+            ["你好 世界", "机器翻译 很 有用"],
+            tokenizer=tiny_tokenizer,
+        )
+        batch = TranslationCollator(pad_token_id=tiny_tokenizer.pad_token_id)([dataset[0], dataset[1]])
+        assert batch["src_ids"].ndim == 2
+        assert batch["tgt_input_ids"].ndim == 2
+        assert batch["labels"].shape == batch["tgt_input_ids"].shape
+
+    def test_attention_mask(self, tiny_tokenizer):
+        """测试 attention mask 正确。"""
+        dataset = TranslationDataset(
+            ["Hello", "Machine translation is useful"],
+            ["你好", "机器翻译 很 有用"],
+            tokenizer=tiny_tokenizer,
+        )
+        batch = TranslationCollator(pad_token_id=tiny_tokenizer.pad_token_id)([dataset[0], dataset[1]])
+        assert batch["src_padding_mask"].dtype == torch.bool
+        assert torch.equal(batch["src_attention_mask"], (~batch["src_padding_mask"]).long())
+
+    def test_dynamic_batch_sampler(self):
+        """测试动态 batch 不超过 token 预算。"""
+        sampler = DynamicBatchSampler([5, 6, 20, 21], max_tokens_per_batch=24, shuffle=False)
+        batches = list(sampler)
+        assert batches
+        for batch in batches:
+            assert max([5, 6, 20, 21][idx] for idx in batch) * len(batch) <= 24

tests/test_evaluation.py

@@ -1,54 +1,54 @@
-"""
-单元测试 — 评估模块 (Person D 实现后需通过)
-"""
-
-import pytest
-import torch
-
-
-class TestMetrics:
-    """测试评估指标。"""
-
-    def test_bleu_perfect(self):
-        """完美翻译的 BLEU 应接近 100。"""
-        pass
-
-    def test_bleu_zero(self):
-        """完全不相关的翻译 BLEU 应接近 0。"""
-        pass
-
-    def test_chrf_score(self):
-        """测试 chrF 分数范围合理。"""
-        pass
-
-
-class TestDecoding:
-    """测试解码策略。"""
-
-    def test_greedy_decode_shape(self):
-        """测试贪心解码输出维度。"""
-        pass
-
-    def test_greedy_produces_eos(self):
-        """测试贪心解码能生成 EOS。"""
-        pass
-
-    def test_beam_search_better_than_greedy(self):
-        """束搜索质量应不低于贪心。"""
-        pass
-
-    def test_no_repeat_ngram(self):
-        """测试 n-gram 不重复约束。"""
-        pass
-
-
-class TestEvaluator:
-    """测试评估器。"""
-
-    def test_translate_single(self):
-        """测试单句翻译。"""
-        pass
-
-    def test_translate_batch(self):
-        """测试批量翻译。"""
-        pass
+"""
+单元测试 — 评估模块 (Person D 实现后需通过)
+"""
+
+import pytest
+import torch
+
+
+class TestMetrics:
+    """测试评估指标。"""
+
+    def test_bleu_perfect(self):
+        """完美翻译的 BLEU 应接近 100。"""
+        pass
+
+    def test_bleu_zero(self):
+        """完全不相关的翻译 BLEU 应接近 0。"""
+        pass
+
+    def test_chrf_score(self):
+        """测试 chrF 分数范围合理。"""
+        pass
+
+
+class TestDecoding:
+    """测试解码策略。"""
+
+    def test_greedy_decode_shape(self):
+        """测试贪心解码输出维度。"""
+        pass
+
+    def test_greedy_produces_eos(self):
+        """测试贪心解码能生成 EOS。"""
+        pass
+
+    def test_beam_search_better_than_greedy(self):
+        """束搜索质量应不低于贪心。"""
+        pass
+
+    def test_no_repeat_ngram(self):
+        """测试 n-gram 不重复约束。"""
+        pass
+
+
+class TestEvaluator:
+    """测试评估器。"""
+
+    def test_translate_single(self):
+        """测试单句翻译。"""
+        pass
+
+    def test_translate_batch(self):
+        """测试批量翻译。"""
+        pass

tests/test_model.py

@@ -1,60 +1,60 @@
-"""
-单元测试 — 模型模块 (Person B 实现后需通过)
-"""
-
-import pytest
-import torch
-
-
-class TestMultiHeadAttention:
-    """测试多头注意力。"""
-
-    def test_output_shape(self):
-        """测试输出维度正确。"""
-        # TODO: 创建 MHA，输入 [2, 10, 512]，验证输出 [2, 10, 512]
-        pass
-
-    def test_padding_mask(self):
-        """测试 padding mask 效果。"""
-        pass
-
-
-class TestPositionalEncoding:
-    """测试位置编码。"""
-
-    def test_sinusoidal_shape(self):
-        """测试正弦位置编码输出维度。"""
-        pass
-
-    def test_rope_shape(self):
-        """测试 RoPE 输出维度。"""
-        pass
-
-    def test_rope_relative_distance(self):
-        """测试 RoPE 捕获相对距离。"""
-        pass
-
-
-class TestTransformerModel:
-    """测试完整 Transformer 模型。"""
-
-    def test_forward_shape(self):
-        """测试前向传播输出维度。"""
-        # TODO: 构建小模型，输入 mock 数据，验证 logits shape [B, T, V]
-        pass
-
-    def test_encode_shape(self):
-        """测试 encode 输出维度。"""
-        pass
-
-    def test_decode_step_shape(self):
-        """测试 decode_step 输出维度。"""
-        pass
-
-    def test_parameter_count(self):
-        """测试参数数量计算。"""
-        pass
-
-    def test_causal_mask(self):
-        """测试因果掩码生成。"""
-        pass
+"""
+单元测试 — 模型模块 (Person B 实现后需通过)
+"""
+
+import pytest
+import torch
+
+
+class TestMultiHeadAttention:
+    """测试多头注意力。"""
+
+    def test_output_shape(self):
+        """测试输出维度正确。"""
+        # TODO: 创建 MHA，输入 [2, 10, 512]，验证输出 [2, 10, 512]
+        pass
+
+    def test_padding_mask(self):
+        """测试 padding mask 效果。"""
+        pass
+
+
+class TestPositionalEncoding:
+    """测试位置编码。"""
+
+    def test_sinusoidal_shape(self):
+        """测试正弦位置编码输出维度。"""
+        pass
+
+    def test_rope_shape(self):
+        """测试 RoPE 输出维度。"""
+        pass
+
+    def test_rope_relative_distance(self):
+        """测试 RoPE 捕获相对距离。"""
+        pass
+
+
+class TestTransformerModel:
+    """测试完整 Transformer 模型。"""
+
+    def test_forward_shape(self):
+        """测试前向传播输出维度。"""
+        # TODO: 构建小模型，输入 mock 数据，验证 logits shape [B, T, V]
+        pass
+
+    def test_encode_shape(self):
+        """测试 encode 输出维度。"""
+        pass
+
+    def test_decode_step_shape(self):
+        """测试 decode_step 输出维度。"""
+        pass
+
+    def test_parameter_count(self):
+        """测试参数数量计算。"""
+        pass
+
+    def test_causal_mask(self):
+        """测试因果掩码生成。"""
+        pass

tests/test_training.py

@@ -1,46 +1,46 @@
-"""
-单元测试 — 训练模块 (Person C 实现后需通过)
-"""
-
-import pytest
-import torch
-
-
-class TestLabelSmoothedLoss:
-    """测试标签平滑损失。"""
-
-    def test_no_smoothing_equals_ce(self):
-        """smoothing=0 时应等于标准交叉熵。"""
-        pass
-
-    def test_smoothing_reduces_confidence(self):
-        """标签平滑应降低损失对正确预测的自信度。"""
-        pass
-
-    def test_ignore_padding(self):
-        """应忽略 padding 位置的损失。"""
-        pass
-
-
-class TestOptimizer:
-    """测试优化器构建。"""
-
-    def test_build_adamw(self):
-        """测试构建 AdamW。"""
-        pass
-
-    def test_weight_decay_groups(self):
-        """测试权重衰减参数组正确。"""
-        pass
-
-
-class TestScheduler:
-    """测试学习率调度器。"""
-
-    def test_warmup_phase(self):
-        """测试 warmup 阶段 lr 线性增长。"""
-        pass
-
-    def test_decay_phase(self):
-        """测试衰减阶段 lr 下降。"""
-        pass
+"""
+单元测试 — 训练模块 (Person C 实现后需通过)
+"""
+
+import pytest
+import torch
+
+
+class TestLabelSmoothedLoss:
+    """测试标签平滑损失。"""
+
+    def test_no_smoothing_equals_ce(self):
+        """smoothing=0 时应等于标准交叉熵。"""
+        pass
+
+    def test_smoothing_reduces_confidence(self):
+        """标签平滑应降低损失对正确预测的自信度。"""
+        pass
+
+    def test_ignore_padding(self):
+        """应忽略 padding 位置的损失。"""
+        pass
+
+
+class TestOptimizer:
+    """测试优化器构建。"""
+
+    def test_build_adamw(self):
+        """测试构建 AdamW。"""
+        pass
+
+    def test_weight_decay_groups(self):
+        """测试权重衰减参数组正确。"""
+        pass
+
+
+class TestScheduler:
+    """测试学习率调度器。"""
+
+    def test_warmup_phase(self):
+        """测试 warmup 阶段 lr 线性增长。"""
+        pass
+
+    def test_decay_phase(self):
+        """测试衰减阶段 lr 下降。"""
+        pass