# Example configuration for RL training (RL stage assumes you have a reward model)

model:
  architecture_type: transformer
  vocab_size: 50257
  hidden_dim: 256
  num_layers: 8
  num_heads: 8
  dropout: 0.1
  max_seq_length: 512
  init_std: 0.02

dataset:
  dataset_name: allenai/real_toxicity_prompts
  split: train
  prompt_column: text
  max_samples: 2000
  cache_dir: .cache/datasets
  tokenizer_threads: 1  # Number of background threads for tokenization (1-32 recommended)

batch_size: 4
num_epochs: 1
gradient_accumulation_steps: 8
max_grad_norm: 0.5

optimizer:
  optimizer_type: adamw
  learning_rate: 1e-5
  weight_decay: 0.0

scheduler:
  scheduler_type: linearWarmup
  warmup_steps: 50

dtype: bfloat16
device: cuda

checkpoint_dir: checkpoints/rl
save_every_steps: 100
save_best_model: false
keep_last_n_checkpoints: 2

eval_every_steps: 100
eval_samples: 100

log_every_steps: 10
aim_repo: .aim

# RL-specific settings
rl_method: ppo  # or "dpo"
reward_model_path: checkpoints/reward_model.pt  # Path to your reward model
ppo_epochs: 4
ppo_clip_ratio: 0.2
entropy_coeff: 0.01
value_loss_coeff: 1.0
generation_max_length: 256

seed: 42
num_workers: 0
pin_memory: true