datasets:
  - name: chatalpaca_predictive_state
    repo_id: BRlkl/chatalpaca-multiturn-predictive-state3
    source_split: train
    format: predictive_state_messages
    validation_ratio: 0.02
    split_seed: 17
    min_turns: 2
    max_turns: 6
    max_message_chars: 6000
    use_base_chat_template: true

model:
  base_model_name: google/t5gemma-l-l-prefixlm-it
  initial_model_path: BRlkl/test
  dtype: bfloat16
  attn_implementation: sdpa
  magicnorm_eps: 1.0e-6
  z_slots: 256
  num_time_tokens: 0
  use_explicit_time_features: false
  gate_attention_heads: 4
  max_observation_tokens: 1024
  max_decoder_tokens: 1024
  thought_loop_proposal_mode: observation_hidden_compression
  preserve_observation_encoder_manifold: true
  observation_encoder_use_state_context: true
  latent_attention_mask_mode: full
  initial_update_gate_bias: 20.0

training:
  seed: 17
  num_workers: 2
  gradient_checkpointing: true
  mixed_precision: bf16
  max_grad_norm: 1.0
  weight_decay: 0.01
  backbone_learning_rate: 5.0e-6
  new_module_learning_rate: 1.0e-4
  adam_beta1: 0.9
  adam_beta2: 0.95
  adam_epsilon: 1.0e-8
  fused_adamw: true
  freeze_gate_head: true
  log_every_steps: 1
  eval_every_steps: 50
  checkpoint_every_steps: 5000
  eval_max_batches: 16
  max_train_examples: 10
  max_validation_examples: 10

  freeze_decoder_first_fraction: 0.0

  response_loss_weight: 0.45
  probe_loss_weight: 0.35
  probe_hard_loss_weight: 0.15
  noop_probe_loss_weight: 0.05
  probe_hard_top_fraction: 0.20

  max_probe_branches_per_step: 8
  max_probe_noop_ticks: 4
  require_probe_steps: 1
  drop_llm_not_stated_probes: true
  drop_llm_positive_answer_not_in_prefix: true
  drop_llm_future_answer_leaks: true

  validation_behavior_enabled: true
  validation_behavior_exact_only: true
  validation_behavior_max_batches: 4
  validation_behavior_max_probe_generations: 192
  validation_probe_max_new_tokens: 96
  validation_probe_extra_new_tokens: 8
  validation_noop_survival_threshold_ticks: 64
  validation_self_fed_behavior_enabled: true
  validation_behavior_self_fed_max_batches: 2
  validation_behavior_self_fed_max_probe_generations: 96
  wandb_train_metric_keys:
    - train/loss_total
    - train/loss_probe
    - train/probe_first_token_exact_match
  wandb_validation_metric_keys:
    - loss_total
    - validation_goal_loss
    - behavior/probe_exact_match
    - behavior/self_fed/probe_exact_match
    - behavior/noop_survival_exact
    - behavior/secret_word_exact
    - behavior/entity_binding_exact
    - behavior/correction_latest_state_accuracy
    - behavior/not_stated_accuracy

  checkpoint_selection_metric: validation_goal_loss
  checkpoint_selection_mode: min
  validation_probe_miss_penalty: 2.0
  validation_noop_miss_penalty: 1.0
  validation_self_fed_miss_penalty: 1.0
  validation_not_stated_miss_penalty: 0.5
  validation_entity_miss_penalty: 0.75
  validation_correction_miss_penalty: 0.75
  checkpoint_min_probe_exact_match: 0.50
  checkpoint_min_secret_word_exact: 0.50
  checkpoint_min_entity_binding_exact: 0.50
  checkpoint_min_not_stated_accuracy: 0.50
  checkpoint_min_correction_latest_state_accuracy: 0.50
  checkpoint_min_noop_survival_exact: 0.50
  checkpoint_min_self_fed_probe_accuracy: 0.40

  self_feedback_start_ratio: 0.0
  self_feedback_end_ratio: 0.15
  self_feedback_start_fraction: 0.65
  self_feedback_end_fraction: 0.95
  self_feedback_seed: 17
  self_feedback_max_new_tokens: 160
  self_feedback_extra_new_tokens: 8

  student_observation_seed: 17
  student_prev_assistant_dropout_prob: 0.0
  student_prev_assistant_keep_last_chars:

phase:
  micro_batch_size: 16
  eval_batch_size: 16
  gradient_accumulation_steps: 1
  num_train_epochs: 100
  warmup_ratio: 0.03
  shuffle_train: true

cache:
  preprocessed_root: cache/preprocessed_pre_sft_multiturn_predictive

paths:
  run_root: runs_pre_sft_multiturn_predictive
  export_root: exports_multiturn_predictive

inference:
  format: predictive_state_multiturn
  use_base_chat_template: true

wandb:
  enabled: true
  project: samantha-pre-sft
  run_name: t5gemma2-thoughtloop-pre-sft-predictive-state

hub:
  model_repo_id: BRlkl/test_multiturn_predictive_state_overfit
  private: false