#!/usr/bin/env bash
# Treina um LoRA do Llama-3.1-8B-Instruct com os atendimentos exportados (estágio 2),
# usando MLX no Apple Silicon. Ajustado para 16 GB de RAM (M4): modelo 4-bit, batch 1.
#
# Pré-requisito: dataset exportado (backend: node scripts/exportarDatasetTreinamento.js)
# Uso: ./treinar_lora_mlx.sh [epocas]   (default 3)
set -euo pipefail
cd "$(dirname "$0")"

MODEL="mlx-community/Meta-Llama-3.1-8B-Instruct-4bit"
DATA_DIR="data"
ADAPTER_DIR="adapters"
EPOCAS="${1:-3}"

if [[ ! -f "$DATA_DIR/train.jsonl" ]]; then
  echo "ERRO: $DATA_DIR/train.jsonl não existe. Rode antes: cd .. && node scripts/exportarDatasetTreinamento.js" >&2
  exit 1
fi

if [[ ! -d .venv ]]; then
  python3 -m venv .venv
fi
source .venv/bin/activate
pip -q install --upgrade mlx-lm

N_TRAIN=$(wc -l < "$DATA_DIR/train.jsonl" | tr -d ' ')
ITERS=$(( N_TRAIN * EPOCAS ))
if (( ITERS < 100 )); then ITERS=100; fi

echo "Treinando: $N_TRAIN amostras, $EPOCAS épocas => $ITERS iterações (batch 1)"

mlx_lm.lora \
  --model "$MODEL" \
  --train \
  --data "$DATA_DIR" \
  --adapter-path "$ADAPTER_DIR" \
  --batch-size 1 \
  --num-layers 16 \
  --iters "$ITERS" \
  --max-seq-length 4096 \
  --grad-checkpoint \
  --save-every 100 \
  --steps-per-eval 100

echo
echo "Adapter salvo em $ADAPTER_DIR/. Teste rápido:"
echo "  mlx_lm.generate --model $MODEL --adapter-path $ADAPTER_DIR --prompt 'Boa tarde, minha internet caiu'"
echo "Para publicar no Ollama: ./publicar_no_ollama.sh"