config_for_2b.yaml

model:
  type: Vishwamai
  size: 2b
  layers: 24
  hidden_size: 2048
  attention_heads: 16
  dropout: 0.1

training:
  batch_size: 32
  learning_rate: 3e-5
  epochs: 10
  optimizer: AdamW
  scheduler: linear

data:
  train_dataset: datasets/train.json
  validation_dataset: datasets/dev.json

evaluation:
  metrics:
    - name: Perplexity
    - name: BLEU
    - name: ROUGE