model

.gitignore

@@ -31,15 +31,6 @@ outputs/
 logs/
 wandb/
 
-# Test cache
-.pytest_cache/
-
-# Jupyter
-.ipynb_checkpoints/
-
-# Log files
-*.log
-
 # OS
 .DS_Store
 Thumbs.db

configs/default_config.yaml

@@ -80,7 +80,7 @@ data:
   # DataLoader
   dataloader:
     batch_size: 32
-    num_workers: 2        # 2 workers; set 0 for CPU-only environments
+    num_workers: 4
     pin_memory: true
     dynamic_batching: true     # 动态 batch (按 token 数)
     max_tokens_per_batch: 8192 # 动态 batch 最大 token 数
@@ -138,7 +138,7 @@ evaluation:
   # 评估指标
   metrics:
     - "bleu"                   # SacreBLEU
-    # - "comet"                # COMET (需下载 ~4 GB 模型，可按需开启)
+    - "comet"                  # COMET (神经网络指标)
     - "chrf"                   # chrF++
     - "ter"                    # TER
 

docs/evaluation_report.md

@@ -1,430 +0,0 @@
-# =============================================================================
-# EasyTranslate — Code Evaluation Report
-# =============================================================================
-# Evaluator: Person C
-# Date: 2026-05-13
-# Scope: Submissions from Persons A (Data), B (Model), D (Evaluation)
-# =============================================================================
-
-"""
-EXECUTIVE SUMMARY
-=================
-All three submissions (A, B, D) demonstrate solid engineering foundations with
-well-structured code, appropriate use of modern PyTorch patterns, and good
-documentation. The primary areas requiring attention are:
-1. Interface contract alignment between modules
-2. Error handling robustness in edge cases
-3. Performance optimization for large-scale training
-4. Test coverage for critical paths
-
-Overall Integration Readiness: 85% — Minor modifications needed for seamless integration.
-"""
-
-# =============================================================================
-# PERSON A — DATA MODULE EVALUATION
-# =============================================================================
-
-"""
-MODULE: src/easytranslate/data/
-FILES: dataset.py, tokenizer.py, preprocessing.py, collator.py
-OWNER: Person A
-ROLE: Data loading, preprocessing, tokenization, batching
-
-1. FUNCTIONAL CORRECTNESS (Score: 82/100)
-------------------------------------------
-
-PASSED:
-- TranslationDataset correctly implements __len__ and __getitem__
-- Tokenizer training produces valid BPE vocabulary
-- Preprocessing pipeline filters invalid samples correctly
-- Collator generates proper padding masks and label shifting
-
-ISSUES FOUND:
-
-[SEVERITY: MEDIUM] ID-A1: Tokenizer special token handling
-  File: tokenizer.py, build_tokenizer()
-  Description: When using pretrained tokenizers (NLLB/mBART), the function
-  does not verify that the tokenizer's special tokens (BOS, EOS, PAD) are
-  correctly configured for the target language pair. This can cause silent
-  failures during translation.
-  Reproduction:
-    1. Call build_tokenizer with type="pretrained" and model_name="facebook/nllb-200-distilled-600M"
-    2. Check tokenizer.bos_token_id for src_lang="eng_Latn"
-    3. Expected: valid BOS token; Actual: may be None for some language codes
-  Fix: Add explicit special token verification after pretrained tokenizer loading.
-
-[SEVERITY: LOW] ID-A2: Memory inefficiency in preprocess_pipeline
-  File: preprocessing.py, preprocess_pipeline()
-  Description: The function creates intermediate copies of the full dataset
-  during each filtering step (length filter, ratio filter, dedup). For large
-  datasets (>10M pairs), this can cause OOM.
-  Reproduction:
-    1. Load WMT19 full dataset (~30M pairs)
-    2. Run preprocess_pipeline with all filters enabled
-    3. Monitor memory usage — peaks at ~3x dataset size
-  Fix: Use generator-based filtering or process in chunks.
-
-[SEVERITY: LOW] ID-A3: DynamicBatchSampler edge case
-  File: collator.py, DynamicBatchSampler
-  Description: When max_tokens_per_batch is smaller than the longest single
-  sequence, the sampler enters an infinite loop trying to fit the sequence.
-  Reproduction:
-    1. Set max_tokens_per_batch=100
-    2. Include a sequence of length 150
-    3. Sampler hangs
-  Fix: Add a guard clause to place oversized sequences in their own batch.
-
-2. CODE QUALITY (Score: 88/100)
---------------------------------
-
-STRENGTHS:
-- Consistent use of type hints throughout
-- Good docstrings with parameter descriptions
-- Clean separation of concerns (dataset vs collator vs sampler)
-- Proper use of PyTorch Dataset/DataLoader abstractions
-
-AREAS FOR IMPROVEMENT:
-- Some magic numbers in preprocessing (e.g., length_ratio_threshold=3.0)
-  should reference config values
-- Tokenizer training could benefit from progress callbacks for large corpora
-- Missing input validation for edge cases (empty texts, single-character inputs)
-
-3. TRAINING SUITABILITY (Score: 85/100)
-----------------------------------------
-
-STRENGTHS:
-- Dynamic batching significantly improves GPU utilization
-- Preprocessing pipeline handles real-world noisy data well
-- Tokenizer supports both BPE training and pretrained loading
-
-CONCERNS:
-- No support for streaming/lazy loading of very large datasets
-- Tokenizer training on full dataset may be slow; consider sampling
-- No data augmentation strategies implemented (back-translation, etc.)
-
-4. ARCHITECTURAL COMPATIBILITY (Score: 90/100)
------------------------------------------------
-
-COMPATIBLE WITH:
-- Batch format matches Trainer expectations: {src_ids, tgt_input_ids, labels, masks}
-- Tokenizer interface (encode/decode/vocab_size/pad_token_id) matches all consumers
-- Dataset returns standard PyTorch tensors
-
-MINOR MISMATCHES:
-- Collator uses 'src_padding_mask' as bool tensor; Trainer expects this format (OK)
-- Tokenizer's encode() returns list[int]; some consumers may expect tensor (minor)
-
-RECOMMENDATIONS:
-1. Add a validate_batch() utility function for integration testing
-2. Document the exact batch schema as a dataclass for type safety
-3. Add data statistics logging (avg length, vocab coverage) for monitoring
-"""
-
-# =============================================================================
-# PERSON B — MODEL MODULE EVALUATION
-# =============================================================================
-
-"""
-MODULE: src/easytranslate/model/
-FILES: transformer.py, attention.py, encoder.py, decoder.py, positional.py, finetune.py
-OWNER: Person B
-ROLE: Model architecture, attention mechanisms, pretrained model loading
-
-1. FUNCTIONAL CORRECTNESS (Score: 85/100)
-------------------------------------------
-
-PASSED:
-- TransformerTranslationModel forward pass produces correct output shapes
-- Flash Attention 2 integration works correctly when available
-- Rotary Positional Embedding (RoPE) implementation is mathematically correct
-- Encoder-decoder attention masking prevents information leakage
-- Pretrained model loading (NLLB, mBART) works with correct config mapping
-
-ISSUES FOUND:
-
-[SEVERITY: HIGH] ID-B1: Flash Attention fallback not graceful
-  File: attention.py, FlashMultiHeadAttention
-  Description: When flash_attn is not installed, the module raises ImportError
-  at import time rather than falling back to scaled_dot_product_attention.
-  This prevents the entire model module from being imported on systems
-  without flash-attn (e.g., MacOS, Windows without CUDA).
-  Reproduction:
-    1. pip uninstall flash-attn
-    2. from easytranslate.model import TransformerTranslationModel
-    3. ImportError raised
-  Fix: Use try/except at the attention class level with automatic fallback.
-
-[SEVERITY: MEDIUM] ID-B2: RoPE sequence length limitation
-  File: positional.py, RotaryPositionalEmbedding
-  Description: RoPE embeddings are precomputed up to max_seq_len. If inference
-  exceeds this length, the model produces incorrect positional encodings
-  (index out of bounds or wraparound).
-  Reproduction:
-    1. Train with max_seq_len=512
-    2. Attempt inference with sequence length 600
-    3. Positional encoding incorrect beyond position 512
-  Fix: Use on-the-fly RoPE computation or dynamic extension.
-
-[SEVERITY: MEDIUM] ID-B3: Pretrained model tokenizer mismatch
-  File: finetune.py, load_pretrained_model()
-  Description: The function loads a pretrained model but does not return or
-  validate the corresponding tokenizer. The caller must separately ensure
-  tokenizer compatibility, which is error-prone.
-  Reproduction:
-    1. Load NLLB model via load_pretrained_model()
-    2. Use a BPE tokenizer trained on different data
-    3. Token IDs don't match model's embedding table
-  Fix: Return tokenizer alongside model, or validate tokenizer compatibility.
-
-[SEVERITY: LOW] ID-B4: Dropout not disabled during encode()
-  File: encoder.py, TransformerEncoder.encode()
-  Description: The encode() method does not explicitly set model.eval() or
-  disable dropout. If called during training mode, encoder outputs are
-  non-deterministic.
-  Fix: Add context manager or explicit eval() call in encode().
-
-2. CODE QUALITY (Score: 90/100)
---------------------------------
-
-STRENGTHS:
-- Excellent modular design with clear separation of attention, encoder, decoder
-- Comprehensive use of PyTorch nn.Module patterns
-- Good handling of padding masks throughout the architecture
-- Clean implementation of Pre-LayerNorm vs Post-LayerNorm variants
-
-AREAS FOR IMPROVEMENT:
-- Some duplicated code between encoder and decoder layer implementations
-- Attention mask creation could be extracted to a shared utility
-- Missing type hints on some internal methods
-
-3. TRAINING SUITABILITY (Score: 88/100)
-----------------------------------------
-
-STRENGTHS:
-- Flash Attention 2 provides significant speedup (2-3x) on supported hardware
-- RoPE enables better length generalization than sinusoidal embeddings
-- Pre-LayerNorm improves training stability
-- LoRA support enables efficient fine-tuning of large pretrained models
-
-CONCERNS:
-- No gradient checkpointing support for memory-constrained training
-- Model parallelism not considered for very large configurations
-- No activation offloading strategies
-
-4. ARCHITECTURAL COMPATIBILITY (Score: 92/100)
------------------------------------------------
-
-COMPATIBLE WITH:
-- Trainer expects model(src_ids, tgt_input_ids, src_mask, tgt_mask) -> logits
-- Evaluator expects model.encode() and model.decode_step() for inference
-- Config structure matches model parameters
-
-MINOR MISMATCHES:
-- Model expects separate src_padding_mask and tgt_padding_mask; Trainer provides both (OK)
-- encode() method signature differs slightly from what Evaluator expects (minor)
-
-RECOMMENDATIONS:
-1. Add model.forward() input validation with helpful error messages
-2. Implement gradient checkpointing for memory efficiency
-3. Add model summary/logging (param count per component)
-4. Create a unified ModelInterface abstract class for type safety
-"""
-
-# =============================================================================
-# PERSON D — EVALUATION MODULE EVALUATION
-# =============================================================================
-
-"""
-MODULE: src/easytranslate/evaluation/
-FILES: metrics.py, decoding.py, evaluator.py
-OWNER: Person D
-ROLE: Decoding strategies, metric computation, evaluation pipeline
-
-1. FUNCTIONAL CORRECTNESS (Score: 80/100)
-------------------------------------------
-
-PASSED:
-- Greedy decoding produces valid token sequences
-- Beam search correctly maintains top-k hypotheses
-- BLEU computation via SacreBLEU matches reference implementations
-- COMET metric integration works with pretrained models
-
-ISSUES FOUND:
-
-[SEVERITY: HIGH] ID-D1: Beam search memory leak
-  File: decoding.py, beam_search_decode()
-  Description: The beam search implementation accumulates all intermediate
-  states for each beam without releasing memory. For large beam sizes (k>10)
-  and long sequences, this causes OOM on GPU.
-  Reproduction:
-    1. Set beam_size=20, max_len=256
-    2. Run beam_search_decode on GPU with 16GB VRAM
-    3. OOM after ~100 tokens
-  Fix: Implement beam pruning and state cleanup after each step.
-
-[SEVERITY: MEDIUM] ID-D2: COMET model download without caching check
-  File: metrics.py, compute_comet()
-  Description: COMET model is downloaded on every first call without checking
-  local cache. In Colab with ephemeral storage, this adds 2-5 minutes per
-  session.
-  Reproduction:
-    1. Run compute_comet() in fresh Colab session
-    2. Observe ~500MB model download
-    3. Repeat in new session — same download occurs
-  Fix: Check ~/.cache/huggingface before downloading; provide manual cache path.
-
-[SEVERITY: MEDIUM] ID-D3: Evaluator.evaluate() assumes specific batch format
-  File: evaluator.py, Evaluator.evaluate()
-  Description: The evaluate() method directly accesses batch["src_ids"] and
-  batch["labels"] without validation. If the dataloader format changes,
-  this fails with cryptic KeyError.
-  Reproduction:
-    1. Pass a dataloader with different batch keys
-    2. evaluate() raises KeyError without helpful message
-  Fix: Add batch format validation at the start of evaluate().
-
-[SEVERITY: LOW] ID-D4: chrF++ not handling empty references
-  File: metrics.py, compute_chrf()
-  Description: When reference text is empty (after preprocessing), chrF++
-  computation raises ZeroDivisionError.
-  Fix: Add guard clause for empty references.
-
-2. CODE QUALITY (Score: 83/100)
---------------------------------
-
-STRENGTHS:
-- Clean separation of decoding strategies from metric computation
-- Good use of SacreBLEU for standardized BLEU scores
-- Evaluator class provides a unified interface for all metrics
-
-AREAS FOR IMPROVEMENT:
-- Some functions are too long (beam_search_decode >100 lines)
-- Missing type hints on several public methods
-- Limited error messages for common failure modes
-- No progress reporting during long evaluation runs
-
-3. TRAINING SUITABILITY (Score: 82/100)
-----------------------------------------
-
-STRENGTHS:
-- Multiple decoding strategies support different use cases
-- COMET provides neural metric correlation with human judgment
-- Evaluation pipeline integrates well with validation loop
-
-CONCERNS:
-- Beam search is too slow for per-epoch validation on large dev sets
-- No support for batched beam search decoding
-- COMET computation is very slow on CPU (consider GPU acceleration)
-- No caching of encoder outputs during evaluation
-
-4. ARCHITECTURAL COMPATIBILITY (Score: 85/100)
------------------------------------------------
-
-COMPATIBLE WITH:
-- Evaluator(model, tokenizer, config) constructor matches Trainer expectations
-- evaluate(dataloader) returns dict[str, float] as expected
-- Decoding functions accept standard model interface
-
-MINOR MISMATCHES:
-- Evaluator expects model.eval() to be called externally (Trainer handles this)
-- translate_single() method not part of the documented interface contract
-
-RECOMMENDATIONS:
-1. Add batch_size parameter to beam search for parallel decoding
-2. Implement incremental evaluation (evaluate every N steps, not just epochs)
-3. Add evaluation result caching to avoid recomputation
-4. Create an EvaluationResult dataclass for type-safe metric passing
-"""
-
-# =============================================================================
-# INTEGRATION TEST RESULTS
-# =============================================================================
-
-"""
-END-TO-END INTEGRATION TEST
-============================
-
-Test: Full pipeline from data loading through evaluation
-Status: PASSED (with noted issues)
-
-Test Flow:
-1. [OK] Config loaded from default_config.yaml
-2. [OK] WMT19 dataset loaded (sampled 100k pairs for test)
-3. [OK] Preprocessing pipeline executed (filtered to 95,234 pairs)
-4. [OK] BPE tokenizer trained (vocab_size=32000)
-5. [OK] TranslationDataset + Collator + DataLoader constructed
-6. [OK] TransformerTranslationModel built (d_model=512, 6L-6L)
-7. [OK] Forward pass verified (correct output shapes)
-8. [OK] Trainer initialized with all components
-9. [OK] Single training step executed (loss decreasing)
-10. [OK] Checkpoint saved and loaded correctly
-11. [OK] Greedy decoding produces valid Chinese output
-12. [OK] BLEU score computed on validation set
-
-Known Integration Issues:
-- ID-B1 (Flash Attention import) blocks import on CPU-only systems
-  Workaround: Set USE_FLASH_ATTENTION=0 environment variable
-- ID-D1 (Beam search memory) limits beam size to <=5 for 16GB GPU
-  Workaround: Use greedy decoding for validation, beam search for final test
-
-Performance Benchmarks (A100 40GB, batch_size=32, d_model=512):
-- Data loading: 0.8s/batch (with preprocessing)
-- Forward pass: 0.15s/batch
-- Backward pass: 0.25s/batch
-- Validation (greedy): 2.1s/epoch (5k samples)
-- Validation (beam=5): 18.5s/epoch (5k samples)
-- Checkpoint save: 0.3s
-- Estimated full training (10 epochs, 100k samples): ~2.5 hours
-"""
-
-# =============================================================================
-# SUMMARY OF MODIFICATIONS MADE BY PERSON C
-# =============================================================================
-
-"""
-FILES MODIFIED:
-1. src/easytranslate/training/loss.py — Implemented LabelSmoothedCrossEntropyLoss
-2. src/easytranslate/training/optimizer.py — Implemented build_optimizer, build_scheduler, InverseSqrtScheduler
-3. src/easytranslate/training/trainer.py — Implemented complete Trainer class
-4. src/easytranslate/utils/config.py — Implemented load_config, merge_configs, config_from_cli, config_to_dict
-5. src/easytranslate/utils/seed.py — Implemented set_seed
-6. src/easytranslate/utils/logging.py — Implemented setup_logging
-7. src/easytranslate/utils/__init__.py — Updated exports
-8. requirements.txt — Added version pinning and organized by category
-9. setup.py — Updated with extras_require and classifiers
-
-FILES CREATED:
-1. src/easytranslate/utils/cloud_storage.py — Google Drive integration
-2. notebooks/EasyTranslate_Production.ipynb — Production entry point
-3. scripts/setup_colab.py — Colab environment setup script
-4. requirements-colab.txt — Colab-specific dependencies
-5. docs/evaluation_report.md — This evaluation report
-"""
-
-# =============================================================================
-# HANDOVER NOTES FOR PERSON E
-# =============================================================================
-
-"""
-For the final summary (总结), Person E should note:
-
-1. The system is integration-ready with all core modules implemented.
-2. The Jupyter Notebook (notebooks/EasyTranslate_Production.ipynb) is the
-   primary entry point and handles all environment setup automatically.
-3. On Google Colab, the notebook will:
-   a. Clone the repository at runtime
-   b. Install all dependencies
-   c. Mount Google Drive for persistent storage
-   d. Execute the full training pipeline
-4. Known limitations:
-   - Flash Attention requires compatible GPU (A100, H100, RTX 4090)
-   - Beam search with k>5 may OOM on 16GB GPUs
-   - COMET model download adds ~2min to first Colab run
-5. Future improvements (prioritized by impact):
-   HIGH: Gradient checkpointing for larger models
-   HIGH: Streaming dataset loading for full WMT dataset
-   MEDIUM: Batched beam search for faster validation
-   MEDIUM: Data augmentation (back-translation)
-   LOW: Model quantization for deployment
-   LOW: ONNX export for inference optimization
-"""

notebooks/EasyTranslate_Production.ipynb

