# ============================================================================
# EasyTranslate 默认配置文件
# 基于 Transformer 架构的英中翻译模型
# ============================================================================

# ---------- 模型配置 ----------
model:
  # 模型类型: "transformer_scratch" | "finetune_nllb" | "finetune_mbart"
  type: "transformer_scratch"

  # === Transformer from scratch 配置 ===
  transformer:
    d_model: 512               # 模型维度
    nhead: 8                   # 多头注意力头数
    num_encoder_layers: 6      # 编码器层数
    num_decoder_layers: 6      # 解码器层数
    dim_feedforward: 2048      # FFN 中间维度
    dropout: 0.1               # Dropout 概率
    activation: "gelu"         # 激活函数: "relu" | "gelu"
    max_seq_len: 512           # 最大序列长度
    use_flash_attention: true  # 是否使用 Flash Attention 2
    use_rotary_embedding: true # 是否使用旋转位置编码 (RoPE)
    pre_norm: true             # Pre-LayerNorm (更稳定)

  # === 预训练模型微调配置 ===
  pretrained:
    model_name: "facebook/nllb-200-distilled-600M"  # HuggingFace 模型名
    src_lang: "eng_Latn"       # 源语言代码
    tgt_lang: "zho_Hans"       # 目标语言代码
    use_lora: true             # 是否使用 LoRA 微调
    lora:
      r: 16                    # LoRA 秩
      alpha: 32                # LoRA alpha
      dropout: 0.05            # LoRA dropout
      target_modules:          # LoRA 目标模块
        - "q_proj"
        - "v_proj"
        - "k_proj"
        - "o_proj"

# ---------- 分词器配置 ----------
tokenizer:
  # 分词器类型: "bpe" | "sentencepiece" | "pretrained"
  type: "bpe"
  vocab_size: 32000            # 词汇表大小
  min_frequency: 2             # 最小词频
  special_tokens:
    pad: "<pad>"
    unk: "<unk>"
    bos: "<s>"
    eos: "</s>"
  max_length: 512              # 最大 token 长度

# ---------- 数据配置 ----------
data:
  # 数据集: "wmt" | "opus" | "custom"
  dataset_name: "wmt"
  wmt:
    year: "19"                 # WMT 年份
    language_pair: "zh-en"     # 语言对
  opus:
    subset: "UNPC"             # OPUS 子集名
  custom:
    train_src: "data/train.en" # 自定义训练集源语言
    train_tgt: "data/train.zh"
    val_src: "data/val.en"
    val_tgt: "data/val.zh"
    test_src: "data/test.en"
    test_tgt: "data/test.zh"

  # 数据处理
  preprocessing:
    lowercase_src: false       # 源语言是否小写化
    remove_punctuation: false  # 是否去除标点
    max_src_len: 256           # 源语言最大长度
    max_tgt_len: 256           # 目标语言最大长度
    filter_by_length: true     # 是否按长度过滤
    length_ratio_threshold: 3.0 # 长度比阈值

  # DataLoader
  dataloader:
    batch_size: 32
    num_workers: 2        # 2 workers; set 0 for CPU-only environments
    pin_memory: true
    dynamic_batching: true     # 动态 batch (按 token 数)
    max_tokens_per_batch: 8192 # 动态 batch 最大 token 数

# ---------- 训练配置 ----------
training:
  # 基础训练参数
  epochs: 30
  max_steps: -1                # -1 表示按 epoch 训练
  gradient_accumulation_steps: 4
  fp16: true                   # 混合精度训练
  bf16: false                  # BF16 (A100+)
  gradient_checkpointing: false

  # 优化器
  optimizer:
    type: "adamw"              # "adam" | "adamw" | "adafactor"
    lr: 3.0e-4
    weight_decay: 0.01
    betas: [0.9, 0.98]
    eps: 1.0e-8

  # 学习率调度
  scheduler:
    type: "cosine_with_warmup" # "cosine_with_warmup" | "inverse_sqrt" | "linear"
    warmup_steps: 4000
    min_lr: 1.0e-6

  # 正则化
  regularization:
    label_smoothing: 0.1       # 标签平滑
    dropout: 0.1

  # 检查点
  checkpoint:
    save_dir: "checkpoints/"
    save_every_n_steps: 5000
    save_best: true            # 保存最佳模型
    metric_for_best: "bleu"    # 选择最佳模型的指标
    max_checkpoints: 5         # 最大保存检查点数

  # 早停
  early_stopping:
    enabled: true
    patience: 5
    min_delta: 0.1

  # 分布式训练
  distributed:
    strategy: "ddp"            # "ddp" | "fsdp" | "deepspeed"
    deepspeed_config: "configs/deepspeed_config.json"

# ---------- 评估配置 ----------
evaluation:
  # 评估指标
  metrics:
    - "bleu"                   # SacreBLEU
    # - "comet"                # COMET (需下载 ~4 GB 模型，可按需开启)
    - "chrf"                   # chrF++
    - "ter"                    # TER

  # 解码策略
  decoding:
    strategy: "beam_search"    # "greedy" | "beam_search" | "sampling"
    beam_size: 5
    length_penalty: 1.0
    no_repeat_ngram_size: 3
    max_decode_len: 256

    # Sampling 参数
    sampling:
      temperature: 0.7
      top_k: 50
      top_p: 0.9

  # 评估频率
  eval_every_n_steps: 1000
  eval_on_epoch_end: true

# ---------- 日志配置 ----------
logging:
  # 日志工具: "wandb" | "tensorboard" | "both"
  backend: "tensorboard"
  project_name: "EasyTranslate"
  log_every_n_steps: 100
  log_dir: "logs/"

# ---------- 推理/部署配置 ----------
inference:
  model_path: "checkpoints/best_model"
  device: "cuda"
  batch_size: 16
  quantization: null           # null | "int8" | "int4"

# ---------- 实验配置 ----------
experiment:
  seed: 42
  name: "baseline"
  output_dir: "outputs/"