# Shared across all stages 
globals: &globals
  workspace_path: null  # null: defaults to current working directory (CWD)
  model_name: "unsloth/Qwen2.5-Coder-3B" 
  fim_prefix_token: "<|fim_prefix|>"
  fim_middle_token: "<|fim_middle|>"
  fim_suffix_token: "<|fim_suffix|>"
  fim_pad_token: "<|fim_pad|>"
  eos_token: "<|endoftext|>"
  label_pad_token_id: -100
  max_token_sequence_length: 1024
  data_language: "c"
  data_extensions: [".c", ".h"]
  use_unsloth: False

# Preprocess stage
preprocess:
  <<: *globals
  split_mode: "auto"
  train_ratio: 0.8  # ignored when split mode manual
  eval_ratio: 0.1  # ignored when split mode manual  
  test_ratio: 0.1  # ignored when split mode manual
  max_code_blocks_ast_depth: 4 
  min_middle_tokens_length: 8
  max_middle_tokens_length: 128 
  fim_examples_per_subblock_ratio: 1.0
  rand_to_ast_fim_examples_ratio: 0.0  # 0 = no randomly split fim examples are created
  rand_examples_min_prefix_suffix_tokens_length: 8
  tokenizer_batch_size: 32
  raw_data_path: null  # null: defaults to <workspace_path>/data
  tree_sitter_parser_path: null  # null: uses tree_sitter_language_pack parser
  tree_sitter_definitions_path: null  # null: uses package internal tree sitter definitions 
  rng_seed: 0

# Finetune stage 
finetune:
  <<: *globals
  model_attn_implementation: "sdpa"
  lora_r: 32
  lora_alpha: 64
  lora_dropout: 0.0
  lora_bias: "none"
  lora_target_modules: ["q_proj", "v_proj", "k_proj", "o_proj", "gate_proj", "down_proj", "up_proj"]
  selected_checkpoint_strategy : "best"  # "best" or "last"
  trainer_resume_from_checkpoint: null  # null->fresh start, "last"->take last checkpoint
  trainer_clear_checkpoint_dir: false
  trainer_num_train_epochs: 1
  trainer_per_device_train_batch_size: 1
  trainer_per_device_eval_batch_size: 1 
  trainer_gradient_accumulation_steps: 8 
  trainer_learning_rate: 5e-5
  trainer_weight_decay: 0.01
  trainer_max_grad_norm: 1.0
  trainer_lr_scheduler_type: "cosine"
  trainer_warmup_steps: 10
  trainer_gradient_checkpointing: true
  trainer_logging_steps: 10
  trainer_eval_strategy: "steps"
  trainer_eval_steps: 10
  trainer_save_strategy: "steps"
  trainer_save_steps: 10
  trainer_logging_strategy: "steps"
  dataset_shuffle_buffer_size: 50000
  dataset_shuffle_seed: 0

# Evaluate stage
evaluate:
  <<: *globals
  benchmark_sample_size: 250 
  benchmark_shuffle_buffer_size: 10000000
  benchmark_shuffle_seed: 42
  generation_checkpoint: "pipeline"  # "pipeline" -> selected checkpoint from finetune stage, "checkpoint-name" -> specific checkpoint
  generation_batch_size: 5 
  generation_max_new_tokens: 128  
  generation_do_sample: false
  generation_temperature: 0.1  # only active if generation_do_sample: true
  generation_top_p: 0.99  # only active if generation_do_sample: true
  codebleu_ngram_weight: 0.1
  codebleu_weighted_ngram_weight: 0.1
  codebleu_syntax_ast_weight: 0.4
  codebleu_dataflow_weight: 0.4
  sentencebleu_ngram_weight_1: 0.25
  sentencebleu_ngram_weight_2: 0.25
  sentencebleu_ngram_weight_3: 0.25
  sentencebleu_ngram_weight_4: 0.25
  line_match_number_of_lines: 2
  plot_only: false
  benchmark_use_existing_dataset: false

# Convert stage
convert:
  <<: *globals

# mlflow tracking (optional)
tracking:
  <<: *globals
  use_mlflow: false 
  mlflow_tracking_path: null   # null -> <workspace_path>/outputs/mlflow
  mlflow_experiment_name: "codefinetuner"
  mlflow_model_logging_strategy: "none"  # "none", "adapter", "gguf", or "all"