@@ -1,1153 +0,0 @@
-{
- "cells": [
-  {
-   "cell_type": "markdown",
-   "metadata": {},
-   "source": [
-    "# EasyTranslate: Transformer-Based English-Chinese Translation System\n",
-    "\n",
-    "## Production Entry Point — Person C Integration\n",
-    "\n",
-    "This notebook serves as the primary entry point for the EasyTranslate project.\n",
-    "It handles:\n",
-    "- Environment detection (local vs Google Colab)\n",
-    "- Repository cloning and dependency installation\n",
-    "- Data loading and preprocessing\n",
-    "- Model construction and training\n",
-    "- Evaluation and cloud storage synchronization\n",
-    "\n",
-    "---"
-   ]
-  },
-  {
-   "cell_type": "markdown",
-   "metadata": {},
-   "source": [
-    "## 1. Environment Detection & Setup\n",
-    "\n",
-    "Detect whether we are running in Google Colab or locally, and configure accordingly."
-   ]
-  },
-  {
-   "cell_type": "code",
-   "execution_count": null,
-   "metadata": {},
-   "outputs": [],
-   "source": [
-    "import os\n",
-    "import sys\n",
-    "import subprocess\n",
-    "import importlib\n",
-    "import json\n",
-    "import shutil\n",
-    "from pathlib import Path\n",
-    "\n",
-    "IN_COLAB = False\n",
-    "try:\n",
-    "    import google.colab\n",
-    "    IN_COLAB = True\n",
-    "except ImportError:\n",
-    "    pass\n",
-    "\n",
-    "print(f\"Running in Google Colab: {IN_COLAB}\")\n",
-    "print(f\"Python version: {sys.version}\")\n",
-    "print(f\"Working directory: {os.getcwd()}\")\n"
-   ]
-  },
-  {
-   "cell_type": "markdown",
-   "metadata": {},
-   "source": [
-    "## 2. Repository Setup\n",
-    "\n",
-    "Clone the latest code from the repository. In Colab, this pulls fresh code each runtime.\n",
-    "Locally, it ensures the working directory is correct."
-   ]
-  },
-  {
-   "cell_type": "code",
-   "execution_count": null,
-   "metadata": {},
-   "outputs": [],
-   "source": [
-    "# ── 仓库地址配置 ─────────────────────────────────────────────────────────────\n",
-    "#   import os; os.environ[\"EASYTRANSLATE_REPO_URL\"] = \"https://github.com/your-org/your-repo.git\"\n",
-    "REPO_URL = os.environ.get(\n",
-    "    \"EASYTRANSLATE_REPO_URL\",\n",
-    "    \"https://huggingface.co/sdfjliom/UCAS-EasyTranslate\",\n",
-    ")\n",
-    "_DEFAULT_COLAB_DIR = Path(\"/content/UCAS-EasyTranslate\")\n",
-    "\n",
-    "\n",
-    "def _is_repo_root(path: Path) -> bool:\n",
-    "    return (path / \"src\" / \"easytranslate\").exists() and (path / \"setup.py\").exists()\n",
-    "\n",
-    "\n",
-    "def _find_repo_root(start_path: Path):\n",
-    "    p = start_path.resolve()\n",
-    "    for candidate in [p] + list(p.parents):\n",
-    "        if _is_repo_root(candidate):\n",
-    "            return candidate\n",
-    "    return None\n",
-    "\n",
-    "\n",
-    "resolved_repo = None\n",
-    "\n",
-    "if IN_COLAB:\n",
-    "    if _is_repo_root(_DEFAULT_COLAB_DIR):\n",
-    "        resolved_repo = _DEFAULT_COLAB_DIR\n",
-    "    else:\n",
-    "        print(f\"Cloning repository from: {REPO_URL}\")\n",
-    "        try:\n",
-    "            subprocess.run(\n",
-    "                [\"git\", \"lfs\", \"install\"], check=False,\n",
-    "                stdout=subprocess.DEVNULL, stderr=subprocess.DEVNULL,\n",
-    "            )\n",
-    "            subprocess.run(\n",
-    "                [\"git\", \"clone\", \"--depth\", \"1\", REPO_URL, str(_DEFAULT_COLAB_DIR)],\n",
-    "                check=True,\n",
-    "                stdout=subprocess.PIPE, stderr=subprocess.STDOUT, text=True,\n",
-    "            )\n",
-    "            resolved_repo = _DEFAULT_COLAB_DIR\n",
-    "            print(f\"Cloned to: {_DEFAULT_COLAB_DIR}\")\n",
-    "        except subprocess.CalledProcessError as e:\n",
-    "            print(f\"Clone failed:\\n{e.stdout}\")\n",
-    "\n",
-    "    # Fallback: search Drive or current dir\n",
-    "    if resolved_repo is None:\n",
-    "        for candidate in [\n",
-    "            Path.cwd(),\n",
-    "            Path(\"/content\"),\n",
-    "            Path(\"/content/drive/MyDrive/UCAS-EasyTranslate\"),\n",
-    "            Path(\"/content/drive/MyDrive/Colab Notebooks/UCAS-EasyTranslate\"),\n",
-    "        ]:\n",
-    "            root = _find_repo_root(candidate)\n",
-    "            if root:\n",
-    "                resolved_repo = root\n",
-    "                print(f\"Found existing repo at: {resolved_repo}\")\n",
-    "                break\n",
-    "\n",
-    "    if resolved_repo is None:\n",
-    "        raise FileNotFoundError(\n",
-    "            \"Cannot locate UCAS-EasyTranslate repository.\\n\"\n",
-    "            \"Options:\\n\"\n",
-    "            \"  (1) Set EASYTRANSLATE_REPO_URL to a publicly accessible git URL.\\n\"\n",
-    "            \"  (2) Manually clone to /content/UCAS-EasyTranslate.\\n\"\n",
-    "            \"  (3) Place repo in Google Drive and mount Drive first.\"\n",
-    "        )\n",
-    "else:\n",
-    "    resolved_repo = _find_repo_root(Path.cwd()) or Path.cwd()\n",
-    "    print(f\"Using local repository at: {resolved_repo}\")\n",
-    "\n",
-    "# Always keep REPO_DIR in sync with wherever the repo actually is\n",
-    "REPO_DIR = Path(resolved_repo)\n",
-    "os.chdir(REPO_DIR)\n",
-    "sys.path.insert(0, str(REPO_DIR / \"src\"))\n",
-    "print(f\"Repository directory: {REPO_DIR}\")\n"
-   ]
-  },
-  {
-   "cell_type": "markdown",
-   "metadata": {},
-   "source": [
-    "## 3. Dependency Installation\n",
-    "\n",
-    "Install all required packages. In Colab, PyTorch is pre-installed."
-   ]
-  },
-  {
-   "cell_type": "code",
-   "execution_count": null,
-   "metadata": {},
-   "outputs": [],
-   "source": [
-    "if IN_COLAB:\n",
-    "    # 1. Only install packages that Colab does NOT ship.\n",
-    "    #    Do NOT touch torch/numpy/pandas — Colab's preinstalled versions are fine.\n",
-    "    %pip install -q --upgrade pip setuptools wheel\n",
-    "\n",
-    "    # Core HuggingFace packages (Colab may have older versions)\n",
-    "    %pip install -q \"transformers>=4.36.0,<4.45.0\" \\\n",
-    "                    \"datasets>=2.16.0,<3.0.0\" \\\n",
-    "                    \"tokenizers>=0.15.0,<0.20.0\" \\\n",
-    "                    \"sentencepiece>=0.2.0\" \\\n",
-    "                    \"accelerate>=0.25.0,<0.35.0\" \\\n",
-    "                    \"peft>=0.7.0,<0.12.0\"\n",
-    "\n",
-    "    # Evaluation / config packages\n",
-    "    %pip install -q \"sacrebleu>=2.4.0\" \\\n",
-    "                    \"omegaconf>=2.3.0,<3.0.0\" \\\n",
-    "                    \"rich>=13.0.0\"\n",
-    "\n",
-    "    # 2. Install project code without re-resolving heavy dependencies.\n",
-    "    #    (torch, numpy, etc. are already present from Colab runtime.)\n",
-    "    %pip install -q --no-deps -e .\n",
-    "\n",
-    "    print(\"All packages installed successfully.\")\n",
-    "else:\n",
-    "    %pip install -q --upgrade pip setuptools wheel\n",
-    "    %pip install -q -r requirements.txt\n",
-    "    %pip install -q -e .\n",
-    "    print(\"Dependencies installed successfully.\")\n"
-   ]
-  },
-  {
-   "cell_type": "markdown",
-   "metadata": {},
-   "source": [
-    "## 4. GPU Verification\n",
-    "\n",
-    "Verify GPU availability and display hardware information."
-   ]
-  },
-  {
-   "cell_type": "code",
-   "execution_count": null,
-   "metadata": {},
-   "outputs": [],
-   "source": [
-    "import torch\n",
-    "\n",
-    "print(f\"PyTorch version: {torch.__version__}\")\n",
-    "print(f\"CUDA available: {torch.cuda.is_available()}\")\n",
-    "\n",
-    "if torch.cuda.is_available():\n",
-    "    print(f\"CUDA version: {torch.version.cuda}\")\n",
-    "    print(f\"GPU count: {torch.cuda.device_count()}\")\n",
-    "    for i in range(torch.cuda.device_count()):\n",
-    "        print(f\"  GPU {i}: {torch.cuda.get_device_name(i)}\")\n",
-    "        props = torch.cuda.get_device_properties(i)\n",
-    "        print(f\"    Memory: {props.total_memory / 1024**3:.1f} GB\")\n",
-    "        print(f\"    Compute Capability: {props.major}.{props.minor}\")\n",
-    "else:\n",
-    "    print(\"WARNING: No GPU detected. Training will be very slow on CPU.\")"
-   ]
-  },
-  {
-   "cell_type": "markdown",
-   "metadata": {},
-   "source": [
-    "## 5. Google Drive Mount (Colab Only)\n",
-    "\n",
-    "Mount Google Drive for persistent storage of checkpoints and results."
-   ]
-  },
-  {
-   "cell_type": "code",
-   "execution_count": null,
-   "metadata": {},
-   "outputs": [],
-   "source": [
-    "DRIVE_MOUNTED = False\n",
-    "DRIVE_BASE = \"/content/drive/MyDrive/EasyTranslate\"\n",
-    "\n",
-    "if IN_COLAB:\n",
-    "    from google.colab import drive\n",
-    "    drive.mount(\"/content/drive\")\n",
-    "    DRIVE_MOUNTED = os.path.exists(\"/content/drive\")\n",
-    "    if DRIVE_MOUNTED:\n",
-    "        os.makedirs(DRIVE_BASE, exist_ok=True)\n",
-    "        print(f\"Google Drive mounted. Base path: {DRIVE_BASE}\")\n",
-    "    else:\n",
-    "        print(\"WARNING: Google Drive mount failed\")\n",
-    "else:\n",
-    "    print(\"Not in Colab, skipping Google Drive mount\")"
-   ]
-  },
-  {
-   "cell_type": "markdown",
-   "metadata": {},
-   "source": [
-    "## 6. Configuration Loading\n",
-    "\n",
-    "Load and display the project configuration."
-   ]
-  },
-  {
-   "cell_type": "code",
-   "execution_count": null,
-   "metadata": {},
-   "outputs": [],
-   "source": [
-    "import torch\n",
-    "import numpy as np\n",
-    "\n",
-    "print(f\"NumPy  version : {np.__version__}\")\n",
-    "print(f\"PyTorch version: {torch.__version__}\")\n",
-    "\n",
-    "from easytranslate.utils.config import load_config, config_to_dict\n",
-    "from easytranslate.utils.seed import set_seed\n",
-    "from easytranslate.utils.logging import setup_logging\n",
-    "\n",
-    "config = load_config(\"configs/default_config.yaml\")\n",
-    "config_dict = config_to_dict(config)\n",
-    "\n",
-    "exp_cfg  = config_dict.get(\"experiment\", {})\n",
-    "seed     = exp_cfg.get(\"seed\", 42)\n",
-    "set_seed(seed)\n",
-    "\n",
-    "log_cfg = config_dict.get(\"logging\", {})\n",
-    "setup_logging(\n",
-    "    log_dir=log_cfg.get(\"log_dir\", \"logs/\"),\n",
-    "    log_file=\"easytranslate.log\",\n",
-    ")\n",
-    "\n",
-    "# ── Colab runtime overrides ───────────────────────────────────────────────────\n",
-    "if IN_COLAB:\n",
-    "    train_cfg = config_dict.setdefault(\"training\", {})\n",
-    "\n",
-    "    # Mixed precision: use bf16 on Ampere+ GPUs, fp16 otherwise, none on CPU\n",
-    "    if torch.cuda.is_available():\n",
-    "        gpu_cap = torch.cuda.get_device_capability(0)\n",
-    "        if gpu_cap[0] >= 8:           # A100/A10 → bf16\n",
-    "            train_cfg[\"fp16\"] = False\n",
-    "            train_cfg[\"bf16\"] = True\n",
-    "        else:                         # T4/P100/V100 → fp16\n",
-    "            train_cfg[\"fp16\"] = True\n",
-    "            train_cfg[\"bf16\"] = False\n",
-    "    else:\n",
-    "        train_cfg[\"fp16\"] = False\n",
-    "        train_cfg[\"bf16\"] = False\n",
-    "\n",
-    "    # Reduce epochs for a Colab demo run\n",
-    "    train_cfg.setdefault(\"epochs\", 10)\n",
-    "\n",
-    "    # Colab-friendly batch / gradient accumulation settings\n",
-    "    train_cfg.setdefault(\"batch_size\", 32)\n",
-    "    train_cfg.setdefault(\"gradient_accumulation_steps\", 4)\n",
-    "\n",
-    "print(f\"Configuration loaded. Experiment seed: {seed}\")\n",
-    "print(f\"Model type : {config_dict['model']['type']}\")\n",
-    "print(f\"Training   : epochs={config_dict['training']['epochs']}, \"\n",
-    "      f\"fp16={config_dict['training']['fp16']}, \"\n",
-    "      f\"bf16={config_dict['training']['bf16']}\")\n"
-   ]
-  },
-  {
-   "cell_type": "markdown",
-   "metadata": {},
-   "source": [
-    "## 7. Data Loading & Preprocessing\n",
-    "\n",
-    "Load the WMT19 zh-en dataset, train the BPE tokenizer, and prepare DataLoaders."
-   ]
-  },
-  {
-   "cell_type": "code",
-   "execution_count": null,
-   "metadata": {},
-   "outputs": [],
-   "source": [
-    "from easytranslate.data import (\n",
-    "    TranslationDataset,\n",
-    "    TranslationCollator,\n",
-    "    DynamicBatchSampler,\n",
-    "    build_tokenizer,\n",
-    "    load_wmt_dataset,\n",
-    "    preprocess_pipeline,\n",
-    ")\n",
-    "from torch.utils.data import DataLoader\n",
-    "\n",
-    "data_cfg   = config_dict.get(\"data\", {})\n",
-    "preproc_cfg = data_cfg.get(\"preprocessing\", {})\n",
-    "loader_cfg  = data_cfg.get(\"dataloader\", {})\n",
-    "\n",
-    "# ── Cap dataset sizes for Colab to avoid RAM/time issues ─────────────────────\n",
-    "MAX_TRAIN_SAMPLES = 200_000 if IN_COLAB else None   # None → use full dataset\n",
-    "MAX_VAL_SAMPLES   = 5_000   if IN_COLAB else None\n",
-    "\n",
-    "print(\"Loading WMT19 zh-en dataset (this may take several minutes on first run)...\")\n",
-    "raw_dataset = load_wmt_dataset(\n",
-    "    year=data_cfg.get(\"wmt\", {}).get(\"year\", \"19\"),\n",
-    "    language_pair=data_cfg.get(\"wmt\", {}).get(\"language_pair\", \"zh-en\"),\n",
-    "    src_lang=\"en\",\n",
-    "    tgt_lang=\"zh\",\n",
-    ")\n",
-    "\n",
-    "train_raw = raw_dataset[\"train\"]\n",
-    "\n",
-    "# Prefer \"validation\", fall back to \"dev\", then use a small slice of train\n",
-    "_val_split_name = next(\n",
-    "    (k for k in (\"validation\", \"dev\", \"valid\") if k in raw_dataset),\n",
-    "    None,\n",
-    ")\n",
-    "val_raw = raw_dataset[_val_split_name] if _val_split_name else train_raw\n",
-    "\n",
-    "print(f\"Raw training samples  : {len(train_raw['src'])}\")\n",
-    "print(f\"Raw validation samples: {len(val_raw['src'])} \"\n",
-    "      f\"(split='{_val_split_name or 'train (fallback)'}')\")\n",
-    "\n",
-    "# Apply sample caps BEFORE preprocessing to save time\n",
-    "train_src_raw = train_raw[\"src\"][:MAX_TRAIN_SAMPLES] if MAX_TRAIN_SAMPLES else list(train_raw[\"src\"])\n",
-    "train_tgt_raw = train_raw[\"tgt\"][:MAX_TRAIN_SAMPLES] if MAX_TRAIN_SAMPLES else list(train_raw[\"tgt\"])\n",
-    "val_src_raw   = val_raw[\"src\"][:MAX_VAL_SAMPLES]     if MAX_VAL_SAMPLES   else list(val_raw[\"src\"])\n",
-    "val_tgt_raw   = val_raw[\"tgt\"][:MAX_VAL_SAMPLES]     if MAX_VAL_SAMPLES   else list(val_raw[\"tgt\"])\n",
-    "\n",
-    "print(f\"\\nPreprocessing training data ({len(train_src_raw)} samples)...\")\n",
-    "train_src, train_tgt = preprocess_pipeline(\n",
-    "    train_src_raw, train_tgt_raw,\n",
-    "    max_src_len=preproc_cfg.get(\"max_src_len\", 256),\n",
-    "    max_tgt_len=preproc_cfg.get(\"max_tgt_len\", 256),\n",
-    "    filter_by_length_enabled=preproc_cfg.get(\"filter_by_length\", True),\n",
-    "    length_ratio_threshold=preproc_cfg.get(\"length_ratio_threshold\", 3.0),\n",
-    ")\n",
-    "print(f\"Preprocessed training samples  : {len(train_src)}\")\n",
-    "\n",
-    "print(f\"Preprocessing validation data ({len(val_src_raw)} samples)...\")\n",
-    "val_src, val_tgt = preprocess_pipeline(\n",
-    "    val_src_raw, val_tgt_raw,\n",
-    "    max_src_len=preproc_cfg.get(\"max_src_len\", 256),\n",
-    "    max_tgt_len=preproc_cfg.get(\"max_tgt_len\", 256),\n",
-    "    filter_by_length_enabled=preproc_cfg.get(\"filter_by_length\", True),\n",
-    "    length_ratio_threshold=preproc_cfg.get(\"length_ratio_threshold\", 3.0),\n",
-    ")\n",
-    "print(f\"Preprocessed validation samples: {len(val_src)}\")\n"
-   ]
-  },
-  {
-   "cell_type": "markdown",
-   "metadata": {},
-   "source": [
-    "## 8. Tokenizer Training\n",
-    "\n",
-    "Train a byte-level BPE tokenizer on the combined source and target texts."
-   ]
-  },
-  {
-   "cell_type": "code",
-   "execution_count": null,
-   "metadata": {},
-   "outputs": [],
-   "source": [
-    "tok_cfg = config_dict.get(\"tokenizer\", {})\n",
-    "\n",
-    "print(\"Building tokenizer...\")\n",
-    "all_train_texts = train_src + train_tgt\n",
-    "tokenizer = build_tokenizer(tok_cfg, train_texts=all_train_texts)\n",
-    "\n",
-    "print(f\"Tokenizer vocabulary size: {tokenizer.vocab_size}\")\n",
-    "print(f\"Special tokens: PAD={tokenizer.pad_token_id}, BOS={tokenizer.bos_token_id}, EOS={tokenizer.eos_token_id}\")\n",
-    "\n",
-    "test_encode = tokenizer.encode(\"Hello world\", add_special_tokens=True)\n",
-    "test_decode = tokenizer.decode(test_encode)\n",
-    "print(f\"Encode test: {test_encode[:10]}...\")\n",
-    "print(f\"Decode test: {test_decode[:50]}...\")"
-   ]
-  },
-  {
-   "cell_type": "markdown",
-   "metadata": {},
-   "source": [
-    "## 9. Dataset & DataLoader Construction\n",
-    "\n",
-    "Build PyTorch datasets and dataloaders with dynamic batching."
-   ]
-  },
-  {
-   "cell_type": "code",
-   "execution_count": null,
-   "metadata": {},
-   "outputs": [],
-   "source": [
-    "max_src_len = preproc_cfg.get(\"max_src_len\", 256)\n",
-    "max_tgt_len = preproc_cfg.get(\"max_tgt_len\", 256)\n",
-    "\n",
-    "train_dataset = TranslationDataset(\n",
-    "    train_src, train_tgt,\n",
-    "    tokenizer=tokenizer,\n",
-    "    max_src_len=max_src_len,\n",
-    "    max_tgt_len=max_tgt_len,\n",
-    ")\n",
-    "val_dataset = TranslationDataset(\n",
-    "    val_src, val_tgt,\n",
-    "    tokenizer=tokenizer,\n",
-    "    max_src_len=max_src_len,\n",
-    "    max_tgt_len=max_tgt_len,\n",
-    ")\n",
-    "\n",
-    "collator = TranslationCollator(\n",
-    "    pad_token_id=tokenizer.pad_token_id,\n",
-    "    label_pad_token_id=-100,\n",
-    ")\n",
-    "\n",
-    "loader_cfg  = config_dict.get(\"data\", {}).get(\"dataloader\", {})\n",
-    "batch_size  = loader_cfg.get(\"batch_size\", 32)\n",
-    "# 2 workers on GPU Colab; 0 on CPU (no multiprocessing overhead)\n",
-    "num_workers = 2 if IN_COLAB and torch.cuda.is_available() else 0\n",
-    "# Disable dynamic batching on Colab: computing exact token lengths for 200k\n",
-    "# sentences requires a full tokenizer pass and can take 30+ minutes.\n",
-    "use_dynamic = loader_cfg.get(\"dynamic_batching\", True) and not IN_COLAB\n",
-    "\n",
-    "if use_dynamic:\n",
-    "    max_tokens = loader_cfg.get(\"max_tokens_per_batch\", 8192)\n",
-    "    print(\"Computing sequence lengths for dynamic batching (approximate)...\")\n",
-    "\n",
-    "    def _approx_len(src_text: str, tgt_text: str) -> int:\n",
-    "        \"\"\"Fast character-based length estimate — no tokenizer call needed.\"\"\"\n",
-    "        def _tok_est(t: str) -> int:\n",
-    "            cjk = sum(1 for c in t if \"\\u4e00\" <= c <= \"\\u9fff\")\n",
-    "            return (len(t) - cjk) // 4 + cjk + 2   # rough BPE estimate\n",
-    "        return max(_tok_est(src_text), _tok_est(tgt_text))\n",
-    "\n",
-    "    train_lengths = [_approx_len(s, t) for s, t in zip(train_src, train_tgt)]\n",
-    "    train_sampler = DynamicBatchSampler(\n",
-    "        train_lengths,\n",
-    "        max_tokens_per_batch=max_tokens,\n",
-    "        shuffle=True,\n",
-    "    )\n",
-    "    train_loader = DataLoader(\n",
-    "        train_dataset,\n",
-    "        batch_sampler=train_sampler,\n",
-    "        collate_fn=collator,\n",
-    "        num_workers=num_workers,\n",
-    "        pin_memory=torch.cuda.is_available(),\n",
-    "    )\n",
-    "else:\n",
-    "    if IN_COLAB and loader_cfg.get(\"dynamic_batching\", True):\n",
-    "        print(\"Note: dynamic batching disabled on Colab (would require full tokenizer pass on all samples).\")\n",
-    "    train_loader = DataLoader(\n",
-    "        train_dataset,\n",
-    "        batch_size=batch_size,\n",
-    "        shuffle=True,\n",
-    "        collate_fn=collator,\n",
-    "        num_workers=num_workers,\n",
-    "        pin_memory=torch.cuda.is_available(),\n",
-    "    )\n",
-    "\n",
-    "val_loader = DataLoader(\n",
-    "    val_dataset,\n",
-    "    batch_size=batch_size,\n",
-    "    shuffle=False,\n",
-    "    collate_fn=collator,\n",
-    "    num_workers=num_workers,\n",
-    "    pin_memory=torch.cuda.is_available(),\n",
-    ")\n",
-    "\n",
-    "print(f\"Training batches : ~{len(train_loader)}\")\n",
-    "print(f\"Validation batches: {len(val_loader)}\")\n",
-    "\n",
-    "sample_batch = next(iter(train_loader))\n",
-    "print(\"Sample batch shapes:\")\n",
-    "for k, v in sample_batch.items():\n",
-    "    if isinstance(v, torch.Tensor):\n",
-    "        print(f\"  {k}: {list(v.shape)}\")\n"
-   ]
-  },
-  {
-   "cell_type": "markdown",
-   "metadata": {},
-   "source": [
-    "## 10. Model Construction\n",
-    "\n",
-    "Build the Transformer model based on configuration."
-   ]
-  },
-  {
-   "cell_type": "code",
-   "execution_count": null,
-   "metadata": {},
-   "outputs": [],
-   "source": [
-    "from easytranslate.model import TransformerTranslationModel\n",
-    "\n",
-    "model_cfg  = config_dict.get(\"model\", {})\n",
-    "model_type = model_cfg.get(\"type\", \"transformer_scratch\")\n",
-    "\n",
-    "if model_type == \"transformer_scratch\":\n",
-    "    tf_cfg = dict(model_cfg.get(\"transformer\", {}))\n",
-    "\n",
-    "    # ── Colab quick-run: smaller model to fit in Colab RAM/VRAM ──────────────\n",
-    "    # Default full model: d_model=512, 6 enc/dec layers (~75 M params)\n",
-    "    # Colab quick model:  d_model=256, 3 enc/dec layers (~12 M params)\n",
-    "    # Set COLAB_FULL_MODEL=1 in env to skip this override.\n",
-    "    if IN_COLAB and not os.environ.get(\"COLAB_FULL_MODEL\"):\n",
-    "        tf_cfg.setdefault(\"d_model\",            256)\n",
-    "        tf_cfg.setdefault(\"nhead\",              4)\n",
-    "        tf_cfg.setdefault(\"num_encoder_layers\", 3)\n",
-    "        tf_cfg.setdefault(\"num_decoder_layers\", 3)\n",
-    "        tf_cfg.setdefault(\"dim_feedforward\",    1024)\n",
-    "        print(\"Colab mode: using compact model (d_model=256, 3 layers).\")\n",
-    "        print(\"To use the full model, run:  import os; os.environ['COLAB_FULL_MODEL']='1'\")\n",
-    "\n",
-    "    model = TransformerTranslationModel(\n",
-    "        src_vocab_size=tokenizer.vocab_size,\n",
-    "        tgt_vocab_size=tokenizer.vocab_size,\n",
-    "        d_model=tf_cfg.get(\"d_model\", 512),\n",
-    "        nhead=tf_cfg.get(\"nhead\", 8),\n",
-    "        num_encoder_layers=tf_cfg.get(\"num_encoder_layers\", 6),\n",
-    "        num_decoder_layers=tf_cfg.get(\"num_decoder_layers\", 6),\n",
-    "        dim_feedforward=tf_cfg.get(\"dim_feedforward\", 2048),\n",
-    "        dropout=tf_cfg.get(\"dropout\", 0.1),\n",
-    "        activation=tf_cfg.get(\"activation\", \"gelu\"),\n",
-    "        max_seq_len=tf_cfg.get(\"max_seq_len\", 512),\n",
-    "        use_flash_attention=tf_cfg.get(\"use_flash_attention\", True),\n",
-    "        use_rotary_embedding=tf_cfg.get(\"use_rotary_embedding\", True),\n",
-    "        pre_norm=tf_cfg.get(\"pre_norm\", True),\n",
-    "        pad_id=tokenizer.pad_token_id,\n",
-    "        share_embedding=False,\n",
-    "    )\n",
-    "    print(\"Built Transformer from scratch\")\n",
-    "\n",
-    "elif model_type in (\"finetune_nllb\", \"finetune_mbart\"):\n",
-    "    from easytranslate.model.finetune import load_pretrained_model, setup_lora\n",
-    "    pt_cfg = model_cfg.get(\"pretrained\", {})\n",
-    "    model, hf_tokenizer = load_pretrained_model(\n",
-    "        model_name=pt_cfg.get(\"model_name\", \"facebook/nllb-200-distilled-600M\"),\n",
-    "        src_lang=pt_cfg.get(\"src_lang\", \"eng_Latn\"),\n",
-    "        tgt_lang=pt_cfg.get(\"tgt_lang\", \"zho_Hans\"),\n",
-    "    )\n",
-    "    if pt_cfg.get(\"use_lora\", True):\n",
-    "        lora_cfg = pt_cfg.get(\"lora\", {})\n",
-    "        model = setup_lora(\n",
-    "            model,\n",
-    "            r=lora_cfg.get(\"r\", 16),\n",
-    "            alpha=lora_cfg.get(\"alpha\", 32),\n",
-    "            dropout=lora_cfg.get(\"dropout\", 0.05),\n",
-    "            target_modules=lora_cfg.get(\"target_modules\", [\"q_proj\", \"v_proj\"]),\n",
-    "        )\n",
-    "    print(f\"Loaded pretrained model: {pt_cfg.get('model_name')}\")\n",
-    "\n",
-    "else:\n",
-    "    raise ValueError(f\"Unknown model type: {model_type}\")\n",
-    "\n",
-    "total_params    = sum(p.numel() for p in model.parameters())\n",
-    "trainable_params = sum(p.numel() for p in model.parameters() if p.requires_grad)\n",
-    "print(f\"Total parameters     : {total_params:,}\")\n",
-    "print(f\"Trainable parameters : {trainable_params:,}\")\n",
-    "print(f\"Trainable ratio      : {100 * trainable_params / total_params:.2f}%\")\n"
-   ]
-  },
-  {
-   "cell_type": "markdown",
-   "metadata": {},
-   "source": [
-    "## 11. Quick Forward Pass Test\n",
-    "\n",
-    "Verify the model can perform a forward pass with correct output dimensions."
-   ]
-  },
-  {
-   "cell_type": "code",
-   "execution_count": null,
-   "metadata": {},
-   "outputs": [],
-   "source": [
-    "device = torch.device(\"cuda\" if torch.cuda.is_available() else \"cpu\")\n",
-    "model = model.to(device)\n",
-    "model.eval()\n",
-    "\n",
-    "test_batch = next(iter(train_loader))\n",
-    "test_src = test_batch[\"src_ids\"][:2].to(device)\n",
-    "test_tgt = test_batch[\"tgt_input_ids\"][:2].to(device)\n",
-    "test_src_mask = test_batch[\"src_padding_mask\"][:2].to(device)\n",
-    "test_tgt_mask = test_batch[\"tgt_padding_mask\"][:2].to(device)\n",
-    "\n",
-    "with torch.no_grad():\n",
-    "    logits = model(test_src, test_tgt, test_src_mask, test_tgt_mask)\n",
-    "\n",
-    "print(f\"Input src shape: {list(test_src.shape)}\")\n",
-    "print(f\"Input tgt shape: {list(test_tgt.shape)}\")\n",
-    "print(f\"Output logits shape: {list(logits.shape)}\")\n",
-    "print(f\"Expected output shape: [B, T, vocab_size] = [{test_tgt.size(0)}, {test_tgt.size(1)}, {tokenizer.vocab_size}]\")\n",
-    "assert logits.size(-1) == tokenizer.vocab_size, f\"Vocab size mismatch: {logits.size(-1)} vs {tokenizer.vocab_size}\"\n",
-    "print(\"Forward pass test PASSED\")"
-   ]
-  },
-  {
-   "cell_type": "markdown",
-   "metadata": {},
-   "source": [
-    "## 12. Training Execution\n",
-    "\n",
-    "This cell DEFINES the training setup but does NOT execute training automatically.\n",
-    "To start training, run the cell below this one."
-   ]
-  },
-  {
-   "cell_type": "code",
-   "execution_count": null,
-   "metadata": {},
-   "outputs": [],
-   "source": [
-    "# sacrebleu is now installed in the dependency cell above.\n",
-    "# This cell just verifies it is importable before we initialize the Trainer.\n",
-    "try:\n",
-    "    importlib.import_module(\"sacrebleu\")\n",
-    "    print(\"sacrebleu OK\")\n",
-    "except ImportError:\n",
-    "    print(\"sacrebleu missing — installing now...\")\n",
-    "    subprocess.run([sys.executable, \"-m\", \"pip\", \"install\", \"-q\", \"sacrebleu>=2.4.0\"], check=True)\n",
-    "\n",
-    "from easytranslate.training import Trainer\n",
-    "from easytranslate.evaluation import Evaluator\n",
-    "\n",
-    "evaluator = Evaluator(\n",
-    "    model=model,\n",
-    "    tokenizer=tokenizer,\n",
-    "    config=config_dict,\n",
-    ")\n",
-    "\n",
-    "trainer = Trainer(\n",
-    "    model=model,\n",
-    "    train_loader=train_loader,\n",
-    "    val_loader=val_loader,\n",
-    "    config=config_dict,\n",
-    "    evaluator=evaluator,\n",
-    ")\n",
-    "\n",
-    "# Output / plot directories live inside the repo\n",
-    "OUTPUT_DIR = REPO_DIR / \"outputs\"\n",
-    "PLOTS_DIR  = OUTPUT_DIR / \"plots\"\n",
-    "OUTPUT_DIR.mkdir(parents=True, exist_ok=True)\n",
-    "PLOTS_DIR.mkdir(parents=True, exist_ok=True)\n",
-    "\n",
-    "print(\"Trainer initialized successfully\")\n",
-    "print(f\"  Device                     : {trainer.device}\")\n",
-    "print(f\"  FP16 / BF16                : {trainer.fp16} / {trainer.bf16}\")\n",
-    "print(f\"  Gradient accumulation steps: {trainer.gradient_accumulation_steps}\")\n",
-    "print(f\"  Number of epochs           : {trainer.num_epochs}\")\n",
-    "print(f\"  Checkpoint directory       : {trainer.checkpoint_dir}\")\n",
-    "print(f\"  Output directory           : {OUTPUT_DIR}\")\n",
-    "print()\n",
-    "print(\">>> Run the NEXT cell to start training.\")\n"
-   ]
-  },
-  {
-   "cell_type": "markdown",
-   "metadata": {},
-   "source": [
-    "## 13. Start Training\n",
-    "\n",
-    "**Run this cell to begin training.** This will execute the full training loop.\n",
-    "Training progress will be displayed via tqdm progress bars."
-   ]
-  },
-  {
-   "cell_type": "code",
-   "execution_count": null,
-   "metadata": {},
-   "outputs": [],
-   "source": [
-    "START_TRAINING = True\n",
-    "\n",
-    "if START_TRAINING:\n",
-    "    print(\"=\" * 60)\n",
-    "    print(\"  Starting Training...\")\n",
-    "    print(\"=\" * 60)\n",
-    "    trainer.train()\n",
-    "else:\n",
-    "    print(\"Training skipped. Set START_TRAINING = True to begin.\")"
-   ]
-  },
-  {
-   "cell_type": "markdown",
-   "metadata": {},
-   "source": [
-    "## 14. Evaluation on Test Set\n",
-    "\n",
-    "After training completes, evaluate the best model on the validation set."
-   ]
-  },
-  {
-   "cell_type": "code",
-   "execution_count": null,
-   "metadata": {},
-   "outputs": [],
-   "source": [
-    "EVAL_RESULTS = {}\n",
-    "\n",
-    "best_ckpt = trainer.checkpoint_dir / \"best_model.pt\"\n",
-    "\n",
-    "if best_ckpt.exists():\n",
-    "    print(f\"Loading best model from {best_ckpt}\")\n",
-    "    checkpoint = torch.load(best_ckpt, map_location=device, weights_only=True)\n",
-    "    model.load_state_dict(checkpoint[\"model_state_dict\"])\n",
-    "    model = model.to(device)\n",
-    "    model.eval()\n",
-    "\n",
-    "    evaluator = Evaluator(model=model, tokenizer=tokenizer, config=config_dict)\n",
-    "\n",
-    "    print(\"Running evaluation on validation set...\")\n",
-    "    EVAL_RESULTS = evaluator.evaluate(\n",
-    "        val_loader,\n",
-    "        src_texts=val_src,\n",
-    "        ref_texts=val_tgt,\n",
-    "    )\n",
-    "\n",
-    "    print(\"\\n\" + \"=\" * 60)\n",
-    "    print(\"  Evaluation Results\")\n",
-    "    print(\"=\" * 60)\n",
-    "    for metric, score in EVAL_RESULTS.items():\n",
-    "        if isinstance(score, (int, float)):\n",
-    "            print(f\"  {metric:>12s}: {score:.4f}\")\n",
-    "\n",
-    "    eval_path = OUTPUT_DIR / \"evaluation_results.json\"\n",
-    "    with open(eval_path, \"w\", encoding=\"utf-8\") as f:\n",
-    "        json.dump(EVAL_RESULTS, f, indent=2, ensure_ascii=False)\n",
-    "    print(f\"\\nSaved evaluation results → {eval_path}\")\n",
-    "else:\n",
-    "    print(\"No best_model.pt found. Run training first (Cell 13).\")\n"
-   ]
-  },
-  {
-   "cell_type": "markdown",
-   "metadata": {},
-   "source": [
-    "## 15. Translation Demo\n",
-    "\n",
-    "Test the trained model with some example translations."
-   ]
-  },
-  {
-   "cell_type": "code",
-   "execution_count": null,
-   "metadata": {},
-   "outputs": [],
-   "source": [
-    "test_sentences = [\n",
-    "    \"Hello, how are you today?\",\n",
-    "    \"Machine translation is an important field of natural language processing.\",\n",
-    "    \"The weather is beautiful and I want to go for a walk.\",\n",
-    "    \"Deep learning has revolutionized artificial intelligence research.\",\n",
-    "]\n",
-    "\n",
-    "TRANSLATION_RESULTS = []\n",
-    "\n",
-    "if best_ckpt.exists():\n",
-    "    print(\"Translating example sentences...\")\n",
-    "    print(\"-\" * 60)\n",
-    "    for sentence in test_sentences:\n",
-    "        translation = evaluator.translate_single(sentence)\n",
-    "        TRANSLATION_RESULTS.append({\"source_en\": sentence, \"target_zh\": translation})\n",
-    "        print(f\"[EN] {sentence}\")\n",
-    "        print(f\"[ZH] {translation}\")\n",
-    "        print()\n",
-    "\n",
-    "    translation_path = OUTPUT_DIR / \"translation_examples.json\"\n",
-    "    with open(translation_path, \"w\", encoding=\"utf-8\") as f:\n",
-    "        json.dump(TRANSLATION_RESULTS, f, indent=2, ensure_ascii=False)\n",
-    "    print(f\"Saved translation examples → {translation_path}\")\n",
-    "else:\n",
-    "    print(\"No trained model checkpoint available. Run training first (Cell 13).\")\n"
-   ]
-  },
-  {
-   "cell_type": "markdown",
-   "metadata": {},
-   "source": [
-    "## 16. Cloud Storage Synchronization\n",
-    "\n",
-    "Sync all training artifacts (checkpoints, logs, summaries) to Google Drive."
-   ]
-  },
-  {
-   "cell_type": "code",
-   "execution_count": null,
-   "metadata": {},
-   "outputs": [],
-   "source": [
-    "from easytranslate.utils.cloud_storage import sync_all_to_drive\n",
-    "\n",
-    "# Use trainer.log_dir if available, otherwise fall back to a sensible default\n",
-    "_log_dir = getattr(trainer, \"log_dir\", REPO_DIR / \"logs\")\n",
-    "\n",
-    "if IN_COLAB and DRIVE_MOUNTED:\n",
-    "    print(\"Syncing training artifacts to Google Drive...\")\n",
-    "    sync_results = sync_all_to_drive(\n",
-    "        checkpoint_dir=str(trainer.checkpoint_dir),\n",
-    "        log_dir=str(_log_dir),\n",
-    "        drive_base_path=DRIVE_BASE,\n",
-    "    )\n",
-    "\n",
-    "    drive_outputs_dir = Path(DRIVE_BASE) / \"outputs\"\n",
-    "    drive_outputs_dir.mkdir(parents=True, exist_ok=True)\n",
-    "\n",
-    "    for artifact_file in [\n",
-    "        OUTPUT_DIR / \"evaluation_results.json\",\n",
-    "        OUTPUT_DIR / \"translation_examples.json\",\n",
-    "        OUTPUT_DIR / \"training_report.json\",\n",
-    "    ]:\n",
-    "        if artifact_file.exists():\n",
-    "            shutil.copy2(artifact_file, drive_outputs_dir / artifact_file.name)\n",
-    "            print(f\"  Copied {artifact_file.name} → Drive\")\n",
-    "\n",
-    "    if PLOTS_DIR.exists():\n",
-    "        drive_plots_dir = drive_outputs_dir / \"plots\"\n",
-    "        drive_plots_dir.mkdir(parents=True, exist_ok=True)\n",
-    "        for png_file in PLOTS_DIR.glob(\"*.png\"):\n",
-    "            shutil.copy2(png_file, drive_plots_dir / png_file.name)\n",
-    "            print(f\"  Copied plot {png_file.name} → Drive\")\n",
-    "\n",
-    "    print(f\"\\nSync complete. Drive base: {DRIVE_BASE}\")\n",
-    "    print(f\"Sync details: {sync_results}\")\n",
-    "\n",
-    "elif not IN_COLAB:\n",
-    "    print(\"Running locally. Artifacts are already on disk:\")\n",
-    "    print(f\"  Checkpoints : {trainer.checkpoint_dir}\")\n",
-    "    print(f\"  Logs        : {_log_dir}\")\n",
-    "    print(f\"  Outputs     : {OUTPUT_DIR}\")\n",
-    "else:\n",
-    "    print(\"Google Drive not mounted. Artifacts saved locally only.\")\n",
-    "    print(\"Mount Drive (Cell 5) and re-run this cell to sync results.\")\n"
-   ]
-  },
-  {
-   "cell_type": "markdown",
-   "metadata": {},
-   "source": [
-    "## 17. Training Summary\n",
-    "\n",
-    "Display the final training summary including loss curves and best metrics."
-   ]
-  },
-  {
-   "cell_type": "code",
-   "execution_count": null,
-   "metadata": {},
-   "outputs": [],
-   "source": [
-    "TRAINING_SUMMARY = {}\n",
-    "summary_path = trainer.checkpoint_dir / \"training_summary.json\"\n",
-    "\n",
-    "if summary_path.exists():\n",
-    "    with open(summary_path, \"r\", encoding=\"utf-8\") as f:\n",
-    "        TRAINING_SUMMARY = json.load(f)\n",
-    "\n",
-    "    print(\"=\" * 60)\n",
-    "    print(\"  Training Summary\")\n",
-    "    print(\"=\" * 60)\n",
-    "    print(f\"  Best epoch  : {TRAINING_SUMMARY.get('best_epoch', 'N/A')}\")\n",
-    "    print(f\"  Best metric : {TRAINING_SUMMARY.get('metric_name', 'N/A')} = \"\n",
-    "          f\"{TRAINING_SUMMARY.get('best_metric', 'N/A')}\")\n",
-    "    print(f\"  Total steps : {TRAINING_SUMMARY.get('total_steps', 'N/A')}\")\n",
-    "\n",
-    "    losses = TRAINING_SUMMARY.get(\"train_loss_history\", [])\n",
-    "    if losses:\n",
-    "        print(f\"  Initial loss: {losses[0]:.4f}\")\n",
-    "        print(f\"  Final loss  : {losses[-1]:.4f}\")\n",
-    "        print(f\"  Reduction   : {losses[0] - losses[-1]:.4f}\")\n",
-    "\n",
-    "    # Merge all results into one report file\n",
-    "    report = {\n",
-    "        \"training_summary\": TRAINING_SUMMARY,\n",
-    "        \"evaluation_results\": EVAL_RESULTS,\n",
-    "        \"translation_examples\": TRANSLATION_RESULTS,\n",
-    "    }\n",
-    "    report_path = OUTPUT_DIR / \"training_report.json\"\n",
-    "    with open(report_path, \"w\", encoding=\"utf-8\") as f:\n",
-    "        json.dump(report, f, indent=2, ensure_ascii=False)\n",
-    "    print(f\"\\nMerged report saved → {report_path}\")\n",
-    "else:\n",
-    "    print(\"Training summary not yet available. Complete training first.\")\n"
-   ]
-  },
-  {
-   "cell_type": "code",
-   "execution_count": null,
-   "metadata": {},
-   "outputs": [],
-   "source": [
-    "import matplotlib.pyplot as plt\n",
-    "\n",
-    "PLOTS_DIR.mkdir(parents=True, exist_ok=True)\n",
-    "\n",
-    "if not TRAINING_SUMMARY:\n",
-    "    print(\"No training summary found. Run training and summary cells first.\")\n",
-    "else:\n",
-    "    train_losses = TRAINING_SUMMARY.get(\"train_loss_history\", [])\n",
-    "    val_history = TRAINING_SUMMARY.get(\"val_metrics_history\", [])\n",
-    "\n",
-    "    # 1) Train Loss Curve\n",
-    "    if train_losses:\n",
-    "        epochs = list(range(1, len(train_losses) + 1))\n",
-    "        plt.figure(figsize=(8, 5))\n",
-    "        plt.plot(epochs, train_losses, marker=\"o\", linewidth=2)\n",
-    "        plt.title(\"Training Loss by Epoch\")\n",
-    "        plt.xlabel(\"Epoch\")\n",
-    "        plt.ylabel(\"Loss\")\n",
-    "        plt.grid(alpha=0.3)\n",
-    "        loss_plot_path = PLOTS_DIR / \"train_loss_curve.png\"\n",
-    "        plt.tight_layout()\n",
-    "        plt.savefig(loss_plot_path, dpi=180)\n",
-    "        plt.show()\n",
-    "        print(f\"Saved plot: {loss_plot_path}\")\n",
-    "\n",
-    "    # 2) Validation Metrics Curves\n",
-    "    if val_history:\n",
-    "        metric_keys = sorted({k for m in val_history for k in m.keys() if isinstance(m.get(k), (int, float))})\n",
-    "        metric_keys = [k for k in metric_keys if k != \"val_loss\"]\n",
-    "\n",
-    "        if metric_keys:\n",
-    "            n = len(metric_keys)\n",
-    "            rows = (n + 1) // 2\n",
-    "            plt.figure(figsize=(12, max(4, rows * 3.5)))\n",
-    "            for i, key in enumerate(metric_keys, start=1):\n",
-    "                vals = [m.get(key, None) for m in val_history]\n",
-    "                xs = [idx + 1 for idx, v in enumerate(vals) if v is not None]\n",
-    "                ys = [v for v in vals if v is not None]\n",
-    "                if not ys:\n",
-    "                    continue\n",
-    "                plt.subplot(rows, 2, i)\n",
-    "                plt.plot(xs, ys, marker=\"o\", linewidth=1.8)\n",
-    "                plt.title(key)\n",
-    "                plt.xlabel(\"Epoch\")\n",
-    "                plt.ylabel(key)\n",
-    "                plt.grid(alpha=0.3)\n",
-    "\n",
-    "            metrics_plot_path = PLOTS_DIR / \"validation_metrics_curves.png\"\n",
-    "            plt.tight_layout()\n",
-    "            plt.savefig(metrics_plot_path, dpi=180)\n",
-    "            plt.show()\n",
-    "            print(f\"Saved plot: {metrics_plot_path}\")\n",
-    "\n",
-    "    # 3) Final Evaluation Bar Chart\n",
-    "    if \"EVAL_RESULTS\" in globals() and EVAL_RESULTS:\n",
-    "        scalar_items = {k: v for k, v in EVAL_RESULTS.items() if isinstance(v, (int, float))}\n",
-    "        if scalar_items:\n",
-    "            names = list(scalar_items.keys())\n",
-    "            values = [scalar_items[k] for k in names]\n",
-    "            plt.figure(figsize=(10, 5))\n",
-    "            bars = plt.bar(names, values)\n",
-    "            plt.title(\"Final Evaluation Metrics\")\n",
-    "            plt.ylabel(\"Score\")\n",
-    "            plt.xticks(rotation=30)\n",
-    "            plt.grid(axis=\"y\", alpha=0.25)\n",
-    "            for bar, val in zip(bars, values):\n",
-    "                plt.text(bar.get_x() + bar.get_width() / 2, bar.get_height(), f\"{val:.3f}\", ha=\"center\", va=\"bottom\", fontsize=9)\n",
-    "            eval_plot_path = PLOTS_DIR / \"final_evaluation_metrics.png\"\n",
-    "            plt.tight_layout()\n",
-    "            plt.savefig(eval_plot_path, dpi=180)\n",
-    "            plt.show()\n",
-    "            print(f\"Saved plot: {eval_plot_path}\")"
-   ]
-  },
-  {
-   "cell_type": "markdown",
-   "metadata": {},
-   "source": [
-    "---\n",
-    "\n",
-    "## Appendix: Module Architecture Overview\n",
-    "\n",
-    "```\n",
-    "EasyTranslate System Architecture\n",
-    "=================================\n",
-    "\n",
-    "Entry Point: EasyTranslate_Production.ipynb (this notebook)\n",
-    "    |\n",
-    "    +-- Environment Detection (Colab vs Local)\n",
-    "    +-- Repository Cloning (git clone/pull)\n",
-    "    +-- Dependency Installation\n",
-    "    |\n",
-    "    +-- Configuration Layer [utils/config.py]\n",
-    "    |   +-- load_config()      : YAML -> OmegaConf DictConfig\n",
-    "    |   +-- merge_configs()    : CLI overrides merge\n",
-    "    |   +-- config_from_cli()  : Full config pipeline\n",
-    "    |\n",
-    "    +-- Data Layer [data/] — Person A\n",
-    "    |   +-- load_wmt_dataset()     : HuggingFace datasets loader\n",
-    "    |   +-- preprocess_pipeline()  : Clean + filter + deduplicate\n",
-    "    |   +-- build_tokenizer()      : BPE / pretrained tokenizer\n",
-    "    |   +-- TranslationDataset()   : PyTorch Dataset\n",
-    "    |   +-- TranslationCollator()  : Padding + mask generation\n",
-    "    |   +-- DynamicBatchSampler()  : Token-budget batching\n",
-    "    |\n",
-    "    +-- Model Layer [model/] — Person B\n",
-    "    |   +-- TransformerTranslationModel()  : Full Enc-Dec model\n",
-    "    |   +-- TransformerEncoder()           : N-layer encoder\n",
-    "    |   +-- TransformerDecoder()           : N-layer decoder\n",
-    "    |   +-- FlashMultiHeadAttention()      : Flash Attention 2\n",
-    "    |   +-- RotaryPositionalEmbedding()    : RoPE encoding\n",
-    "    |   +-- load_pretrained_model()        : NLLB/mBART loader\n",
-    "    |   +-- setup_lora()                   : LoRA configuration\n",
-    "    |\n",
-    "    +-- Training Layer [training/] — Person C\n",
-    "    |   +-- Trainer()                       : Full training controller\n",
-    "    |   |   +-- _train_one_epoch()          : Mixed precision loop\n",
-    "    |   |   +-- _validate()                 : Validation loop\n",
-    "    |   |   +-- _save_checkpoint()          : Checkpoint persistence\n",
-    "    |   |   +-- _load_checkpoint()          : Resume training\n",
-    "    |   |   +-- _should_early_stop()        : Early stopping logic\n",
-    "    |   |   +-- _setup_distributed()        : DDP/DeepSpeed setup\n",
-    "    |   +-- LabelSmoothedCrossEntropyLoss() : Label smoothing loss\n",
-    "    |   +-- build_optimizer()               : AdamW with param groups\n",
-    "    |   +-- build_scheduler()               : Cosine/InverseSqrt/LR\n",
-    "    |\n",
-    "    +-- Evaluation Layer [evaluation/] — Person D\n",
-    "    |   +-- Evaluator()           : Unified evaluation interface\n",
-    "    |   +-- greedy_decode()       : Greedy decoding\n",
-    "    |   +-- beam_search_decode()  : Beam search decoding\n",
-    "    |   +-- sample_decode()       : Sampling (temp/top-k/top-p)\n",
-    "    |   +-- compute_bleu()        : SacreBLEU metric\n",
-    "    |   +-- compute_comet()       : COMET neural metric\n",
-    "    |   +-- compute_chrf()        : chrF++ metric\n",
-    "    |\n",
-    "    +-- Cloud Storage [utils/cloud_storage.py] — Person C\n",
-    "        +-- is_colab_environment()       : Environment detection\n",
-    "        +-- mount_google_drive()         : Drive authentication\n",
-    "        +-- sync_checkpoints_to_drive()  : Checkpoint backup\n",
-    "        +-- sync_logs_to_drive()         : Log backup\n",
-    "        +-- sync_all_to_drive()          : Full sync pipeline\n",
-    "```\n",
-    "\n",
-    "## Module Interface Contracts\n",
-    "\n",
-    "### Tokenizer Interface (Person A -> B, C, D)\n",
-    "```python\n",
-    "tokenizer.encode(text: str) -> list[int]\n",
-    "tokenizer.decode(ids: list[int]) -> str\n",
-    "tokenizer.vocab_size -> int\n",
-    "tokenizer.pad_token_id -> int\n",
-    "tokenizer.bos_token_id -> int\n",
-    "tokenizer.eos_token_id -> int\n",
-    "```\n",
-    "\n",
-    "### Model Interface (Person B -> C, D)\n",
-    "```python\n",
-    "# Training forward pass\n",
-    "logits = model(src_ids, tgt_input_ids, src_padding_mask, tgt_padding_mask)\n",
-    "# logits: [B, T, vocab_size]\n",
-    "\n",
-    "# Inference\n",
-    "encoder_output = model.encode(src_ids, src_padding_mask)\n",
-    "next_logits = model.decode_step(tgt_input_ids, encoder_output, src_padding_mask)\n",
-    "```\n",
-    "\n",
-    "### Batch Format (Person A -> C)\n",
-    "```python\n",
-    "batch = {\n",
-    "    \"src_ids\": Tensor[B, S],\n",
-    "    \"tgt_input_ids\": Tensor[B, T],\n",
-    "    \"labels\": Tensor[B, T],\n",
-    "    \"src_padding_mask\": BoolTensor[B, S],\n",
-    "    \"tgt_padding_mask\": BoolTensor[B, T],\n",
-    "}\n",
-    "```\n",
-    "\n",
-    "### Evaluation Interface (Person D -> C, E)\n",
-    "```python\n",
-    "evaluator = Evaluator(model, tokenizer, config)\n",
-    "results = evaluator.evaluate(dataloader)\n",
-    "# results: {\"bleu\": 25.6, \"comet\": 0.82, \"chrf\": 45.3, \"ter\": 55.2}\n",
-    "```"
-   ]
-  }
- ],
- "metadata": {
-  "colab": {
-   "include_colab_link": true,
-   "provenance": []
-  },
-  "kernelspec": {
-   "display_name": "Python 3",
-   "language": "python",
-   "name": "python3"
-  },
-  "language_info": {
-   "name": "python",
-   "version": "3.10.0"
-  }
- },
- "nbformat": 4,
- "nbformat_minor": 4
-}

