#!/bin/bash
# ======================================
# 之之母模型 v2.0 · 全自动训练脚本
# 由铸渊生成 · 2026-05-18
# 用法：bash run_zhizhi_v2.sh
# ======================================

set -e

TRAIN_DIR="/root/autodl-tmp"
DATA_FILE="$TRAIN_DIR/sft_v2_nosys.jsonl"
YAML_FILE="$TRAIN_DIR/zhizhi_v2.yaml"
LOG_FILE="$TRAIN_DIR/train_zhizhi_v2.log"
OUTPUT_DIR="$TRAIN_DIR/output/zhizhi_v2"

echo "========================================"
echo "  之之母模型 v2.0 · 训练启动"
echo "  铸渊配置 · 2026-05-18"
echo "========================================"
echo ""

# 1. 检查数据
if [ ! -f "$DATA_FILE" ]; then
    echo "❌ 数据文件不存在: $DATA_FILE"
    exit 1
fi
DATA_LINES=$(wc -l < "$DATA_FILE")
echo "✅ 数据: $DATA_FILE ($DATA_LINES 条)"
echo ""

# 2. 创建配置文件
echo "📝 写入配置..."

cat > "$TRAIN_DIR/dataset_info.json" << 'JSONEOF'
{"zhizhi_v2":{"file_name":"/root/autodl-tmp/sft_v2_nosys.jsonl","formatting":"sharegpt","columns":{"messages":"messages"}}}
JSONEOF

echo "  ✅ dataset_info.json"

cat > "$TRAIN_DIR/ds_config.json" << 'JSONEOF'
{"train_batch_size":"auto","train_micro_batch_size_per_gpu":1,"gradient_accumulation_steps":8,"optimizer":{"type":"AdamW","params":{"lr":5e-6,"betas":[0.9,0.999],"eps":1e-8,"weight_decay":0.01}},"scheduler":{"type":"WarmupLR","params":{"warmup_min_lr":0,"warmup_max_lr":5e-6,"warmup_num_steps":10}},"zero_optimization":{"stage":2,"offload_optimizer":{"device":"cpu","pin_memory":true}},"gradient_clipping":1.0,"bf16":{"enabled":"auto"}}
JSONEOF

echo "  ✅ ds_config.json"

cat > "$YAML_FILE" << 'YAMLEOF'
model_name_or_path: Qwen/Qwen2.5-7B
dataset: zhizhi_v2
dataset_dir: /root/autodl-tmp
template: qwen
cutoff_len: 4096
finetuning_type: full
per_device_train_batch_size: 1
gradient_accumulation_steps: 8
learning_rate: 5.0e-6
num_train_epochs: 3.0
lr_scheduler_type: cosine
warmup_ratio: 0.03
optim: adamw_torch
bf16: true
ddp_find_unused_parameters: false
deepspeed: /root/autodl-tmp/ds_config.json
output_dir: /root/autodl-tmp/output/zhizhi_v2
logging_steps: 5
save_steps: 58
plot_loss: true
overwrite_output_dir: true
preprocessing_num_workers: 16
report_to: none
YAMLEOF

echo "  ✅ zhizhi_v2.yaml"
echo ""

# 3. 检查/安装 LLaMA-Factory
echo "🔍 检查 LLaMA-Factory..."
if python3 -c "import llamafactory; print(llamafactory.__version__)" 2>/dev/null; then
    echo "✅ LLaMA-Factory 已安装"
else
    echo "⏳ 安装 LLaMA-Factory 0.8.3..."
    pip install "llamafactory==0.8.3" -q
    echo "✅ 安装完成"
fi
echo ""

# 4. 检查是否有模型缓存
echo "🔍 检查模型 Qwen2.5-7B..."
python3 -c "from transformers import AutoConfig; AutoConfig.from_pretrained('Qwen/Qwen2.5-7B')" 2>/dev/null && \
    echo "✅ 模型可用" || \
    echo "⏳ 首次运行会自动下载模型（约15GB）"
echo ""

# 5. dry_run验证
echo "🔍 配置验证中..."
cd "$TRAIN_DIR"
export CUDA_VISIBLE_DEVICES=0,1,2,3
DRY_RESULT=$(llamafactory-cli train zhizhi_v2.yaml --dry_run 2>&1)
if echo "$DRY_RESULT" | grep -qE "(Error|error|Traceback)"; then
    echo "❌ 配置验证失败:"
    echo "$DRY_RESULT"
    echo ""
    echo "请将上面的报错信息发给铸渊"
    exit 1
fi
echo "✅ 配置验证通过"
echo ""

# 6. 启动训练
echo "🚀 启动训练..."
echo "  日志: $LOG_FILE"
echo "  输出: $OUTPUT_DIR"
echo "  预计: 5-10分钟"
echo ""

nohup llamafactory-cli train zhizhi_v2.yaml > "$LOG_FILE" 2>&1 &
PID=$!
echo "📌 进程ID: $PID"
echo ""

# 7. 等待并显示进度
echo "📊 训练日志（等待10秒后开始输出，按 Ctrl+C 停止查看但训练继续）:"
echo ""
sleep 10
tail -30 "$LOG_FILE" 2>/dev/null

echo ""
echo "========================================"
echo "  训练已在后台运行 (PID=$PID)"
echo ""
echo "  查看实时日志: tail -f $LOG_FILE"
echo "  查看进程:     ps aux | grep llamafactory"
echo "  停止训练:     kill $PID"
echo "========================================"