filipelopesmedbr/icd11-llm-ministral-8b

# ============================================================
# Configuração Axolotl para finetuning do modelo Mistral 8B
# (arquivo: axolotl_2_a40_runpod_config.yaml)
# ============================================================

# ============================================================
# 1. Modelo Base e Configurações Gerais
# ------------------------------------------------------------
base_model: mistralai/Ministral-8B-Instruct-2410

# ------------------------------------------------------------
# 2. Quantização / dtype
#   - load_in_4bit: false (vamos usar LoRA puro)
#   - load_in_8bit: false
#   - bf16: true    (carregar em bfloat16 para A40)
#   - fp16: false
# ------------------------------------------------------------
load_in_8bit: false
load_in_4bit: false         
bf16: true
fp16: false

# ------------------------------------------------------------
# 3. Chat template e tokens especiais
# ------------------------------------------------------------
chat_template: mistral_v2v3
special_tokens:
  bos_token: "<s>"
  eos_token: "</s>"

# ============================================================
# 4. Paths dos Datasets (cada split numa linha separada)
#  
# Atenção: cada bloco dentro de `datasets:` deve ter um "path: <string>"
# e um "split: train" ou "split: validation".  
# ============================================================
datasets:
  - path: dataset
    split: train
    type: alpaca
    field_instruction: instruction
    field_input: input
    field_output: output
    field_weight: weight

validation_datasets:
  - path: dataset
    split: validation
    type: alpaca
    field_instruction: instruction
    field_input: input
    field_output: output

# (Opcional) Se quiser um split separado de test para avaliação final,
# basta descomentar este bloco e criar “datasets/test.jsonl”.
test_datasets:
  - path: dataset
    split: test
    type: alpaca
    field_instruction: instruction
    field_input: input
    field_output: output

# ============================================================
# 5. Configuração de Treino LoRA
# ------------------------------------------------------------
adapter: lora
lora_r: 8
lora_alpha: 16
lora_dropout: 0.05
lora_target_modules:
  - gate_proj
  - down_proj
  - up_proj
  - q_proj
  - v_proj
  - k_proj
  - o_proj

# Módulos LoRA a salvar para inferência leve
lora_modules_to_save:
  - embed_tokens
  - lm_head

# ============================================================
# 6. Batch / Gradiente / Checkpointing
# ------------------------------------------------------------
micro_batch_size: 20
gradient_accumulation_steps: 4
gradient_checkpointing: true

# ============================================================
# 7. Otimizador / Scheduler
# ------------------------------------------------------------
learning_rate: 0.0001
optimizer: adamw_bnb_8bit
lr_scheduler: cosine

# ============================================================
# 8. Épocas / Sequence Len / Packing
# ------------------------------------------------------------
num_epochs: 3
sequence_len: 1536
pad_to_sequence_len: false
sample_packing: true

# ============================================================
# 9. Log / Checkpoints / Avaliações
# ------------------------------------------------------------
logging_steps: 40
evals_per_epoch: 4
save_per_epoch: 4

# ============================================================
# 10. Diretórios de Saída
# ------------------------------------------------------------
output_dir: ./outputs/cid11-agent-mistral-8b
save_safetensors: true

# ============================================================
# 11. Atenção / Configurações Adicionais
# ------------------------------------------------------------
device_map: auto
attn_implementation: flash_attention_2
flash_attention: false

resume_from_checkpoint: null
save_total_limit: 3

# ============================================================
# 12. Métricas de Avaliação Personalizadas
# ------------------------------------------------------------
# O Axolotl monitora internamente “eval_loss”.  
# Se quiser acrescentar “eval_accuracy” (por ex., comparar
# o conteúdo dentro de <code>…</code>), tem que fazer por fora.
#
# Basta manter estas duas linhas (não removê-las), mesmo que
# você tenha retirado o bloco callbacks. Depois, rodar um script
# externo que leia o arquivo “outputs/predictions.jsonl” e 
# compare com “datasets/validation.jsonl” para obter a acurácia.
monitor: eval_loss

# — Forçar o Axolotl a gerar predições em cada validação ——
predict_with_generate: true

# — Se quiser salvar essas predições no disco, deixe true ——
save_preds: true

# ============================================================
# 13. (Não há mais callback_modules nem callbacks aqui)
# ============================================================

# ============================================================
# 14. Estrutura esperada de pastas no repositório:
#
# CIEL-Clinical-Concepts-to-ICD-11/
# ├── axolotl_2_a40_runpod_config.yaml   ← este arquivo
# ├── datasets/
# │   ├── train.jsonl
# │   ├── validation.jsonl
# │   └── test.jsonl   (opcional se você quisesse finetune+teste juntos)
# ├── scripts/
# │   └── custom_metrics_callback.py     (se for usar callback externamente)
# ├── outputs/                           (aqui ficarão checkpoints e predictions.jsonl)
# └── train.py  (opcional, não é mandatório para Axolotl CLI)
# ============================================================