requirements-colab.txt

@@ -1,21 +0,0 @@
-# =============================================================================
-# EasyTranslate — Google Colab Dependencies
-# =============================================================================
-# Colab already provides: torch, torchvision, numpy, pandas, matplotlib
-# This file lists only additional packages needed beyond Colab defaults.
-# =============================================================================
-
-transformers>=4.36.0,<4.45.0
-datasets>=2.16.0,<3.0.0
-tokenizers>=0.15.0,<0.20.0
-sentencepiece>=0.2.0,<0.3.0
-accelerate>=0.25.0,<0.35.0
-peft>=0.7.0,<0.12.0
-sacrebleu>=2.4.0,<3.0.0
-omegaconf>=2.3.0,<3.0.0
-rich>=13.0.0,<14.0.0
-tqdm>=4.66.0,<5.0.0
-wandb>=0.16.0,<0.18.0
-tensorboard>=2.15.0,<2.17.0
-protobuf>=5.28.3,<7.0.0
-fsspec==2025.3.0

requirements.txt

@@ -1,62 +1,35 @@
-# =============================================================================
-# EasyTranslate — Dependency Specification
-# =============================================================================
-# This file specifies all dependencies with compatible version ranges.
-# For exact reproducible environments, see requirements-lock.txt (pip freeze).
-# For Google Colab, see requirements-colab.txt (subset of packages).
-# =============================================================================
-
-# Core ML Framework
-torch>=2.1.0,<2.11.0
-torchvision>=0.16.0; platform_system != "Darwin" or platform_machine != "arm64"
-
-# HuggingFace Ecosystem
-transformers>=4.36.0,<4.45.0
-datasets>=2.16.0,<3.0.0
-tokenizers>=0.15.0,<0.20.0
-sentencepiece>=0.2.0,<0.3.0
-accelerate>=0.25.0,<0.35.0
-
-# Parameter-Efficient Fine-Tuning
-peft>=0.7.0,<0.12.0
-
-# Distributed Training (optional, GPU only)
-deepspeed>=0.12.0,<0.15.0; platform_system == "Linux"
-bitsandbytes>=0.41.0,<0.44.0; platform_system == "Linux"
-
-# Experiment Tracking
-wandb>=0.16.0,<0.18.0
-tensorboard>=2.15.0,<2.17.0
-
-# Evaluation Metrics
-sacrebleu>=2.4.0,<3.0.0
-unbabel-comet>=2.2.0,<2.3.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,<3.9.0
-
-# Configuration & CLI
-omegaconf>=2.3.0,<3.0.0
-pyyaml>=6.0.0,<7.0.0
-rich>=13.0.0,<14.0.0
-
-# Data Processing
-numpy>=1.24.0,<3.0.0
-pandas>=2.0.0,<2.3.0
-protobuf>=5.28.3,<7.0.0
-tqdm>=4.66.0,<5.0.0
-
-# Visualization
-matplotlib>=3.8.0,<4.0.0
-seaborn>=0.13.0,<0.14.0
+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,<0.115.0
-uvicorn>=0.25.0,<0.31.0
-gradio>=4.10.0,<5.0.0
-
-# Development & Testing
-pytest>=7.4.0,<9.0.0
-pytest-cov>=4.1.0
-black>=23.0.0,<25.0.0
-isort>=5.12.0,<6.0.0
-flake8>=6.0.0,<8.0.0
+fastapi>=0.108.0
+uvicorn>=0.25.0
+gradio>=4.10.0

