# @package _global_
defaults:
  - /base_config
  - override /trainer: quick
  - override /trainer/optim: adam_with_one_cycle
  - override /model/sent/encoder: lstm
  - override /model/fragment: transformer
  - override /model/doc: transformer
  - _self_

text_proc:
  fragment_size: 16
  tokenizer:
    add_bos: false
    vocab_path: /train/sentencepiece.bpe.model

trainer:
  switch_tasks_every: 10
  emb_grad_scale: 20
  print_gpu_memory_stat_every: 1000
  tasks:
    - DOC_RETR
    - SENT_RETR

  optim:
    max_grad_norm: 3.0
    optim_kind: ADAMW
    use_zero_optim: false
    lr_scheduler_kwargs:
      div_factor: 10
      final_div_factor: 1
    emb:
      lr: 0.0005
    sent:
      lr: 0.0009
    fragment:
      lr: 0.00006
    doc:
      lr: 0.00005
model:
  freeze_base_sents_layer: false
  max_chunk_size: 2048
  max_tokens_in_chunk: 64000
  scale: 40
  margin: 0.2
  sent:
    scale: 30
    margin: 0.3
    max_chunk_size: 1024
  fragment:
    add_beg_seq_token: true
    num_layers: 2
    full_intermediate: true
    hidden_size: 768
    share_attn: false
  doc:
    add_beg_seq_token: true
    num_layers: 3
    full_intermediate: true
    hidden_size: 768

batches:
  reinit_last_iter: false
  sents_batch_iterator_conf:
    async_generators: 2
    batch_generator_conf:
      max_sent_size: 128
      batch_size: 2000
      dont_use_hns: true


  docs_batch_iterator_conf:
    async_generators: 6
    batch_generator_conf:
      input_dir: /train/data/docs
      batch_total_sents_cnt: 36000
      max_sents_per_doc: 1024
