# @package _global_

defaults:
  - override /criterion: cross_entropy
  - override /callbacks: htr
  - override /datamodule: iam_lines
  - override /network: null
  - override /model: lit_transformer
  - override /lr_scheduler: null
  - override /optimizer: null

tags: [lines]
epochs: &epochs 64
ignore_index: &ignore_index 3
# summary: [[1, 1, 56, 1024], [1, 89]]

logger:
  wandb:
    tags: ${tags}

criterion:
  ignore_index: *ignore_index
  # label_smoothing: 0.05

callbacks:
  stochastic_weight_averaging:
    _target_: pytorch_lightning.callbacks.StochasticWeightAveraging
    swa_epoch_start: 0.75
    swa_lrs: 1.0e-5
    annealing_epochs: 10
    annealing_strategy: cos
    device: null

optimizer:
  _target_: adan_pytorch.Adan
  lr: 3.0e-4
  betas: [0.02, 0.08, 0.01]
  weight_decay: 0.02

lr_scheduler:
  _target_: torch.optim.lr_scheduler.ReduceLROnPlateau
  mode: min
  factor: 0.8
  patience: 10
  threshold: 1.0e-4
  threshold_mode: rel
  cooldown: 0
  min_lr: 1.0e-5
  eps: 1.0e-8
  verbose: false
  interval: epoch
  monitor: val/cer

datamodule:
  batch_size: 16
  train_fraction: 0.95

network:
  _target_: text_recognizer.networks.ConvTransformer
  encoder:
    _target_: text_recognizer.networks.image_encoder.ImageEncoder
    encoder:
      _target_: text_recognizer.networks.convnext.ConvNext
      dim: 16
      dim_mults: [2, 4, 32]
      depths: [3, 3, 6]
      downsampling_factors: [[2, 2], [2, 2], [2, 2]]
      attn:
        _target_: text_recognizer.networks.convnext.TransformerBlock
        attn:
          _target_: text_recognizer.networks.convnext.Attention
          dim: &dim 512
          heads: 4
          dim_head: 64
          scale: 8
        ff:
          _target_: text_recognizer.networks.convnext.FeedForward
          dim: *dim
          mult: 2
    pixel_embedding:
      _target_: "text_recognizer.networks.transformer.embeddings.axial.\
        AxialPositionalEmbeddingImage"
      dim: *dim
      axial_shape: [7, 128]
  decoder:
    _target_: text_recognizer.networks.text_decoder.TextDecoder
    dim: *dim
    num_classes: 58
    pad_index: *ignore_index
    decoder:
      _target_: text_recognizer.networks.transformer.Decoder
      dim: *dim
      depth: 6
      block:
        _target_: "text_recognizer.networks.transformer.decoder_block.\
          DecoderBlock"
        self_attn:
          _target_: text_recognizer.networks.transformer.Attention
          dim: *dim
          num_heads: 8
          dim_head: &dim_head 64
          dropout_rate: &dropout_rate 0.2
          causal: true
        cross_attn:
          _target_: text_recognizer.networks.transformer.Attention
          dim: *dim
          num_heads: 8
          dim_head: *dim_head
          dropout_rate: *dropout_rate
          causal: false
        norm:
          _target_: text_recognizer.networks.transformer.RMSNorm
          dim: *dim
        ff:
          _target_: text_recognizer.networks.transformer.FeedForward
          dim: *dim
          dim_out: null
          expansion_factor: 2
          glu: true
          dropout_rate: *dropout_rate
      rotary_embedding:
        _target_: text_recognizer.networks.transformer.RotaryEmbedding
        dim: *dim_head

model:
  max_output_len: 89

trainer:
  gradient_clip_val: 1.0
  max_epochs: *epochs
  accumulate_grad_batches: 1
  limit_train_batches: 1.0
  limit_val_batches: 1.0
  limit_test_batches: 1.0