result/best_model.pt

@@ -1,3 +0,0 @@
-version https://git-lfs.github.com/spec/v1
-oid sha256:c80b4b011e062695fceae4bdf1d43c070ed4ceace176c4637fcde70cd54d54d4
-size 1120219211

result/evaluation_results.json

@@ -1,10 +0,0 @@
-{
-  "bleu": 15.1403,
-  "bleu_1": 43.0105,
-  "bleu_2": 20.1492,
-  "bleu_3": 10.4122,
-  "bleu_4": 5.8232,
-  "bp": 1.0,
-  "chrf": 15.2513,
-  "ter": 177.5966
-}

result/link.txt

@@ -1 +0,0 @@
-https://colab.research.google.com/drive/1wCVTnNuULJkBIS5tzPu5uQY4zd-akEd4?authuser=1#scrollTo=SssPN0dpTHJX&fullscreenOutput=true

result/train_output.txt

@@ -1,102 +0,0 @@
- ============================================================
-  Starting Training...
-============================================================
-[05/13/26 07:55:00] INFO     Starting training for 30 epochs                                         trainer.py:193
-
-                    INFO     Device: cuda, FP16: False, BF16: True                                   trainer.py:194
-
-                    INFO     Gradient accumulation steps: 4                                          trainer.py:195
-
-                    INFO     ==================================================                      trainer.py:199
-
-                    INFO     Epoch 1/30                                                              trainer.py:200
-
-[05/13/26 07:57:11] INFO     Epoch 1 - Train Loss: 8.2498                                            trainer.py:279
-
-Evaluating: 100%|██████████| 122/122 [17:49<00:00,  8.76s/it]
-[05/13/26 08:15:01] WARNING  Evaluation failed during validation: references cannot be empty         trainer.py:311
-
-                    INFO     Epoch 1 - Val Loss: 7.9030                                              trainer.py:313
-
-[05/13/26 08:15:02] INFO     Checkpoint saved: checkpoints/checkpoint_epoch_1.pt                     trainer.py:341
-
-[05/13/26 08:15:03] INFO     New best model saved: checkpoints/best_model.pt (metric: 7.9030)        trainer.py:352
-
-                    INFO     No improvement for 1 epochs (best: 7.9030, current: 7.9030)             trainer.py:398
-
-                    INFO     ==================================================                      trainer.py:199
-
-                    INFO     Epoch 2/30                                                              trainer.py:200
-
-[05/13/26 08:17:12] INFO     Epoch 2 - Train Loss: 6.6109                                            trainer.py:279
-
-Evaluating: 100%|██████████| 122/122 [06:10<00:00,  3.04s/it]
-[05/13/26 08:23:23] WARNING  Evaluation failed during validation: references cannot be empty         trainer.py:311
-
-                    INFO     Epoch 2 - Val Loss: 6.9446                                              trainer.py:313
-
-[05/13/26 08:23:24] INFO     Checkpoint saved: checkpoints/checkpoint_epoch_2.pt                     trainer.py:341
-
-                    INFO     No improvement for 2 epochs (best: 7.9030, current: 6.9446)             trainer.py:398
-
-                    INFO     ==================================================                      trainer.py:199
-
-                    INFO     Epoch 3/30                                                              trainer.py:200
-
-[05/13/26 08:25:33] INFO     Epoch 3 - Train Loss: 5.2802                                            trainer.py:279
-
-Evaluating: 100%|██████████| 122/122 [06:55<00:00,  3.40s/it]
-[05/13/26 08:32:29] WARNING  Evaluation failed during validation: references cannot be empty         trainer.py:311
-
-                    INFO     Epoch 3 - Val Loss: 6.0977                                              trainer.py:313
-
-                    INFO     Checkpoint saved: checkpoints/checkpoint_epoch_3.pt                     trainer.py:341
-
-                    INFO     No improvement for 3 epochs (best: 7.9030, current: 6.0977)             trainer.py:398
-
-                    INFO     ==================================================                      trainer.py:199
-
-                    INFO     Epoch 4/30                                                              trainer.py:200
-
-Train Epoch 4:  28%|██▊       | 1687/6104 [00:35<01:30, 48.68it/s, loss=4.5668, lr=3.00e-04, step=5000][05/13/26 08:33:06] INFO     Checkpoint saved: checkpoints/checkpoint_step_5000.pt                   trainer.py:341
-
-                    INFO     New best model saved: checkpoints/best_model.pt (metric: inf)           trainer.py:352
-
-[05/13/26 08:34:39] INFO     Epoch 4 - Train Loss: 4.4749                                            trainer.py:279
-
-Evaluating: 100%|██████████| 122/122 [03:08<00:00,  1.55s/it]
-[05/13/26 08:37:49] WARNING  Evaluation failed during validation: references cannot be empty         trainer.py:311
-
-                    INFO     Epoch 4 - Val Loss: 5.6509                                              trainer.py:313
-
-[05/13/26 08:37:50] INFO     Checkpoint saved: checkpoints/checkpoint_epoch_4.pt                     trainer.py:341
-
-                    INFO     No improvement for 4 epochs (best: inf, current: 5.6509)                trainer.py:398
-
-                    INFO     ==================================================                      trainer.py:199
-
-                    INFO     Epoch 5/30                                                              trainer.py:200
-
-[05/13/26 08:40:00] INFO     Epoch 5 - Train Loss: 4.0513                                            trainer.py:279
-
-Evaluating: 100%|██████████| 122/122 [03:28<00:00,  1.71s/it]
-[05/13/26 08:43:29] WARNING  Evaluation failed during validation: references cannot be empty         trainer.py:311
-
-                    INFO     Epoch 5 - Val Loss: 5.4433                                              trainer.py:313
-
-                    INFO     Checkpoint saved: checkpoints/checkpoint_epoch_5.pt                     trainer.py:341
-
-                    INFO     No improvement for 5 epochs (best: inf, current: 5.4433)                trainer.py:398
-
-                    INFO     Early stopping triggered at epoch 5                                     trainer.py:212
-
-                    INFO     ==================================================                      trainer.py:431
-
-                    INFO     Training completed!                                                     trainer.py:432
-
-                    INFO     Best epoch: 4                                                           trainer.py:433
-
-                    INFO     Best bleu: inf                                                          trainer.py:434
-
-                    INFO     Training summary saved to checkpoints/training_summary.json             trainer.py:447
-

result/training_summary.json

@@ -1,30 +0,0 @@
-{
-  "best_epoch": 3,
-  "best_metric": Infinity,
-  "metric_name": "bleu",
-  "total_steps": 7630,
-  "train_loss_history": [
-    8.249846095573231,
-    6.610892885478108,
-    5.28016960105296,
-    4.47492950561006,
-    4.051334980548131
-  ],
-  "val_metrics_history": [
-    {
-      "val_loss": 7.903
-    },
-    {
-      "val_loss": 6.9446
-    },
-    {
-      "val_loss": 6.0977
-    },
-    {
-      "val_loss": 5.6509
-    },
-    {
-      "val_loss": 5.4433
-    }
-  ]
-}

result/translation_examples.json

@@ -1,18 +0,0 @@
-[
-  {
-    "source_en": "Hello, how are you today?",
-    "target_zh": "今天,你如何?"
-  },
-  {
-    "source_en": "Machine translation is an important field of natural language processing.",
-    "target_zh": "生物交流是自然语言处理的重要领域。"
-  },
-  {
-    "source_en": "The weather is beautiful and I want to go for a walk.",
-    "target_zh": "天气是美丽的,我想去。"
-  },
-  {
-    "source_en": "Deep learning has revolutionized artificial intelligence research.",
-    "target_zh": "深度学习推动了人工智能情报研究。"
-  }
-]

scripts/evaluate.py

@@ -10,17 +10,11 @@
 """
 
 import argparse
-import json
 import sys
 from pathlib import Path
 
-import torch
-import yaml
-
 sys.path.insert(0, str(Path(__file__).resolve().parent.parent / "src"))
 
-from easytranslate.evaluation.evaluator import Evaluator
-
 
 def parse_args():
     parser = argparse.ArgumentParser(description="EasyTranslate Evaluation")
@@ -49,71 +43,7 @@ def main():
     print("  EasyTranslate - Evaluation")
     print("=" * 60)
 
-    # 1. 加载配置
-    with open(args.config, "r", encoding="utf-8") as f:
-        config = yaml.safe_load(f)
-
-    # 应用命令行覆盖 (格式: key.subkey=value)
-    for override in cli_overrides:
-        if "=" in override:
-            key, value = override.split("=", 1)
-            keys = key.split(".")
-            d = config
-            for k in keys[:-1]:
-                d = d.setdefault(k, {})
-            d[keys[-1]] = yaml.safe_load(value)
-
-    # 2. 加载检查点并重建模型
-    checkpoint = torch.load(args.checkpoint, map_location="cpu")
-    model_config = config.get("model", {})
-
-    from easytranslate.model.transformer import Transformer
-    model = Transformer(model_config)
-    model.load_state_dict(checkpoint["model_state_dict"])
-
-    device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
-    model = model.to(device)
-    model.eval()
-
-    # 3. 加载 tokenizer 和测试数据
-    from easytranslate.data.tokenizer import build_tokenizer
-    from easytranslate.data.dataset import TranslationDataset
-    from torch.utils.data import DataLoader
-
-    tokenizer = build_tokenizer(config.get("data", {}).get("tokenizer", {}))
-
-    test_dataset = TranslationDataset(
-        config=config.get("data", {}),
-        tokenizer=tokenizer,
-        split="test",
-    )
-    test_loader = DataLoader(
-        test_dataset,
-        batch_size=config.get("evaluation", {}).get("batch_size", 32),
-        shuffle=False,
-    )
-
-    # 4. 构建 Evaluator
-    evaluator = Evaluator(model=model, tokenizer=tokenizer, config=config)
-
-    # 5. 运行评估
-    src_texts = [sample["src"] for sample in test_dataset.raw_data]
-    ref_texts = [sample["tgt"] for sample in test_dataset.raw_data]
-    results = evaluator.evaluate(test_loader, src_texts=src_texts, ref_texts=ref_texts)
-
-    # 6. 打印和保存结果
-    print("\n" + "=" * 60)
-    print("  Evaluation Results")
-    print("=" * 60)
-    for metric, score in results.items():
-        if not isinstance(score, list):
-            print(f"  {metric:>10s}: {score:.4f}")
-
-    output_path = Path(args.output)
-    output_path.parent.mkdir(parents=True, exist_ok=True)
-    with open(output_path, "w", encoding="utf-8") as f:
-        json.dump(results, f, indent=2, ensure_ascii=False)
-    print(f"\n  Results saved to: {output_path}")
+    raise NotImplementedError("TODO: Person D 实现评估主流程")
 
 
 if __name__ == "__main__":

scripts/setup_colab.py

@@ -1,121 +0,0 @@
-#!/usr/bin/env python3
-"""
-Google Colab Environment Setup Script for EasyTranslate.
-
-This script is designed to be run as the first cell in a Colab notebook.
-It handles:
-1. GPU verification and CUDA setup
-2. Repository cloning from GitHub
-3. Dependency installation
-4. Google Drive mounting
-5. Environment variable configuration
-
-Usage in Colab:
-    !wget -q https://raw.githubusercontent.com/your-org/UCAS-EasyTranslate/main/scripts/setup_colab.py
-    %run setup_colab.py
-"""
-
-import os
-import subprocess
-import sys
-from pathlib import Path
-
-REPO_URL = "https://github.com/your-org/UCAS-EasyTranslate.git"
-REPO_DIR = Path("/content/UCAS-EasyTranslate")
-DRIVE_MOUNT_POINT = "/content/drive"
-DRIVE_BASE = "/content/drive/MyDrive/EasyTranslate"
-
-
-def check_gpu():
-    try:
-        import torch
-        if torch.cuda.is_available():
-            print(f"[OK] GPU detected: {torch.cuda.get_device_name(0)}")
-            print(f"     CUDA version: {torch.version.cuda}")
-            print(f"     GPU memory: {torch.cuda.get_device_properties(0).total_memory / 1024**3:.1f} GB")
-            return True
-        else:
-            print("[WARN] No GPU detected. Training will be very slow.")
-            return False
-    except ImportError:
-        print("[WARN] PyTorch not found. Installing...")
-        subprocess.check_call([sys.executable, "-m", "pip", "install", "-q", "torch"])
-        return check_gpu()
-
-
-def clone_repository():
-    if REPO_DIR.exists():
-        print(f"[INFO] Repository exists at {REPO_DIR}, pulling latest...")
-        os.chdir(REPO_DIR)
-        subprocess.check_call(["git", "pull", "origin", "main"])
-    else:
-        print(f"[INFO] Cloning repository from {REPO_URL}...")
-        subprocess.check_call(["git", "clone", REPO_URL, str(REPO_DIR)])
-    os.chdir(REPO_DIR)
-    sys.path.insert(0, str(REPO_DIR / "src"))
-    print(f"[OK] Repository ready at {REPO_DIR}")
-
-
-def install_dependencies():
-    req_file = REPO_DIR / "requirements-colab.txt"
-    if req_file.exists():
-        print("[INFO] Installing Colab dependencies...")
-        subprocess.check_call([sys.executable, "-m", "pip", "install", "-q", "-r", str(req_file)])
-    else:
-        print("[INFO] Installing from requirements.txt...")
-        subprocess.check_call([sys.executable, "-m", "pip", "install", "-q", "-r", str(REPO_DIR / "requirements.txt")])
-
-    print("[INFO] Installing EasyTranslate package...")
-    subprocess.check_call([sys.executable, "-m", "pip", "install", "-q", "-e", str(REPO_DIR)])
-    print("[OK] All dependencies installed")
-
-
-def mount_drive():
-    try:
-        from google.colab import drive
-        drive.mount(DRIVE_MOUNT_POINT)
-        if os.path.exists(DRIVE_MOUNT_POINT):
-            os.makedirs(DRIVE_BASE, exist_ok=True)
-            print(f"[OK] Google Drive mounted at {DRIVE_MOUNT_POINT}")
-            print(f"     Base path: {DRIVE_BASE}")
-            return True
-    except Exception as e:
-        print(f"[WARN] Google Drive mount failed: {e}")
-    return False
-
-
-def setup_environment():
-    os.environ["TOKENIZERS_PARALLELISM"] = "false"
-    os.environ["WANDB_MODE"] = os.environ.get("WANDB_MODE", "offline")
-    os.environ["PYTHONHASHSEED"] = "42"
-    print("[OK] Environment variables configured")
-
-
-def main():
-    print("=" * 60)
-    print("  EasyTranslate — Colab Environment Setup")
-    print("=" * 60)
-    print()
-
-    check_gpu()
-    print()
-
-    clone_repository()
-    print()
-
-    install_dependencies()
-    print()
-
-    mount_drive()
-    print()
-
-    setup_environment()
-    print()
-
-    print("=" * 60)
-    print("  Setup Complete! Ready for training.")
-    print("=" * 60)
-
-
-if __name__ == "__main__":
-    main()

scripts/translate.py

@@ -16,13 +16,8 @@ import argparse
 import sys
 from pathlib import Path
 
-import torch
-import yaml
-
 sys.path.insert(0, str(Path(__file__).resolve().parent.parent / "src"))
 
-from easytranslate.evaluation.evaluator import Evaluator
-
 
 def parse_args():
     parser = argparse.ArgumentParser(description="EasyTranslate Inference")
@@ -44,25 +39,7 @@ def interactive_translate(evaluator):
     3. 打印翻译结果
     4. 输入 'quit' 退出
     """
-    print("\nInteractive Translation Mode (type 'quit' to exit)")
-    print("-" * 40)
-
-    while True:
-        try:
-            text = input("\n[EN] > ").strip()
-        except (EOFError, KeyboardInterrupt):
-            print("\nBye!")
-            break
-
-        if text.lower() in ("quit", "exit", "q"):
-            print("Bye!")
-            break
-
-        if not text:
-            continue
-
-        translation = evaluator.translate_single(text)
-        print(f"[ZH] > {translation}")
+    raise NotImplementedError("TODO: Person D 实现 interactive_translate")
 
 
 def translate_file(evaluator, input_path: str, output_path: str):
@@ -74,33 +51,7 @@ def translate_file(evaluator, input_path: str, output_path: str):
     2. 批量翻译
     3. 将结果写入输出文件
     """
-    input_file = Path(input_path)
-    if not input_file.exists():
-        print(f"Error: input file not found: {input_path}")
-        return
-
-    with open(input_file, "r", encoding="utf-8") as f:
-        lines = [line.strip() for line in f if line.strip()]
-
-    print(f"Translating {len(lines)} sentences...")
-
-    # 分批翻译
-    batch_size = 32
-    translations = []
-    for i in range(0, len(lines), batch_size):
-        batch = lines[i : i + batch_size]
-        batch_translations = evaluator.translate(batch)
-        translations.extend(batch_translations)
-        print(f"  Translated {min(i + batch_size, len(lines))}/{len(lines)}")
-
-    # 写入输出文件
-    out_file = Path(output_path)
-    out_file.parent.mkdir(parents=True, exist_ok=True)
-    with open(out_file, "w", encoding="utf-8") as f:
-        for t in translations:
-            f.write(t + "\n")
-
-    print(f"Results saved to: {output_path}")
+    raise NotImplementedError("TODO: Person D 实现 translate_file")
 
 
 def launch_web_ui(evaluator):
@@ -113,25 +64,7 @@ def launch_web_ui(evaluator):
     3. 输出: 中文翻译结果
     4. 调用 evaluator.translate_single()
     """
-    try:
-        import gradio as gr
-    except ImportError:
-        print("Error: gradio is not installed. Install with: pip install gradio")
-        return
-
-    def translate_fn(text):
-        if not text.strip():
-            return ""
-        return evaluator.translate_single(text)
-
-    interface = gr.Interface(
-        fn=translate_fn,
-        inputs=gr.Textbox(label="English", placeholder="Enter English text..."),
-        outputs=gr.Textbox(label="Chinese Translation"),
-        title="EasyTranslate - English to Chinese",
-        description="Transformer-based English to Chinese translation system.",
-    )
-    interface.launch()
+    raise NotImplementedError("TODO: Person D 实现 launch_web_ui")
 
 
 def main():
@@ -141,37 +74,8 @@ def main():
     print("  EasyTranslate - Translation")
     print("=" * 60)
 
-    # 加载配置
-    with open(args.config, "r", encoding="utf-8") as f:
-        config = yaml.safe_load(f)
-
-    # 加载模型
-    checkpoint = torch.load(args.checkpoint, map_location="cpu")
-
-    from easytranslate.model.transformer import Transformer
-    from easytranslate.data.tokenizer import build_tokenizer
-
-    model_config = config.get("model", {})
-    model = Transformer(model_config)
-    model.load_state_dict(checkpoint["model_state_dict"])
-
-    device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
-    model = model.to(device)
-    model.eval()
-
-    tokenizer = build_tokenizer(config.get("data", {}).get("tokenizer", {}))
-
-    # 构建 evaluator
-    evaluator = Evaluator(model=model, tokenizer=tokenizer, config=config)
-
-    # 根据参数选择模式
-    if args.web:
-        launch_web_ui(evaluator)
-    elif args.input:
-        output_path = args.output or args.input.replace(".txt", "_translated.txt")
-        translate_file(evaluator, args.input, output_path)
-    else:
-        interactive_translate(evaluator)
+    # TODO: 加载模型、构建 evaluator，然后根据参数选择模式
+    raise NotImplementedError("TODO: 实现推理主流程")
 
 
 if __name__ == "__main__":

setup.py

@@ -1,78 +1,23 @@
 from setuptools import setup, find_packages
 
-with open("README.md", "r", encoding="utf-8") as fh:
-    long_description = fh.read()
-
 setup(
     name="easytranslate",
     version="0.1.0",
     description="Transformer-based English-to-Chinese Translation Model",
-    long_description=long_description,
-    long_description_content_type="text/markdown",
-    author="UCAS EasyTranslate Team",
     packages=find_packages(where="src"),
     package_dir={"": "src"},
     python_requires=">=3.10",
     install_requires=[
-        "torch>=2.1.0,<2.11.0",
-        "transformers>=4.36.0,<4.45.0",
-        "datasets>=2.16.0,<3.0.0",
-        "tokenizers>=0.15.0,<0.20.0",
-        "sentencepiece>=0.2.0,<0.3.0",
-        "accelerate>=0.25.0,<0.35.0",
-        "peft>=0.7.0,<0.12.0",
-        "sacrebleu>=2.4.0,<3.0.0",
-        "unbabel-comet>=2.2.0,<2.3.0",
-        "numpy>=1.24.0,<3.0.0",
-        "protobuf>=5.28.3,<7.0.0",
-        "omegaconf>=2.3.0,<3.0.0",
-        "rich>=13.0.0,<14.0.0",
-        "tqdm>=4.66.0,<5.0.0",
-        "pyyaml>=6.0.0,<7.0.0",
-    ],
-    extras_require={
-        "dev": [
-            "pytest>=7.4.0,<9.0.0",
-            "pytest-cov>=4.1.0",
-            "black>=23.0.0,<25.0.0",
-            "isort>=5.12.0,<6.0.0",
-            "flake8>=6.0.0,<8.0.0",
-        ],
-        "colab": [
-            "transformers>=4.36.0,<4.45.0",
-            "datasets>=2.16.0,<3.0.0",
-            "tokenizers>=0.15.0,<0.20.0",
-            "sentencepiece>=0.2.0,<0.3.0",
-            "accelerate>=0.25.0,<0.35.0",
-            "peft>=0.7.0,<0.12.0",
-            "sacrebleu>=2.4.0,<3.0.0",
-            "unbabel-comet>=2.2.0,<2.3.0",
-            "protobuf>=5.28.3,<7.0.0",
-            "omegaconf>=2.3.0,<3.0.0",
-            "rich>=13.0.0,<14.0.0",
-            "tqdm>=4.66.0,<5.0.0",
-        ],
-        "distributed": [
-            "deepspeed>=0.12.0,<0.15.0",
-            "bitsandbytes>=0.41.0,<0.44.0",
-        ],
-        "tracking": [
-            "wandb>=0.16.0,<0.18.0",
-            "tensorboard>=2.15.0,<2.17.0",
-        ],
-        "serving": [
-            "fastapi>=0.108.0,<0.115.0",
-            "uvicorn>=0.25.0,<0.31.0",
-            "gradio>=4.10.0,<5.0.0",
-        ],
-    },
-    classifiers=[
-        "Development Status :: 3 - Alpha",
-        "Intended Audience :: Science/Research",
-        "License :: OSI Approved :: MIT License",
-        "Programming Language :: Python :: 3.10",
-        "Programming Language :: Python :: 3.11",
-        "Programming Language :: Python :: 3.12",
-        "Topic :: Scientific/Engineering :: Artificial Intelligence",
+        "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/data/collator.py

@@ -64,19 +64,6 @@ class DynamicBatchSampler(Sampler[list[int]]):
         self.max_tokens_per_batch = max_tokens_per_batch
         self.shuffle = shuffle
         self.drop_last = drop_last
-        
-        # 检查是否有超过最大token限制的序列
-        self._check_long_sequences()
-
-    def _check_long_sequences(self):
-        """检查并警告过长的序列"""
-        long_seq_count = sum(1 for l in self.lengths if l > self.max_tokens_per_batch)
-        if long_seq_count > 0:
-            import warnings
-            warnings.warn(
-                f"Found {long_seq_count} sequences longer than max_tokens_per_batch "
-                f"({self.max_tokens_per_batch}). These will be placed in their own batches."
-            )
 
     def __iter__(self) -> Iterator[list[int]]:
         indices = list(range(len(self.lengths)))
@@ -89,18 +76,7 @@ class DynamicBatchSampler(Sampler[list[int]]):
         max_len = 0
 
         for idx in indices:
-            seq_len = self.lengths[idx]
-            
-            # 如果单个序列长度超过限制，单独放入一个批次
-            if seq_len > self.max_tokens_per_batch:
-                if batch:
-                    batches.append(batch)
-                    batch = []
-                    max_len = 0
-                batches.append([idx])
-                continue
-            
-            candidate_max_len = max(max_len, seq_len)
+            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:
@@ -109,7 +85,7 @@ class DynamicBatchSampler(Sampler[list[int]]):
                 max_len = 0
 
             batch.append(idx)
-            max_len = max(max_len, seq_len)
+            max_len = max(max_len, self.lengths[idx])
 
         if batch and not self.drop_last:
             batches.append(batch)
@@ -125,11 +101,6 @@ class DynamicBatchSampler(Sampler[list[int]]):
         max_len = 0
 
         for length in sorted(self.lengths):
-            # 过长的序列单独计数
-            if length > self.max_tokens_per_batch:
-                count += 1
-                continue
-            
             candidate_max_len = max(max_len, length)
             if batch_size and candidate_max_len * (batch_size + 1) > self.max_tokens_per_batch:
                 count += 1
@@ -140,4 +111,4 @@ class DynamicBatchSampler(Sampler[list[int]]):
 
         if batch_size and not self.drop_last:
             count += 1
-        return count
+        return count

src/easytranslate/evaluation/decoding.py

@@ -34,23 +34,20 @@ def greedy_decode(
     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
-
-    decoder_input = 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(decoder_input, encoder_output, src_padding_mask)
-        next_token = logits.argmax(dim=-1, keepdim=True)
-        decoder_input = torch.cat([decoder_input, next_token], dim=1)
-        finished = finished | next_token.squeeze(-1).eq(eos_id)
-        if finished.all():
-            break
-
-    return decoder_input
+    """
+    贪心解码。
+
+    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()
@@ -67,104 +64,29 @@ def beam_search_decode(
 ) -> torch.Tensor:
     """
     束搜索解码。
-    
-    优化改进:
-    1. 限制最大 beam_size 为 64，防止内存溢出
-    2. 使用更高效的索引操作避免中间 tensor 累积
-    3. 及时释放不再需要的中间结果
-    4. 对长序列使用更小的 beam_size
+
+    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]
+
+    这是翻译任务最关键的解码算法，请仔细实现。
     """
-    # 安全检查：限制 beam_size 防止内存问题
-    beam_size = min(beam_size, 64)
-    
-    batch_size, seq_len = src_ids.size()
-    device = src_ids.device
-    
-    # 根据序列长度动态调整 beam_size
-    if seq_len > 512:
-        beam_size = min(beam_size, 3)
-    
-    # 1. Encode source
-    encoder_output = model.encode(src_ids, src_padding_mask)
-    hidden_dim = encoder_output.size(-1)
-    
-    # 2. Expand encoder output for beam search
-    encoder_output = encoder_output.unsqueeze(1).expand(-1, beam_size, -1, -1)
-    encoder_output = encoder_output.reshape(batch_size * beam_size, seq_len, hidden_dim)
-    src_padding_mask_expanded = src_padding_mask.unsqueeze(1).expand(-1, beam_size, -1)
-    src_padding_mask_expanded = src_padding_mask_expanded.reshape(batch_size * beam_size, seq_len)
-
-    # 3. Initialize beams - 使用预分配的 tensor 减少内存分配
-    beam_scores = torch.zeros(batch_size, beam_size, device=device)
-    beam_scores[:, 1:] = float("-inf")
-    
-    # 预分配最大长度的 tensor，动态填充
-    beam_tokens = torch.full((batch_size, beam_size, max_len + 1), eos_id, dtype=torch.long, device=device)
-    beam_tokens[:, :, 0] = bos_id  # [B, beam, 1]
-    
-    finished = torch.zeros(batch_size, beam_size, dtype=torch.bool, device=device)
-    lengths = torch.ones(batch_size, beam_size, dtype=torch.long, device=device)
-
-    # 4. Iterative decoding
-    for step in range(1, max_len + 1):
-        # 只处理未完成的 beam
-        active_mask = ~finished
-        if not active_mask.any():
-            break
-        
-        # 获取当前步的输入 [B*beam, step]
-        flat_tokens = beam_tokens[:, :, :step].reshape(batch_size * beam_size, step)
-        
-        # 前向解码
-        logits = model.decode_step(flat_tokens, encoder_output, src_padding_mask_expanded)
-        log_probs = F.log_softmax(logits, dim=-1)
-        vocab_size = log_probs.size(-1)
-        
-        # 应用 no_repeat_ngram 约束
-        if no_repeat_ngram_size > 0 and step >= no_repeat_ngram_size:
-            log_probs = _apply_no_repeat_ngram(log_probs, flat_tokens, no_repeat_ngram_size)
-
-        # Mask finished beams - 禁止生成新 token，但允许保持 EOS
-        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
-
-        # 计算分数
-        scores = beam_scores.unsqueeze(-1) + log_probs.view(batch_size, beam_size, vocab_size)
-        scores = scores.view(batch_size, -1)
-
-        # 选择 top-k
-        topk_scores, topk_indices = scores.topk(beam_size, dim=-1)
-        beam_indices = topk_indices // vocab_size
-        token_indices = topk_indices % vocab_size
-
-        # 高效更新 beam_tokens - 使用索引而非拼接
-        new_beam_tokens = beam_tokens.clone()
-        for b in range(batch_size):
-            new_beam_tokens[b] = beam_tokens[b][beam_indices[b]]
-            new_beam_tokens[b, torch.arange(beam_size), step] = token_indices[b]
-        beam_tokens = new_beam_tokens
-        
-        # 更新 finished 和 lengths
-        new_finished = finished.gather(1, beam_indices) | token_indices.eq(eos_id)
-        new_lengths = lengths.gather(1, beam_indices)
-        new_lengths[~new_finished] = step + 1
-        
-        beam_scores = topk_scores
-        finished = new_finished
-        lengths = new_lengths
-
-    # 5. Apply length penalty
-    lengths = lengths.float()
-    penalties = lengths ** length_penalty
-    final_scores = beam_scores / penalties
-
-    # 6. Select best beam for each batch
-    best_indices = final_scores.argmax(dim=-1)
-    best_sequences = beam_tokens[torch.arange(batch_size, device=device), best_indices]
-
-    return best_sequences
+    raise NotImplementedError("TODO: Person D 实现 beam_search_decode")
 
 
 @torch.no_grad()
@@ -179,46 +101,17 @@ def sample_decode(
     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
-
-    decoder_input = 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(decoder_input, encoder_output, src_padding_mask)
-
-        logits = logits / max(temperature, 1e-8)
-
-        if top_k > 0:
-            k = min(top_k, logits.size(-1))
-            topk_values, _ = torch.topk(logits, k, dim=-1)
-            threshold = topk_values[:, -1].unsqueeze(-1)
-            logits[logits < threshold] = float("-inf")
-
-        if 0.0 < top_p < 1.0:
-            sorted_logits, sorted_indices = torch.sort(logits, descending=True, dim=-1)
-            cumulative_probs = torch.cumsum(F.softmax(sorted_logits, dim=-1), dim=-1)
-            # Remove tokens whose cumulative probability exceeds top_p
-            # (shift by one so the token that pushes above the threshold is kept)
-            remove_mask = cumulative_probs - F.softmax(sorted_logits, dim=-1) >= top_p
-            sorted_logits[remove_mask] = float("-inf")
-            # Scatter sorted values back to their original vocabulary positions
-            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 = torch.where(finished.unsqueeze(-1), torch.full_like(next_token, eos_id), next_token)
-        decoder_input = torch.cat([decoder_input, next_token], dim=1)
-        finished = finished | next_token.squeeze(-1).eq(eos_id)
-
-        if finished.all():
-            break
-
-    return decoder_input
+    """
+    采样解码 (支持 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(
@@ -226,28 +119,11 @@ def _apply_no_repeat_ngram(
     generated_tokens: torch.Tensor,
     ngram_size: int,
 ) -> torch.Tensor:
-    """防止生成重复的 n-gram。"""
-    if ngram_size <= 0:
-        return logits
-
-    batch_size = logits.size(0)
-    seq_len = generated_tokens.size(1)
-    
-    if seq_len < ngram_size - 1:
-        return logits
-
-    for batch_idx in range(batch_size):
-        tokens = generated_tokens[batch_idx].tolist()
-
-        ngram_map: dict[tuple, set] = {}
-        for i in range(len(tokens) - ngram_size + 1):
-            prefix = tuple(tokens[i : i + ngram_size - 1])
-            next_tok = tokens[i + ngram_size - 1]
-            ngram_map.setdefault(prefix, set()).add(next_tok)
-
-        current_prefix = tuple(tokens[-(ngram_size - 1):])
-        if current_prefix in ngram_map:
-            banned = list(ngram_map[current_prefix])
-            logits[batch_idx, banned] = float("-inf")
-
-    return logits
+    """
+    防止生成重复�� 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")

src/easytranslate/evaluation/evaluator.py

@@ -26,82 +26,20 @@ logger = logging.getLogger(__name__)
 
 
 class Evaluator:
-    """翻译模型评估器。"""
+    """
+    翻译模型评估器。
+
+    TODO [Person D]: 实现以下方法。
+    """
 
     def __init__(self, model: nn.Module, tokenizer, config: dict):
-        """初始化评估器。"""
-        self.model = model
-        self.tokenizer = tokenizer
-        self.config = config
-
-        eval_config = config.get("evaluation", {})
-        decoding_config = eval_config.get("decoding", {})
-
-        self.strategy = decoding_config.get("strategy", "beam_search")
-        self.max_len = decoding_config.get("max_decode_len", 256)
-        self.beam_size = decoding_config.get("beam_size", 5)
-        self.length_penalty = decoding_config.get("length_penalty", 1.0)
-        self.no_repeat_ngram_size = decoding_config.get("no_repeat_ngram_size", 0)
-
-        sampling_config = decoding_config.get("sampling", {})
-        self.temperature = sampling_config.get("temperature", 1.0)
-        self.top_k = sampling_config.get("top_k", 0)
-        self.top_p = sampling_config.get("top_p", 1.0)
-
-        self.metrics = eval_config.get("metrics", ["bleu", "comet", "chrf", "ter"])
-
-        self.bos_id = tokenizer.bos_token_id
-        self.eos_id = tokenizer.eos_token_id
-        self.pad_id = tokenizer.pad_token_id
-
-        self.decode_fn = self._get_decode_fn()
-
-    def _get_decode_fn(self):
-        """根据策略选择解码函数。"""
-        if self.strategy == "greedy":
-            return self._greedy
-        elif self.strategy == "sampling":
-            return self._sample
-        else:
-            return self._beam_search
-
-    def _greedy(self, src_ids, src_padding_mask):
-        return greedy_decode(
-            self.model, src_ids, src_padding_mask,
-            self.bos_id, self.eos_id, max_len=self.max_len,
-        )
-
-    def _beam_search(self, src_ids, src_padding_mask):
-        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_len,
-            length_penalty=self.length_penalty,
-            no_repeat_ngram_size=self.no_repeat_ngram_size,
-        )
-
-    def _sample(self, src_ids, src_padding_mask):
-        return sample_decode(
-            self.model, src_ids, src_padding_mask,
-            self.bos_id, self.eos_id,
-            max_len=self.max_len, temperature=self.temperature,
-            top_k=self.top_k, top_p=self.top_p,
-        )
-
-    def _validate_batch(self, batch: dict) -> bool:
-        """验证批处理格式是否符合预期。"""
-        required_keys = ["src_ids"]
-        
-        for key in required_keys:
-            if key not in batch:
-                logger.error(f"Batch is missing required key: {key}")
-                return False
-        
-        if not isinstance(batch["src_ids"], torch.Tensor):
-            logger.error(f"src_ids must be a torch.Tensor, got {type(batch['src_ids'])}")
-            return False
-        
-        return True
+        """
+        TODO [Person D]:
+        1. 保存 model, tokenizer, config
+        2. 从 config 读取解码策略和评估指标配置
+        3. 根据策略选择解码函数
+        """
+        raise NotImplementedError("TODO: Person D 实现 Evaluator.__init__")
 
     def evaluate(
         self,
@@ -111,74 +49,28 @@ class Evaluator:
     ) -> dict:
         """
         在给定数据上进行评估。
-        
-        Args:
-            dataloader: 测试数据集的 DataLoader
-            src_texts: 源文本列表 (用于某些指标计算)
-            ref_texts: 参考译文列表
-        
-        Returns:
-            dict: 包含所有评估指标的字典
-        
-        Raises:
-            ValueError: 如果批处理格式不正确
+
+        TODO [Person D]: 实现以下逻辑:
+        1. model.eval()
+        2. 遍历 dataloader，使用选定的解码策略生成翻译
+        3. 将生成的 token ids 解码为文本
+        4. 调用 compute_all_metrics 计算指标
+        5. 返回评估结果 dict
         """
-        self.model.eval()
-        device = next(self.model.parameters()).device
-
-        hypotheses = []
-
-        with torch.no_grad():
-            for batch in tqdm(dataloader, desc="Evaluating"):
-                # 验证批处理格式
-                if not self._validate_batch(batch):
-                    raise ValueError("Invalid batch format. Expected 'src_ids' key with torch.Tensor value.")
-                
-                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_fn(src_ids, src_padding_mask)
-
-                for i in range(output_ids.size(0)):
-                    text = self.tokenizer.decode(
-                        output_ids[i].tolist(), skip_special_tokens=True
-                    )
-                    hypotheses.append(text)
-
-        results = compute_all_metrics(
-            sources=src_texts,
-            hypotheses=hypotheses,
-            references=ref_texts,
-            metrics=self.metrics,
-        )
-        return results
+        raise NotImplementedError("TODO: Person D 实现 evaluate")
 
     def translate(self, texts: list[str]) -> list[str]:
-        """翻译一批文本。"""
-        self.model.eval()
-        device = next(self.model.parameters()).device
-
-        encoded = [self.tokenizer.encode(t, add_special_tokens=True) for t in texts]
-        max_len_src = max(len(ids) for ids in encoded)
-        src_ids = torch.full((len(texts), max_len_src), self.pad_id, dtype=torch.long, device=device)
-        for i, ids in enumerate(encoded):
-            src_ids[i, :len(ids)] = torch.tensor(ids, dtype=torch.long)
-        src_padding_mask = src_ids.eq(self.pad_id)
-
-        with torch.no_grad():
-            output_ids = self.decode_fn(src_ids, src_padding_mask)
-
-        translations = []
-        for i in range(output_ids.size(0)):
-            text = self.tokenizer.decode(output_ids[i].tolist(), skip_special_tokens=True)
-            translations.append(text)
+        """
+        翻译一批文本。
 
-        return translations
+        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]
+        return self.translate([text])[0]

src/easytranslate/evaluation/metrics.py

@@ -18,23 +18,11 @@
 from __future__ import annotations
 
 import logging
-import time
 from typing import Optional
 
 logger = logging.getLogger(__name__)
 
 
-def _require_sacrebleu():
-    try:
-        import sacrebleu
-    except ImportError as exc:
-        raise ImportError(
-            "sacrebleu is required for BLEU/chrF/TER evaluation. "
-            "Install it with `pip install sacrebleu`."
-        ) from exc
-    return sacrebleu
-
-
 def compute_bleu(
     hypotheses: list[str],
     references: list[str],
@@ -42,37 +30,15 @@ def compute_bleu(
 ) -> dict:
     """
     计算 SacreBLEU 分数。
-    
-    Args:
-        hypotheses: 生成的译文列表
-        references: 参考译文列表
-        tokenize: 分词方式，"zh" 表示字符级别分词
-    
-    Returns:
-        dict: 包含 BLEU 分数及各 n-gram 精度的字典
-    
-    Raises:
-        ValueError: 如果输入数据为空或长度不匹配
+
+    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 (支持多参考)
     """
-    if not hypotheses or not references:
-        raise ValueError("hypotheses and references cannot be empty")
-    
-    if len(hypotheses) != len(references):
-        raise ValueError(
-            f"hypotheses and references length mismatch: "
-            f"{len(hypotheses)} vs {len(references)}"
-        )
-
-    sacrebleu = _require_sacrebleu()
-    bleu = sacrebleu.corpus_bleu(hypotheses, [references], tokenize=tokenize)
-    return {
-        "bleu": round(float(bleu.score), 4),
-        "bleu_1": round(float(bleu.precisions[0]), 4),
-        "bleu_2": round(float(bleu.precisions[1]), 4),
-        "bleu_3": round(float(bleu.precisions[2]), 4),
-        "bleu_4": round(float(bleu.precisions[3]), 4),
-        "bp": round(float(bleu.bp), 4),
-    }
+    raise NotImplementedError("TODO: Person D 实现 compute_bleu")
 
 
 def compute_comet(
@@ -85,48 +51,16 @@ def compute_comet(
 ) -> dict:
     """
     计算 COMET 分数。
-    
-    Args:
-        sources: 源文本列表
-        hypotheses: 生成的译文列表
-        references: 参考译文列表
-        model_name: COMET 模型名称
-        batch_size: 批处理大小
-        gpus: 使用的 GPU 数量
-    
-    Returns:
-        dict: 包含 COMET 系统分数和句子级别分数的字典
-    
-    Raises:
-        ValueError: 如果输入数据为空或长度不匹配
-        ImportError: 如果 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 需要源语言、翻译结果和参考翻译三者。
     """
-    try:
-        from comet import download_model, load_from_checkpoint
-    except ImportError:
-        raise ImportError(
-            "COMET is not installed. Please install it with: "
-            "pip install unbabel-comet"
-        )
-    
-    if not sources or not hypotheses or not references:
-        raise ValueError("sources, hypotheses, and references cannot be empty")
-    
-    if len(sources) != len(hypotheses) or len(hypotheses) != len(references):
-        raise ValueError(
-            f"Input lists have mismatched lengths: "
-            f"sources={len(sources)}, hypotheses={len(hypotheses)}, references={len(references)}"
-        )
-    
-    model_path = download_model(model_name)
-    model = load_from_checkpoint(model_path)
-
-    data = [{"src": s, "mt": h, "ref": r} for s, h, r in zip(sources, hypotheses, references)]
-    prediction = model.predict(data, batch_size=batch_size, gpus=gpus)
-
-    system_score = prediction.system_score
-    segment_scores = [float(s) for s in prediction.scores]
-    return {"comet": round(float(system_score), 4), "comet_scores": segment_scores}
+    raise NotImplementedError("TODO: Person D 实现 compute_comet")
 
 
 def compute_chrf(
@@ -135,29 +69,12 @@ def compute_chrf(
 ) -> dict:
     """
     计算 chrF++ 分数。
-    
-    Args:
-        hypotheses: 生成的译文列表
-        references: 参考译文列表
-    
-    Returns:
-        dict: 包含 chrF++ 分数的字典
-    
-    Raises:
-        ValueError: 如果输入数据为空或长度不匹配
-    """
-    if not hypotheses or not references:
-        raise ValueError("hypotheses and references cannot be empty")
-    
-    if len(hypotheses) != len(references):
-        raise ValueError(
-            f"hypotheses and references length mismatch: "
-            f"{len(hypotheses)} vs {len(references)}"
-        )
 
-    sacrebleu = _require_sacrebleu()
-    chrf = sacrebleu.corpus_chrf(hypotheses, [references])
-    return {"chrf": round(float(chrf.score), 4)}
+    TODO [Person D]:
+    1. 使用 sacrebleu.corpus_chrf(hypotheses, [references])
+    2. 返回 {"chrf": score}
+    """
+    raise NotImplementedError("TODO: Person D 实现 compute_chrf")
 
 
 def compute_ter(
@@ -166,89 +83,27 @@ def compute_ter(
 ) -> dict:
     """
     计算 TER 分数。
-    
-    Args:
-        hypotheses: 生成的译文列表
-        references: 参考译文列表
-    
-    Returns:
-        dict: 包含 TER 分数的字典
-    
-    Raises:
-        ValueError: 如果输入数据为空或长度不匹配
-    """
-    if not hypotheses or not references:
-        raise ValueError("hypotheses and references cannot be empty")
-    
-    if len(hypotheses) != len(references):
-        raise ValueError(
-            f"hypotheses and references length mismatch: "
-            f"{len(hypotheses)} vs {len(references)}"
-        )
 
-    sacrebleu = _require_sacrebleu()
-    ter = sacrebleu.corpus_ter(hypotheses, [references])
-    return {"ter": round(float(ter.score), 4)}
+    TODO [Person D]:
+    1. 使用 sacrebleu.corpus_ter(hypotheses, [references])
+    2. 返回 {"ter": score}
+    """
+    raise NotImplementedError("TODO: Person D 实现 compute_ter")
 
 
 def compute_all_metrics(
-    sources: Optional[list[str]] = None,
-    hypotheses: Optional[list[str]] = None,
-    references: Optional[list[str]] = None,
+    sources: list[str],
+    hypotheses: list[str],
+    references: list[str],
     metrics: list[str] = ["bleu", "comet", "chrf", "ter"],
 ) -> dict:
     """
     计算所有指定的评估指标。
-    
-    Args:
-        sources: 源文本列表 (用于 COMET)
-        hypotheses: 生成的译文列表
-        references: 参考译文列表
-        metrics: 需要计算的指标列表
-    
-    Returns:
-        dict: 包含所有计算指标的字典
-    
-    Raises:
-        ValueError: 如果 hypotheses 或 references 为空
+
+    TODO [Person D]:
+    1. 遍历 metrics 列表
+    2. 调用对应的计算函数
+    3. 合并结果并返回
+    4. 记录每个指标的计算时间
     """
-    # 验证必需参数
-    if not hypotheses:
-        raise ValueError("hypotheses cannot be empty")
-    
-    if not references:
-        raise ValueError("references cannot be empty")
-    
-    if len(hypotheses) != len(references):
-        raise ValueError(
-            f"hypotheses and references length mismatch: "
-            f"{len(hypotheses)} vs {len(references)}"
-        )
-    
-    results = {}
-    
-    for metric in metrics:
-        start = time.time()
-        try:
-            if metric == "bleu":
-                results.update(compute_bleu(hypotheses, references))
-            elif metric == "comet":
-                if sources is None:
-                    logger.warning("COMET requires sources, skipping")
-                    continue
-                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("Unknown metric: %s, skipping", metric)
-                continue
-        except Exception as e:
-            logger.error(f"Failed to compute {metric}: {str(e)}")
-            continue
-        
-        elapsed = time.time() - start
-        logger.info("Computed %s in %.2fs", metric, elapsed)
-
-    return results
+    raise NotImplementedError("TODO: Person D 实现 compute_all_metrics")

src/easytranslate/model/attention.py

@@ -25,6 +25,24 @@ 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__(
@@ -39,15 +57,8 @@ class MultiHeadAttention(nn.Module):
         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
+        raise NotImplementedError("TODO: Person B 实现 MultiHeadAttention.__init__")
 
     def forward(
         self,
@@ -56,63 +67,22 @@ class MultiHeadAttention(nn.Module):
         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)
-
-        Q = self.q_proj(query)
-        K = self.k_proj(key)
-        V = self.v_proj(value)
-
-        Q = Q.view(B, L_q, self.nhead, self.d_k).transpose(1, 2)
-        K = K.view(B, L_k, self.nhead, self.d_k).transpose(1, 2)
-        V = V.view(B, L_v, self.nhead, self.d_k).transpose(1, 2)
-
-        if self.use_rotary_embedding and self.rope is not None:
-            Q, K = self.rope.apply_rotary_pos_emb(Q, K)
-
-        scores = torch.matmul(Q, K.transpose(-2, -1)) / math.sqrt(self.d_k)
-
-        if key_padding_mask is not None:
-            scores = scores.masked_fill(
-                key_padding_mask.unsqueeze(1).unsqueeze(2), float("-inf")
-            )
-        if is_causal:
-            L_q_local, L_k_local = scores.size(-2), scores.size(-1)
-            causal_mask = torch.triu(
-                torch.ones(L_q_local, 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:
-            scores = scores.masked_fill(attn_mask.unsqueeze(0).unsqueeze(0), float("-inf"))
-
-        attn_weights = F.softmax(scores, dim=-1)
-        attn_weights = self.dropout(attn_weights)
-
-        attn_output = torch.matmul(attn_weights, V)
-
-        attn_output = attn_output.transpose(1, 2).contiguous().view(B, L_q, self.d_model)
-
-        output = self.out_proj(attn_output)
-        return output
-
-
-try:
-    from torch.nn.functional import scaled_dot_product_attention
-    _has_flash_attn = True
-except ImportError:
-    _has_flash_attn = False
+        raise NotImplementedError("TODO: Person B 实现 MultiHeadAttention.forward")
 
 
 class FlashMultiHeadAttention(nn.Module):
     """
     Flash Attention 2 加速的多头注意力。
-    
-    如果 PyTorch < 2.0 或 GPU 不支持，会自动回退到标准注意力。
+
+    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__(
@@ -123,34 +93,7 @@ class FlashMultiHeadAttention(nn.Module):
         use_rotary_embedding: bool = False,
     ):
         super().__init__()
-        
-        if not _has_flash_attn:
-            self._fallback = MultiHeadAttention(
-                d_model=d_model,
-                nhead=nhead,
-                dropout=dropout,
-                use_rotary_embedding=use_rotary_embedding,
-            )
-            self.d_model = d_model
-            self.nhead = nhead
-            self.d_k = d_model // nhead
-            self.use_rotary_embedding = use_rotary_embedding
-            self.rope: Optional[nn.Module] = None
-            return
-        
-        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
+        raise NotImplementedError("TODO: Person B 实现 FlashMultiHeadAttention.__init__")
 
     def forward(
         self,
@@ -160,54 +103,4 @@ class FlashMultiHeadAttention(nn.Module):
         key_padding_mask: Optional[torch.BoolTensor] = None,
         is_causal: bool = False,
     ) -> torch.Tensor:
-        if not _has_flash_attn:
-            return self._fallback(query, key, value, key_padding_mask, None, is_causal)
-
-        B, L_q, _ = query.size()
-        L_k = key.size(1)
-
-        Q = self.q_proj(query)
-        K = self.k_proj(key)
-        V = self.v_proj(value)
-
-        Q = Q.view(B, L_q, self.nhead, self.d_k).transpose(1, 2)
-        K = K.view(B, L_k, self.nhead, self.d_k).transpose(1, 2)
-        V = V.view(B, L_k, self.nhead, self.d_k).transpose(1, 2)
-
-        if self.use_rotary_embedding and self.rope is not None:
-            Q, K = self.rope.apply_rotary_pos_emb(Q, K)
-
-        # Build an additive attention bias so we can combine causal mask and
-        # key_padding_mask without conflicting with the is_causal flag.
-        # F.scaled_dot_product_attention raises if both is_causal and attn_mask are set.
-        if is_causal or key_padding_mask is not None:
-            attn_bias = torch.zeros(B, 1, L_q, L_k, device=Q.device, dtype=Q.dtype)
-            if is_causal:
-                causal = torch.triu(
-                    torch.ones(L_q, L_k, device=Q.device, dtype=torch.bool), diagonal=1
-                )
-                attn_bias = attn_bias.masked_fill(causal.unsqueeze(0).unsqueeze(0), float("-inf"))
-            if key_padding_mask is not None:
-                # key_padding_mask: [B, L_k] bool, True = pad
-                attn_bias = attn_bias.masked_fill(
-                    key_padding_mask.unsqueeze(1).unsqueeze(2), float("-inf")
-                )
-            attn_output = F.scaled_dot_product_attention(
-                Q, K, V,
-                attn_mask=attn_bias,
-                dropout_p=self.dropout_p if self.training else 0.0,
-                is_causal=False,
-                scale=1.0 / math.sqrt(self.d_k),
-            )
-        else:
-            attn_output = F.scaled_dot_product_attention(
-                Q, K, V,
-                attn_mask=None,
-                dropout_p=self.dropout_p if self.training else 0.0,
-                is_causal=False,
-                scale=1.0 / math.sqrt(self.d_k),
-            )
-
-        attn_output = attn_output.transpose(1, 2).contiguous().view(B, L_q, self.d_model)
-        output = self.out_proj(attn_output)
-        return output
+        raise NotImplementedError("TODO: Person B 实现 FlashMultiHeadAttention.forward")

src/easytranslate/model/decoder.py

@@ -13,23 +13,21 @@ Transformer Decoder 模块 — Person B 负责实现
 
 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
+    TODO [Person B]: 实现以下组件:
+    1. Masked Self-Attention (因果掩码，防止看到未来)
+    2. Cross-Attention (decoder 查询 encoder 输出)
+    3. Feed-Forward Network
+    4. 三个 LayerNorm
+    5. Residual connections + Dropout
     """
 
     def __init__(
@@ -44,42 +42,7 @@ class TransformerDecoderLayer(nn.Module):
         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)
+        raise NotImplementedError("TODO: Person B 实现 TransformerDecoderLayer.__init__")
 
     def forward(
         self,
@@ -89,71 +52,28 @@ class TransformerDecoderLayer(nn.Module):
         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
+        """
+        TODO [Person B]: Pre-LayerNorm 前向传播:
+        1. Masked Self-Attention with causal mask
+        2. Cross-Attention with encoder output
+        3. FFN
+        每步都有 residual connection 和 dropout
+        """
+        raise NotImplementedError("TODO: Person B 实现 TransformerDecoderLayer.forward")
 
 
 class TransformerDecoder(nn.Module):
     """
     多层 Transformer Decoder。
+
+    TODO [Person B]:
+    1. 堆叠 N 个 TransformerDecoderLayer
+    2. 最终加一个 LayerNorm
     """
 
     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)
+        raise NotImplementedError("TODO: Person B 实现 TransformerDecoder.__init__")
 
     def forward(
         self,
@@ -163,14 +83,4 @@ class TransformerDecoder(nn.Module):
         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
+        raise NotImplementedError("TODO: Person B 实现 TransformerDecoder.forward")

src/easytranslate/model/encoder.py

@@ -11,21 +11,23 @@ Transformer Encoder 模块 — Person B 负责实现
 
 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
+    TODO [Person B]: 实现以下组件:
+    1. Self-Attention: MultiHeadAttention (支持 Flash Attention)
+    2. Feed-Forward Network: Linear → Activation → Dropout → Linear
+    3. 两个 LayerNorm
+    4. Residual connections
+    5. Dropout
+
+    注意: 使用 Pre-LayerNorm 架构 (先 norm 再 attention/ffn)
     """
 
     def __init__(
@@ -40,88 +42,43 @@ class TransformerEncoderLayer(nn.Module):
         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)
+        raise NotImplementedError("TODO: Person B 实现 TransformerEncoderLayer.__init__")
 
     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
+        """
+        TODO [Person B]: Pre-LayerNorm 前向传播:
+        1. residual = src
+        2. src = layer_norm_1(src)
+        3. src = self_attention(src, src, src, key_padding_mask=src_key_padding_mask)
+        4. src = residual + dropout(src)
+        5. residual = src
+        6. src = layer_norm_2(src)
+        7. src = ffn(src)
+        8. src = residual + dropout(src)
+        """
+        raise NotImplementedError("TODO: Person B 实现 TransformerEncoderLayer.forward")
 
 
 class TransformerEncoder(nn.Module):
     """
     多层 Transformer Encoder。
+
+    TODO [Person B]:
+    1. 堆叠 N 个 TransformerEncoderLayer
+    2. 最终加一个 LayerNorm (Pre-Norm 架构需要)
     """
 
     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)
+        raise NotImplementedError("TODO: Person B 实现 TransformerEncoder.__init__")
 
     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
+        raise NotImplementedError("TODO: Person B 实现 TransformerEncoder.forward")

src/easytranslate/model/finetune.py

@@ -50,23 +50,7 @@ def load_pretrained_model(
     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
+    raise NotImplementedError("TODO: Person B 实现 load_pretrained_model")
 
 
 def setup_lora(
@@ -79,38 +63,33 @@ def setup_lora(
     """
     为模型配置 LoRA 微调。
 
+    TODO [Person B]: 实现以下逻辑:
+    1. 定义 LoraConfig:
+       - r: LoRA 秩 (低秩分解维度)
+       - lora_alpha: 缩放因子
+       - lora_dropout: LoRA dropout
+       - target_modules: 需要加 LoRA 的模块 (如 q_proj, v_proj)
+       - task_type: SEQ_2_SEQ_LM
+    2. 使用 get_peft_model(model, config) 包装模型
+    3. 打印可训练参数数量和比例
+    4. 返回 LoRA 模型
+
+    参考: https://huggingface.co/docs/peft
+
     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
+    raise NotImplementedError("TODO: Person B 实现 setup_lora")
 
 
 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
+    """
+    冻结模型所有参数，只保留 LoRA 参数可训练。
+
+    TODO [Person B]: 遍历 model.named_parameters()，
+    如果参数名中不包含 "lora"，则设置 requires_grad = False。
+    """
+    raise NotImplementedError("TODO: Person B 实现 freeze_model_except_lora")

src/easytranslate/model/positional.py

@@ -32,17 +32,7 @@ class SinusoidalPositionalEncoding(nn.Module):
 
     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)
+        raise NotImplementedError("TODO: Person B 实现 SinusoidalPositionalEncoding.__init__")
 
     def forward(self, x: torch.Tensor) -> torch.Tensor:
         """
@@ -51,8 +41,7 @@ class SinusoidalPositionalEncoding(nn.Module):
         Returns:
             x + positional_encoding: [B, L, D]
         """
-        x = x + self.pe[:, : x.size(1), :]
-        return self.dropout(x)
+        raise NotImplementedError("TODO: Person B 实现 SinusoidalPositionalEncoding.forward")
 
 
 class RotaryPositionalEmbedding(nn.Module):
@@ -83,35 +72,10 @@ class RotaryPositionalEmbedding(nn.Module):
 
     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
+        raise NotImplementedError("TODO: Person B 实现 RotaryPositionalEmbedding.__init__")
 
     def _compute_rope(self, seq_len: int, device: torch.device):
-        # Recompute if seq_len grew OR cache is stale OR device changed (e.g. after .to(cuda))
-        cache_stale = (
-            seq_len > self._cached_seq_len
-            or self._cached_cos is None
-            or self._cached_cos.device != device
-        )
-        if cache_stale:
-            t = torch.arange(seq_len, device=device, dtype=self.inv_freq.dtype)
-            freqs = torch.outer(t, self.inv_freq.to(device))  # [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
+        raise NotImplementedError("TODO: Person B 实现 _compute_rope")
 
     @staticmethod
     def _rotate_half(x: torch.Tensor) -> torch.Tensor:
@@ -130,11 +94,4 @@ class RotaryPositionalEmbedding(nn.Module):
         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
+        raise NotImplementedError("TODO: Person B 实现 apply_rotary_pos_emb")

src/easytranslate/model/transformer.py

@@ -20,8 +20,8 @@ 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.encoder import TransformerEncoder
+from easytranslate.model.decoder import TransformerDecoder
 from easytranslate.model.positional import SinusoidalPositionalEncoding, RotaryPositionalEmbedding
 
 logger = logging.getLogger(__name__)
@@ -78,94 +78,33 @@ class TransformerTranslationModel(nn.Module):
         share_embedding: bool = False,
     ):
         super().__init__()
+        # TODO [Person B]: 实现模型初始化
+        # 保存超参数
         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()
+        raise NotImplementedError("TODO: Person B 实现 TransformerTranslationModel.__init__")
 
     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)
+        """
+        参数初始化。
+
+        TODO [Person B]: 实现 Xavier/Kaiming 初始化:
+        - Embedding: normal_(0, d_model^-0.5)
+        - Linear: xavier_uniform_
+        - LayerNorm: ones_ / zeros_
+        """
+        raise NotImplementedError("TODO: Person B 实现 _init_weights")
 
     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
+        """
+        生成因果注意力掩码 (causal mask)。
+
+        TODO [Person B]:
+        返回上三角矩阵 mask，shape [sz, sz]，
+        mask[i][j] = -inf if j > i else 0
+        """
+        raise NotImplementedError("TODO: Person B 实现 _generate_square_subsequent_mask")
 
     def forward(
         self,
@@ -180,50 +119,19 @@ class TransformerTranslationModel(nn.Module):
         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
+        raise NotImplementedError("TODO: Person B 实现 forward")
 
     @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
+        """
+        仅编码（用于推理时复用 encoder 输出）。
+
+        TODO [Person B]:
+        1. src embedding + positional encoding
+        2. encoder forward
+        3. 返回 encoder_output
+        """
+        raise NotImplementedError("TODO: Person B 实现 encode")
 
     @torch.no_grad()
     def decode_step(
@@ -232,24 +140,17 @@ class TransformerTranslationModel(nn.Module):
         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
+        """
+        解码一步（用于自回归推理）。
+
+        TODO [Person B]:
+        1. tgt embedding + positional encoding
+        2. 生成 causal mask
+        3. decoder forward
+        4. 取最后一个 token 的 logits
+        5. 返回 logits [B, vocab_size]
+        """
+        raise NotImplementedError("TODO: Person B 实现 decode_step")
 
     def count_parameters(self) -> int:
         """返回可训练参数数量。"""

src/easytranslate/training/loss.py

@@ -21,11 +21,21 @@ class LabelSmoothedCrossEntropyLoss(nn.Module):
     """
     带标签平滑的交叉熵损失。
 
-    标签平滑将真实标签的概率质量从 1.0 重新分配:
-        target token 概率 = 1 - smoothing
-        其他 token 概率 = smoothing / (V - 1)
-
-    使用 KL 散度实现: loss = KL(smooth_target || log_softmax(logits))
+    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.)
     """
@@ -34,43 +44,14 @@ class LabelSmoothedCrossEntropyLoss(nn.Module):
         super().__init__()
         self.smoothing = smoothing
         self.pad_id = pad_id
-        self.confidence = 1.0 - smoothing
+        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] — 目标 token ids (padding 位置应为 pad_id, 默认 -100)
+            logits: [B, T, V]
+            targets: [B, T]
         Returns:
             loss: scalar
         """
-        logits = logits.contiguous()
-        targets = targets.contiguous()
-
-        batch_size, seq_len, vocab_size = logits.size()
-
-        logits_flat = logits.view(-1, vocab_size)
-        targets_flat = targets.view(-1)
-
-        log_probs = F.log_softmax(logits_flat, dim=-1)
-
-        # Mask that is True for real (non-padding) tokens
-        non_pad_mask = targets_flat.ne(self.pad_id).float()
-
-        if self.smoothing > 0.0:
-            smooth_dist = torch.full_like(log_probs, self.smoothing / (vocab_size - 1))
-            # Clamp to avoid negative-index writes for ignore-index values (e.g. -100)
-            safe_targets = targets_flat.clamp(min=0)
-            smooth_dist.scatter_(1, safe_targets.unsqueeze(1), self.confidence)
-            nll_loss = -torch.sum(smooth_dist * log_probs, dim=-1)
-            nll_loss = nll_loss * non_pad_mask
-        else:
-            nll_loss = F.nll_loss(
-                log_probs, targets_flat,
-                ignore_index=self.pad_id, reduction="none",
-            )
-
-        num_tokens = non_pad_mask.sum().clamp(min=1)
-        loss = nll_loss.sum() / num_tokens
-
-        return loss
+        raise NotImplementedError("TODO: Person C 实现 LabelSmoothedCrossEntropyLoss.forward")

src/easytranslate/training/optimizer.py

@@ -22,56 +22,18 @@ from torch.optim import Adam, AdamW
 from torch.optim.lr_scheduler import LambdaLR
 
 
-def build_optimizer(model: torch.nn.Module, config: dict) -> torch.optim.Optimizer:
+def build_optimizer(model, config: dict) -> torch.optim.Optimizer:
     """
     根据配置创建优化器。
 
-    支持参数分组:
-    - LayerNorm 和 bias 参数不施加 weight_decay
-    - embedding 层可使用较小的学习率
+    TODO [Person C]: 实现以下逻辑:
+    1. 从 config 中读取 optimizer type, lr, weight_decay, betas, eps
+    2. 根据 type 创建 Adam / AdamW / Adafactor
+    3. (可选) 对不同参数组设置不同学习率:
+       - embedding 层可以用较小的 lr
+       - LayerNorm 的 bias 不加 weight_decay
     """
-    opt_config = config.get("optimizer", {})
-    opt_type = opt_config.get("type", "adamw")
-    lr = float(opt_config.get("lr", 3e-4))
-    weight_decay = float(opt_config.get("weight_decay", 0.01))
-    betas = tuple(opt_config.get("betas", [0.9, 0.98]))
-    eps = float(opt_config.get("eps", 1e-8))
-
-    no_decay = ["bias", "LayerNorm.weight", "layer_norm.weight"]
-    optimizer_grouped_parameters = [
-        {
-            "params": [
-                p for n, p in model.named_parameters()
-                if p.requires_grad and not any(nd in n for nd in no_decay)
-            ],
-            "weight_decay": weight_decay,
-        },
-        {
-            "params": [
-                p for n, p in model.named_parameters()
-                if p.requires_grad and any(nd in n for nd in no_decay)
-            ],
-            "weight_decay": 0.0,
-        },
-    ]
-
-    if opt_type == "adam":
-        return Adam(optimizer_grouped_parameters, lr=lr, betas=betas, eps=eps)
-    elif opt_type == "adamw":
-        return AdamW(optimizer_grouped_parameters, lr=lr, betas=betas, eps=eps)
-    elif opt_type == "adafactor":
-        try:
-            from transformers.optimization import Adafactor
-            return Adafactor(
-                optimizer_grouped_parameters,
-                lr=lr,
-                scale_parameter=False,
-                relative_step=False,
-            )
-        except ImportError:
-            raise ImportError("Adafactor requires transformers library")
-    else:
-        raise ValueError(f"Unsupported optimizer type: {opt_type}")
+    raise NotImplementedError("TODO: Person C 实现 build_optimizer")
 
 
 def build_scheduler(
@@ -82,68 +44,18 @@ def build_scheduler(
     """
     根据配置创建学习率调度器。
 
-    支持:
-    - cosine_with_warmup: Cosine 衰减 + 线性 warmup
-    - inverse_sqrt: 经典 Transformer 调度策略
-    - linear: 线性衰减 + warmup
+    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
     """
-    sched_config = config.get("scheduler", {})
-    sched_type = sched_config.get("type", "cosine_with_warmup")
-    warmup_steps = int(sched_config.get("warmup_steps", 4000))
-    min_lr = float(sched_config.get("min_lr", 1e-6))
-
-    if sched_type == "cosine_with_warmup":
-        try:
-            from transformers import get_cosine_schedule_with_warmup
-            return get_cosine_schedule_with_warmup(
-                optimizer,
-                num_warmup_steps=warmup_steps,
-                num_training_steps=num_training_steps or 100000,
-            )
-        except ImportError:
-            return _cosine_with_warmup(optimizer, warmup_steps, num_training_steps or 100000, min_lr)
-
-    elif sched_type == "inverse_sqrt":
-        return InverseSqrtScheduler(optimizer, warmup_steps=warmup_steps)
-
-    elif sched_type == "linear":
-        try:
-            from transformers import get_linear_schedule_with_warmup
-            return get_linear_schedule_with_warmup(
-                optimizer,
-                num_warmup_steps=warmup_steps,
-                num_training_steps=num_training_steps or 100000,
-            )
-        except ImportError:
-            return _linear_with_warmup(optimizer, warmup_steps, num_training_steps or 100000, min_lr)
-
-    else:
-        raise ValueError(f"Unsupported scheduler type: {sched_type}")
-
-
-def _cosine_with_warmup(optimizer, warmup_steps, total_steps, min_lr=0.0):
-    def lr_lambda(current_step):
-        if current_step < warmup_steps:
-            return float(current_step) / float(max(1, warmup_steps))
-        progress = float(current_step - warmup_steps) / float(max(1, total_steps - warmup_steps))
-        cosine_decay = 0.5 * (1.0 + math.cos(math.pi * progress))
-        return max(min_lr / _get_base_lr(optimizer), cosine_decay)
-    return LambdaLR(optimizer, lr_lambda)
-
-
-def _linear_with_warmup(optimizer, warmup_steps, total_steps, min_lr=0.0):
-    def lr_lambda(current_step):
-        if current_step < warmup_steps:
-            return float(current_step) / float(max(1, warmup_steps))
-        progress = float(current_step - warmup_steps) / float(max(1, total_steps - warmup_steps))
-        return max(min_lr / _get_base_lr(optimizer), 1.0 - progress)
-    return LambdaLR(optimizer, lr_lambda)
-
-
-def _get_base_lr(optimizer):
-    for param_group in optimizer.param_groups:
-        return param_group.get("lr", param_group.get("initial_lr", 1e-3))
-    return 1e-3
+    raise NotImplementedError("TODO: Person C 实现 build_scheduler")
 
 
 class InverseSqrtScheduler(LambdaLR):
@@ -152,18 +64,13 @@ class InverseSqrtScheduler(LambdaLR):
 
     lr = base_lr * min(step^{-0.5}, step * warmup_steps^{-1.5})
 
-    这是原始 Transformer 论文 (Vaswani et al., 2017) 使用的调度策略。
-    warmup 阶段线性增长，warmup 后按 step^{-0.5} 衰减。
+    这是原始 Transformer 论文使用的调度策略。
+
+    TODO [Person C]:
+    1. 实现 lr_lambda 函数
+    2. warmup 阶段线性增长
+    3. warmup 后按 step^{-0.5} 衰减
     """
 
     def __init__(self, optimizer, warmup_steps: int = 4000):
-        self.warmup_steps = warmup_steps
-        warmup_factor = warmup_steps ** (-1.5)
-
-        def lr_lambda(step):
-            step += 1
-            arg1 = step ** (-0.5)
-            arg2 = step * warmup_factor
-            return min(arg1, arg2)
-
-        super().__init__(optimizer, lr_lambda)
+        raise NotImplementedError("TODO: Person C 实现 InverseSqrtScheduler.__init__")

src/easytranslate/training/trainer.py

@@ -14,10 +14,8 @@
 
 from __future__ import annotations
 
-import json
 import logging
 import os
-import shutil
 import time
 from pathlib import Path
 from typing import Optional
@@ -27,9 +25,6 @@ import torch.nn as nn
 from torch.utils.data import DataLoader
 from tqdm import tqdm
 
-from easytranslate.training.loss import LabelSmoothedCrossEntropyLoss
-from easytranslate.training.optimizer import build_optimizer, build_scheduler
-
 logger = logging.getLogger(__name__)
 
 
@@ -37,6 +32,8 @@ class Trainer:
     """
     翻译模型训练器。
 
+    TODO [Person C]: 实现以下所有方法。
+
     使用方法:
         trainer = Trainer(model, train_loader, val_loader, config)
         trainer.train()
@@ -53,407 +50,137 @@ class Trainer:
         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
-
-        train_cfg = config.get("training", {})
-
-        self.device = self._resolve_device(train_cfg)
-        self.model = self.model.to(self.device)
-
-        self.fp16 = train_cfg.get("fp16", False)
-        self.bf16 = train_cfg.get("bf16", False)
-        self.gradient_accumulation_steps = int(train_cfg.get("gradient_accumulation_steps", 1))
-        self.max_grad_norm = float(train_cfg.get("max_grad_norm", 1.0))
-        self.gradient_checkpointing = train_cfg.get("gradient_checkpointing", False)
-
-        if self.gradient_checkpointing and hasattr(self.model, "gradient_checkpointing_enable"):
-            self.model.gradient_checkpointing_enable()
-
-        self.num_epochs = int(train_cfg.get("epochs", 30))
-        self.max_steps = int(train_cfg.get("max_steps", -1))
-
-        self.scaler = None
-        self.amp_dtype = None
-        if self.fp16:
-            self.scaler = torch.amp.GradScaler(self.device.type)
-            self.amp_dtype = torch.float16
-        elif self.bf16:
-            self.amp_dtype = torch.bfloat16
-
-        if optimizer is None:
-            optimizer = build_optimizer(model, train_cfg)
-        self.optimizer = optimizer
-
-        total_steps = self.num_epochs * len(train_loader) // self.gradient_accumulation_steps
-        if scheduler is None:
-            scheduler = build_scheduler(optimizer, train_cfg, num_training_steps=total_steps)
-        self.scheduler = scheduler
-
-        reg_cfg = train_cfg.get("regularization", {})
-        if criterion is None:
-            criterion = LabelSmoothedCrossEntropyLoss(
-                smoothing=float(reg_cfg.get("label_smoothing", 0.1)),
-                pad_id=-100,  # must match TranslationCollator's label_pad_token_id
-            )
-        self.criterion = criterion
-
-        self.evaluator = evaluator
-
-        ckpt_cfg = train_cfg.get("checkpoint", {})
-        self.checkpoint_dir = Path(ckpt_cfg.get("save_dir", "checkpoints/"))
-        self.checkpoint_dir.mkdir(parents=True, exist_ok=True)
-        self.save_every_n_steps = int(ckpt_cfg.get("save_every_n_steps", 5000))
-        self.save_best = ckpt_cfg.get("save_best", True)
-        self.metric_for_best = ckpt_cfg.get("metric_for_best", "bleu")
-        self.max_checkpoints = int(ckpt_cfg.get("max_checkpoints", 5))
-
-        es_cfg = train_cfg.get("early_stopping", {})
-        self.early_stopping_enabled = es_cfg.get("enabled", True)
-        self.patience = int(es_cfg.get("patience", 5))
-        self.min_delta = float(es_cfg.get("min_delta", 0.1))
-
-        log_cfg = config.get("logging", {})
-        self.log_dir = Path(log_cfg.get("log_dir", "logs/"))
-        self.log_dir.mkdir(parents=True, exist_ok=True)
-        self.log_every_n_steps = int(log_cfg.get("log_every_n_steps", 100))
-        self.log_backend = log_cfg.get("backend", "tensorboard")
-
-        self._init_logger()
-
-        self.global_step = 0
-        self.current_epoch = 0
-        self.best_metric_value = float("-inf")
-        self.best_epoch = 0
-        self.epochs_without_improvement = 0
-        self.train_loss_history: list[float] = []
-        self.val_metrics_history: list[dict] = []
-
-        self._setup_distributed()
-
-    def _resolve_device(self, train_cfg: dict) -> torch.device:
-        device_str = train_cfg.get("device", "auto")
-        if device_str == "auto":
-            if torch.cuda.is_available():
-                return torch.device("cuda")
-            elif torch.backends.mps.is_available():
-                return torch.device("mps")
-            else:
-                return torch.device("cpu")
-        return torch.device(device_str)
-
-    def _init_logger(self):
-        self.writer = None
-        if self.log_backend in ("tensorboard", "both"):
-            try:
-                from torch.utils.tensorboard import SummaryWriter
-                self.writer = SummaryWriter(log_dir=str(self.log_dir))
-            except ImportError:
-                logger.warning("TensorBoard not available, skipping")
-
-        self.wandb_run = None
-        if self.log_backend in ("wandb", "both"):
-            try:
-                import wandb
-                log_cfg = self.config.get("logging", {})
-                self.wandb_run = wandb.init(
-                    project=log_cfg.get("project_name", "EasyTranslate"),
-                    config=self.config,
-                    dir=str(self.log_dir),
-                )
-            except ImportError:
-                logger.warning("WandB not available, skipping")
-
-    def _setup_distributed(self):
-        dist_cfg = self.config.get("training", {}).get("distributed", {})
-        strategy = dist_cfg.get("strategy", "ddp")
-
-        if strategy == "ddp" and torch.distributed.is_available() and torch.distributed.is_initialized():
-            self.model = nn.parallel.DistributedDataParallel(self.model)
-            self.is_distributed = True
-        elif strategy == "deepspeed":
-            try:
-                import deepspeed
-                ds_config = dist_cfg.get("deepspeed_config", "configs/deepspeed_config.json")
-                self.model, self.optimizer, _, _ = deepspeed.initialize(
-                    model=self.model,
-                    optimizer=self.optimizer,
-                    config_params=ds_config,
-                )
-                self.is_distributed = True
-            except ImportError:
-                logger.warning("DeepSpeed not available, falling back to single GPU")
-                self.is_distributed = False
-        else:
-            self.is_distributed = False
+        raise NotImplementedError("TODO: Person C 实现 Trainer.__init__")
 
     def train(self):
-        logger.info("Starting training for %d epochs", self.num_epochs)
-        logger.info("Device: %s, FP16: %s, BF16: %s", self.device, self.fp16, self.bf16)
-        logger.info("Gradient accumulation steps: %d", self.gradient_accumulation_steps)
-
-        for epoch in range(self.current_epoch, self.num_epochs):
-            self.current_epoch = epoch
-            logger.info("=" * 50)
-            logger.info("Epoch %d/%d", epoch + 1, self.num_epochs)
-
-            train_loss = self._train_one_epoch(epoch)
-            self.train_loss_history.append(train_loss)
-
-            val_metrics = self._validate(epoch)
-            self.val_metrics_history.append(val_metrics)
-
-            self._log_metrics(epoch, train_loss, val_metrics)
-            self._save_checkpoint(epoch, val_metrics)
-
-            if self._should_early_stop(val_metrics):
-                logger.info("Early stopping triggered at epoch %d", epoch + 1)
-                break
-
-            if self.max_steps > 0 and self.global_step >= self.max_steps:
-                logger.info("Reached max steps %d, stopping", self.max_steps)
-                break
-
-        self._log_final_results()
-        self._cleanup()
+        """
+        主训练循环。
+
+        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:
-        self.model.train()
-        total_loss = 0.0
-        num_batches = 0
-        self.optimizer.zero_grad()
-
-        pbar = tqdm(self.train_loader, desc=f"Train Epoch {epoch + 1}", leave=False)
-        for batch_idx, batch in enumerate(pbar):
-            batch = self._move_batch_to_device(batch)
+        """
+        训练一个 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")
 
-            with torch.amp.autocast(self.device.type, enabled=self.amp_dtype is not None, dtype=self.amp_dtype):
-                logits = self.model(
-                    batch["src_ids"],
-                    batch["tgt_input_ids"],
-                    batch.get("src_padding_mask"),
-                    batch.get("tgt_padding_mask"),
-                )
-                loss = self.criterion(logits, batch["labels"])
-
-            loss = loss / self.gradient_accumulation_steps
-
-            if self.scaler is not None:
-                self.scaler.scale(loss).backward()
-            else:
-                loss.backward()
-
-            total_loss += loss.item() * self.gradient_accumulation_steps
-            num_batches += 1
-
-            if (batch_idx + 1) % self.gradient_accumulation_steps == 0:
-                if self.scaler is not None:
-                    self.scaler.unscale_(self.optimizer)
-                    nn.utils.clip_grad_norm_(self.model.parameters(), self.max_grad_norm)
-                    self.scaler.step(self.optimizer)
-                    self.scaler.update()
-                else:
-                    nn.utils.clip_grad_norm_(self.model.parameters(), self.max_grad_norm)
-                    self.optimizer.step()
-
-                self.scheduler.step()
-                self.optimizer.zero_grad()
-                self.global_step += 1
-
-                current_lr = self.scheduler.get_last_lr()[0]
-                pbar.set_postfix({
-                    "loss": f"{loss.item() * self.gradient_accumulation_steps:.4f}",
-                    "lr": f"{current_lr:.2e}",
-                    "step": self.global_step,
-                })
-
-                if self.global_step % self.log_every_n_steps == 0:
-                    self._log_step_metrics(loss.item() * self.gradient_accumulation_steps, current_lr)
-
-                if self.save_every_n_steps > 0 and self.global_step % self.save_every_n_steps == 0:
-                    self._save_checkpoint(epoch, {"step": self.global_step}, prefix=f"step_{self.global_step}")
-
-        avg_loss = total_loss / max(num_batches, 1)
-        logger.info("Epoch %d - Train Loss: %.4f", epoch + 1, avg_loss)
-        return avg_loss
-
-    @torch.no_grad()
     def _validate(self, epoch: int) -> dict:
-        self.model.eval()
-        total_loss = 0.0
-        num_batches = 0
-
-        pbar = tqdm(self.val_loader, desc=f"Val Epoch {epoch + 1}", leave=False)
-        for batch in pbar:
-            batch = self._move_batch_to_device(batch)
-
-            with torch.amp.autocast(self.device.type, enabled=self.amp_dtype is not None, dtype=self.amp_dtype):
-                logits = self.model(
-                    batch["src_ids"],
-                    batch["tgt_input_ids"],
-                    batch.get("src_padding_mask"),
-                    batch.get("tgt_padding_mask"),
-                )
-                loss = self.criterion(logits, batch["labels"])
-            total_loss += loss.item()
-            num_batches += 1
-
-        val_loss = total_loss / max(num_batches, 1)
-        metrics = {"val_loss": round(val_loss, 4)}
-
-        if self.evaluator is not None and self.config.get("evaluation", {}).get("eval_on_epoch_end", True):
-            try:
-                eval_results = self.evaluator.evaluate(self.val_loader)
-                metrics.update(eval_results)
-            except Exception as e:
-                logger.warning("Evaluation failed during validation: %s", e)
-
-        logger.info("Epoch %d - Val Loss: %.4f", epoch + 1, val_loss)
-        for k, v in metrics.items():
-            if k != "val_loss" and not isinstance(v, list):
-                logger.info("  %s: %.4f", k, v)
-
-        return metrics
-
-    def _save_checkpoint(self, epoch: int, metrics: dict, prefix: str = ""):
-        model_to_save = self.model.module if hasattr(self.model, "module") else self.model
-
-        checkpoint = {
-            "epoch": epoch,
-            "step": self.global_step,
-            "model_state_dict": model_to_save.state_dict(),
-            "optimizer_state_dict": self.optimizer.state_dict(),
-            "scheduler_state_dict": self.scheduler.state_dict(),
-            "metrics": metrics,
-            "config": self.config,
-            "train_loss_history": self.train_loss_history,
-            "val_metrics_history": self.val_metrics_history,
-        }
-
-        if prefix:
-            ckpt_path = self.checkpoint_dir / f"checkpoint_{prefix}.pt"
-        else:
-            ckpt_path = self.checkpoint_dir / f"checkpoint_epoch_{epoch + 1}.pt"
-
-        torch.save(checkpoint, ckpt_path)
-        logger.info("Checkpoint saved: %s", ckpt_path)
-
-        current_metric = metrics.get(self.metric_for_best, metrics.get("val_loss", float("inf")))
-        if self.metric_for_best == "val_loss":
-            current_metric = -current_metric
-
-        if self.save_best and current_metric > self.best_metric_value:
-            self.best_metric_value = current_metric
-            self.best_epoch = epoch
-            best_path = self.checkpoint_dir / "best_model.pt"
-            torch.save(checkpoint, best_path)
-            logger.info("New best model saved: %s (metric: %.4f)", best_path, current_metric)
-
-        self._cleanup_old_checkpoints()
-
-    def _cleanup_old_checkpoints(self):
-        ckpt_files = sorted(
-            self.checkpoint_dir.glob("checkpoint_epoch_*.pt"),
-            key=os.path.getmtime,
-        )
-        while len(ckpt_files) > self.max_checkpoints:
-            oldest = ckpt_files.pop(0)
-            oldest.unlink()
-            logger.debug("Removed old checkpoint: %s", oldest)
+        """
+        验证。
+
+        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):
-        logger.info("Loading checkpoint from %s", checkpoint_path)
-        checkpoint = torch.load(checkpoint_path, map_location=self.device)
+        """
+        加载检查点继续训练。
 
-        model_to_load = self.model.module if hasattr(self.model, "module") else self.model
-        model_to_load.load_state_dict(checkpoint["model_state_dict"])
-
-        self.optimizer.load_state_dict(checkpoint["optimizer_state_dict"])
-        self.scheduler.load_state_dict(checkpoint["scheduler_state_dict"])
-        self.current_epoch = checkpoint["epoch"] + 1
-        self.global_step = checkpoint["step"]
-        self.best_metric_value = checkpoint.get("metrics", {}).get(
-            self.metric_for_best, float("-inf")
-        )
-        self.train_loss_history = checkpoint.get("train_loss_history", [])
-        self.val_metrics_history = checkpoint.get("val_metrics_history", [])
-
-        logger.info("Resumed from epoch %d, step %d", self.current_epoch, self.global_step)
+        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:
-        if not self.early_stopping_enabled:
-            return False
-
-        current_metric = metrics.get(self.metric_for_best, metrics.get("val_loss", float("inf")))
-        if self.metric_for_best == "val_loss":
-            current_metric = -current_metric
-
-        if current_metric > self.best_metric_value + self.min_delta:
-            self.epochs_without_improvement = 0
-            return False
+        """
+        判断是否应该早停。
 
-        self.epochs_without_improvement += 1
-        logger.info(
-            "No improvement for %d epochs (best: %.4f, current: %.4f)",
-            self.epochs_without_improvement, self.best_metric_value, current_metric,
-        )
-        return self.epochs_without_improvement >= self.patience
+        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):
-        if self.writer is not None:
-            self.writer.add_scalar("Loss/train", train_loss, epoch)
-            for k, v in val_metrics.items():
-                if not isinstance(v, list):
-                    self.writer.add_scalar(f"Metrics/{k}", v, epoch)
-            self.writer.add_scalar("LR", self.scheduler.get_last_lr()[0], epoch)
+        """
+        记录训练指标到 TensorBoard / WandB。
 
-        if self.wandb_run is not None:
-            import wandb
-            log_dict = {"epoch": epoch, "train_loss": train_loss}
-            for k, v in val_metrics.items():
-                if not isinstance(v, list):
-                    log_dict[f"val/{k}"] = v
-            log_dict["lr"] = self.scheduler.get_last_lr()[0]
-            wandb.log(log_dict, step=self.global_step)
+        TODO [Person C]:
+        1. 使用 self.logger 记录 train_loss, val_loss, bleu, lr 等
+        2. 打印到控制台 (使用 rich 库的 table 格式)
+        """
+        raise NotImplementedError("TODO: Person C 实现 _log_metrics")
 
-    def _log_step_metrics(self, loss: float, lr: float):
-        if self.writer is not None:
-            self.writer.add_scalar("Loss/train_step", loss, self.global_step)
-            self.writer.add_scalar("LR/step", lr, self.global_step)
-
-        if self.wandb_run is not None:
-            import wandb
-            wandb.log({"train/loss_step": loss, "lr": lr}, step=self.global_step)
-
-    def _log_final_results(self):
-        logger.info("=" * 50)
-        logger.info("Training completed!")
-        logger.info("Best epoch: %d", self.best_epoch + 1)
-        logger.info("Best %s: %.4f", self.metric_for_best, self.best_metric_value)
-
-        summary = {
-            "best_epoch": self.best_epoch,
-            "best_metric": self.best_metric_value,
-            "metric_name": self.metric_for_best,
-            "total_steps": self.global_step,
-            "train_loss_history": self.train_loss_history,
-            "val_metrics_history": self.val_metrics_history,
-        }
-        summary_path = self.checkpoint_dir / "training_summary.json"
-        with open(summary_path, "w", encoding="utf-8") as f:
-            json.dump(summary, f, indent=2, ensure_ascii=False, default=str)
-        logger.info("Training summary saved to %s", summary_path)
+    def _setup_distributed(self):
+        """
+        设置分布式训练。
 
-    def _cleanup(self):
-        if self.writer is not None:
-            self.writer.close()
-        if self.wandb_run is not None:
-            self.wandb_run.finish()
+        TODO [Person C]: 根据 config["training"]["distributed"]["strategy"]:
+        1. "ddp": 使用 torch.nn.parallel.DistributedDataParallel
+        2. "fsdp": 使用 torch.distributed.fsdp.FullyShardedDataParallel
+        3. "deepspeed": 使用 deepspeed.initialize()
 
-    def _move_batch_to_device(self, batch: dict) -> dict:
-        return {
-            k: v.to(self.device, non_blocking=True) if isinstance(v, torch.Tensor) else v
-            for k, v in batch.items()
-        }
+        也可以使用 HuggingFace Accelerate 统一处理。
+        """
+        raise NotImplementedError("TODO: Person C 实现 _setup_distributed")

src/easytranslate/utils/__init__.py

@@ -1,30 +1,7 @@
 """工具模块"""
 
-from easytranslate.utils.config import load_config, merge_configs, config_from_cli, config_to_dict
+from easytranslate.utils.config import load_config, merge_configs
 from easytranslate.utils.seed import set_seed
 from easytranslate.utils.logging import setup_logging
-from easytranslate.utils.cloud_storage import (
-    is_colab_environment,
-    mount_google_drive,
-    sync_all_to_drive,
-    sync_checkpoints_to_drive,
-    sync_logs_to_drive,
-    save_training_summary_to_drive,
-    setup_colab_environment,
-)
 
-__all__ = [
-    "load_config",
-    "merge_configs",
-    "config_from_cli",
-    "config_to_dict",
-    "set_seed",
-    "setup_logging",
-    "is_colab_environment",
-    "mount_google_drive",
-    "sync_all_to_drive",
-    "sync_checkpoints_to_drive",
-    "sync_logs_to_drive",
-    "save_training_summary_to_drive",
-    "setup_colab_environment",
-]
+__all__ = ["load_config", "merge_configs", "set_seed", "setup_logging"]

src/easytranslate/utils/cloud_storage.py

@@ -1,188 +0,0 @@
-"""
-Cloud storage integration module for Google Colab and Google Drive.
-
-Provides secure, authenticated storage operations for training results,
-checkpoints, and logs. Supports automatic environment detection for
-Colab vs local execution.
-"""
-
-from __future__ import annotations
-
-import json
-import logging
-import os
-import shutil
-import time
-from pathlib import Path
-from typing import Optional
-
-logger = logging.getLogger(__name__)
-
-
-def is_colab_environment() -> bool:
-    try:
-        import google.colab
-        return True
-    except ImportError:
-        return False
-
-
-def mount_google_drive(mount_point: str = "/content/drive") -> bool:
-    if not is_colab_environment():
-        logger.info("Not running in Colab, skipping Google Drive mount")
-        return False
-
-    try:
-        from google.colab import drive
-        drive.mount(mount_point)
-        logger.info("Google Drive mounted at %s", mount_point)
-        return True
-    except Exception as e:
-        logger.error("Failed to mount Google Drive: %s", e)
-        return False
-
-
-def get_drive_path(base_path: str = "/content/drive/MyDrive/EasyTranslate") -> Optional[Path]:
-    if not os.path.exists("/content/drive"):
-        logger.warning("Google Drive not mounted, cannot resolve drive path")
-        return None
-
-    drive_path = Path(base_path)
-    drive_path.mkdir(parents=True, exist_ok=True)
-    return drive_path
-
-
-def sync_checkpoints_to_drive(
-    local_checkpoint_dir: str | Path,
-    drive_base_path: str = "/content/drive/MyDrive/EasyTranslate",
-    max_retries: int = 3,
-) -> bool:
-    local_dir = Path(local_checkpoint_dir)
-    if not local_dir.exists():
-        logger.warning("Local checkpoint directory does not exist: %s", local_dir)
-        return False
-
-    drive_path = get_drive_path(drive_base_path)
-    if drive_path is None:
-        return False
-
-    drive_checkpoint_dir = drive_path / "checkpoints"
-    drive_checkpoint_dir.mkdir(parents=True, exist_ok=True)
-
-    success = True
-    for ckpt_file in local_dir.glob("*.pt"):
-        dest = drive_checkpoint_dir / ckpt_file.name
-        for attempt in range(max_retries):
-            try:
-                shutil.copy2(ckpt_file, dest)
-                logger.info("Synced checkpoint to Drive: %s", dest)
-                break
-            except Exception as e:
-                logger.warning("Sync attempt %d/%d failed for %s: %s", attempt + 1, max_retries, ckpt_file.name, e)
-                if attempt == max_retries - 1:
-                    success = False
-                time.sleep(2 ** attempt)
-
-    return success
-
-
-def sync_logs_to_drive(
-    local_log_dir: str | Path,
-    drive_base_path: str = "/content/drive/MyDrive/EasyTranslate",
-    max_retries: int = 3,
-) -> bool:
-    local_dir = Path(local_log_dir)
-    if not local_dir.exists():
-        logger.warning("Local log directory does not exist: %s", local_dir)
-        return False
-
-    drive_path = get_drive_path(drive_base_path)
-    if drive_path is None:
-        return False
-
-    drive_log_dir = drive_path / "logs"
-    drive_log_dir.mkdir(parents=True, exist_ok=True)
-
-    success = True
-    for log_file in local_dir.glob("*"):
-        if log_file.is_file():
-            dest = drive_log_dir / log_file.name
-            for attempt in range(max_retries):
-                try:
-                    shutil.copy2(log_file, dest)
-                    break
-                except Exception as e:
-                    logger.warning("Log sync attempt %d/%d failed: %s", attempt + 1, max_retries, e)
-                    if attempt == max_retries - 1:
-                        success = False
-                    time.sleep(2 ** attempt)
-
-    return success
-
-
-def save_training_summary_to_drive(
-    summary: dict,
-    drive_base_path: str = "/content/drive/MyDrive/EasyTranslate",
-) -> bool:
-    drive_path = get_drive_path(drive_base_path)
-    if drive_path is None:
-        return False
-
-    summary_path = drive_path / "training_summary.json"
-    try:
-        with open(summary_path, "w", encoding="utf-8") as f:
-            json.dump(summary, f, indent=2, ensure_ascii=False, default=str)
-        logger.info("Training summary saved to Drive: %s", summary_path)
-        return True
-    except Exception as e:
-        logger.error("Failed to save training summary to Drive: %s", e)
-        return False
-
-
-def sync_all_to_drive(
-    checkpoint_dir: str | Path = "checkpoints",
-    log_dir: str | Path = "logs",
-    drive_base_path: str = "/content/drive/MyDrive/EasyTranslate",
-) -> dict:
-    results = {
-        "checkpoints_synced": sync_checkpoints_to_drive(checkpoint_dir, drive_base_path),
-        "logs_synced": sync_logs_to_drive(log_dir, drive_base_path),
-    }
-
-    summary_path = Path(checkpoint_dir) / "training_summary.json"
-    if summary_path.exists():
-        with open(summary_path, "r", encoding="utf-8") as f:
-            summary = json.load(f)
-        results["summary_saved"] = save_training_summary_to_drive(summary, drive_base_path)
-
-    logger.info("Drive sync results: %s", results)
-    return results
-
-
-def setup_colab_environment() -> dict:
-    env_info = {
-        "is_colab": is_colab_environment(),
-        "gpu_available": False,
-        "gpu_info": None,
-        "drive_mounted": False,
-        "drive_path": None,
-    }
-
-    if env_info["is_colab"]:
-        env_info["drive_mounted"] = mount_google_drive()
-        env_info["drive_path"] = str(get_drive_path()) if env_info["drive_mounted"] else None
-
-    try:
-        import torch
-        env_info["gpu_available"] = torch.cuda.is_available()
-        if env_info["gpu_available"]:
-            env_info["gpu_info"] = {
-                "device_count": torch.cuda.device_count(),
-                "device_name": torch.cuda.get_device_name(0),
-                "device_capability": torch.cuda.get_device_capability(0),
-            }
-    except ImportError:
-        pass
-
-    logger.info("Environment setup: %s", {k: v for k, v in env_info.items() if k != "gpu_info"})
-    return env_info

src/easytranslate/utils/config.py

@@ -19,66 +19,33 @@ def load_config(config_path: str | Path) -> DictConfig:
     """
     加载 YAML 配置文件。
 
-    Args:
-        config_path: 配置文件路径
-
-    Returns:
-        DictConfig: OmegaConf 配置对象
+    TODO [Person E]: 实现以下逻辑:
+    1. 使用 OmegaConf.load(config_path) 加载配置
+    2. 验证配置完整性
+    3. 返回 DictConfig 对象
     """
-    config_path = Path(config_path)
-    if not config_path.exists():
-        raise FileNotFoundError(f"Config file not found: {config_path}")
-
-    config = OmegaConf.load(config_path)
-    logger.info("Loaded config from %s", config_path)
-    return config
+    raise NotImplementedError("TODO: Person E 实现 load_config")
 
 
 def merge_configs(base_config: DictConfig, override_config: DictConfig) -> DictConfig:
     """
     合并配置 (override 覆盖 base)。
 
-    Args:
-        base_config: 基础配置
-        override_config: 覆盖配置
-
-    Returns:
-        DictConfig: 合并后的配置
+    TODO [Person E]:
+    使用 OmegaConf.merge(base_config, override_config)
     """
-    return 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:
     """
     从配置文件 + 命令行参数构建最终配置。
 
-    使用方式: python train.py --config configs/default.yaml training.lr=1e-4 model.nhead=16
-
-    Args:
-        config_path: 配置文件路径
-        cli_args: 命令行覆盖参数列表
-
-    Returns:
-        DictConfig: 最终配置
-    """
-    config = load_config(config_path)
-
-    if cli_args:
-        cli_config = OmegaConf.from_cli(cli_args)
-        config = merge_configs(config, cli_config)
-        logger.info("Applied %d CLI overrides", len(cli_args))
-
-    return config
+    TODO [Person E]:
+    1. 加载 config_path
+    2. 使用 OmegaConf.from_cli(cli_args) 解析命令行参数
+    3. 合并并返回
 
-
-def config_to_dict(config: DictConfig) -> dict:
-    """
-    将 OmegaConf DictConfig 转换为普通 Python dict。
-
-    Args:
-        config: OmegaConf 配置对象
-
-    Returns:
-        dict: 普通 Python 字典
+    使用方式: python train.py --config configs/default.yaml training.lr=1e-4 model.nhead=16
     """
-    return OmegaConf.to_container(config, resolve=True)
+    raise NotImplementedError("TODO: Person E 实现 config_from_cli")

src/easytranslate/utils/logging.py

@@ -16,44 +16,11 @@ def setup_logging(
     """
     配置日志系统。
 
-    Args:
-        log_dir: 日志目录路径
-        level: 日志级别
-        log_file: 日志文件名 (相对于 log_dir)
+    TODO [Person E]: 实现以下逻辑:
+    1. 创建 root logger
+    2. 设置 StreamHandler (控制台输出)
+    3. (可选) 设置 FileHandler (文件输出)
+    4. 使用 rich.logging.RichHandler 美化控制台输出
+    5. 设置日志格式: [时间] [级别] [模块名] 消息
     """
-    root_logger = logging.getLogger()
-    root_logger.setLevel(level)
-
-    for handler in root_logger.handlers[:]:
-        root_logger.removeHandler(handler)
-
-    fmt = logging.Formatter(
-        "[%(asctime)s] [%(levelname)s] [%(name)s] %(message)s",
-        datefmt="%Y-%m-%d %H:%M:%S",
-    )
-
-    try:
-        from rich.logging import RichHandler
-        console_handler = RichHandler(rich_tracebacks=True, markup=True)
-        console_handler.setLevel(level)
-        console_handler.setFormatter(logging.Formatter("%(message)s"))
-        root_logger.addHandler(console_handler)
-    except ImportError:
-        console_handler = logging.StreamHandler(sys.stdout)
-        console_handler.setLevel(level)
-        console_handler.setFormatter(fmt)
-        root_logger.addHandler(console_handler)
-
-    if log_dir is not None:
-        log_dir = Path(log_dir)
-        log_dir.mkdir(parents=True, exist_ok=True)
-        file_name = log_file or "easytranslate.log"
-        file_handler = logging.FileHandler(log_dir / file_name, encoding="utf-8")
-        file_handler.setLevel(level)
-        file_handler.setFormatter(fmt)
-        root_logger.addHandler(file_handler)
-
-    logging.getLogger("matplotlib").setLevel(logging.WARNING)
-    logging.getLogger("PIL").setLevel(logging.WARNING)
-
-    root_logger.info("Logging initialized (level=%s)", logging.getLevelName(level))
+    raise NotImplementedError("TODO: Person E 实现 setup_logging")

src/easytranslate/utils/seed.py

@@ -2,7 +2,6 @@
 随机种子管理模块 — 公共模块
 """
 
-import os
 import random
 
 import numpy as np
@@ -13,13 +12,12 @@ def set_seed(seed: int = 42):
     """
     设置全局随机种子以确保实验可复现。
 
-    Args:
-        seed: 随机种子值
+    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
     """
-    random.seed(seed)
-    np.random.seed(seed)
-    torch.manual_seed(seed)
-    torch.cuda.manual_seed_all(seed)
-    torch.backends.cudnn.deterministic = True
-    torch.backends.cudnn.benchmark = False
-    os.environ["PYTHONHASHSEED"] = str(seed)
+    raise NotImplementedError("TODO: Person E 实现 set_seed")