#!/bin/bash
# AutoDL 环境准备：拉数据 + 下底模 + 注册数据集 + 生成训练配置
# 用法：cd /root/autodl-tmp && bash hv-setup.sh
set -e

BASE="http://8.148.217.89:18642/apks/x7k2m9"
WORK="/root/autodl-tmp/hv"
mkdir -p "$WORK" && cd "$WORK"

echo "== [1/4] 下载 sft 数据包 =="
[ -f hv-sft-data.zip ] || wget -q --show-progress -O hv-sft-data.zip "$BASE/hv-sft-data.zip"
python3 - <<'PY'
import zipfile
zipfile.ZipFile('/root/autodl-tmp/hv/hv-sft-data.zip').extractall('/root/autodl-tmp/hv')
print('extracted -> /root/autodl-tmp/hv/sft')
PY
ls -la "$WORK/sft/"

echo "== [2/4] 下载底模 Qwen3.5-9B（魔搭） =="
pip install -q modelscope 2>&1 | tail -1 || true
python3 - <<'PY'
from modelscope import snapshot_download
p = snapshot_download('Qwen/Qwen3.5-9B', cache_dir='/root/autodl-tmp/models')
open('/root/autodl-tmp/hv/model_dir.txt', 'w').write(p)
print('MODEL_DIR=' + p)
PY

echo "== [3/4] 注册数据集到 LLaMA-Factory =="
python3 - <<'PY'
import json, os, shutil
import llamafactory
lf = os.path.dirname(os.path.dirname(os.path.abspath(llamafactory.__file__)))
print('LLAMA_FACTORY_DIR=' + lf)
src = '/root/autodl-tmp/hv/sft'
dst = os.path.join(lf, 'data')
for f in ['sharegpt_sys.json', 'sharegpt_nosys.json']:
    shutil.copy(os.path.join(src, f), dst)
info_path = os.path.join(dst, 'dataset_info.json')
info = json.load(open(info_path, encoding='utf-8'))
snippet = json.load(open(os.path.join(src, 'dataset_info_snippet.json'), encoding='utf-8'))
info.update(snippet)
json.dump(info, open(info_path, 'w', encoding='utf-8'), ensure_ascii=False, indent=2)
open('/root/autodl-tmp/hv/lf_dir.txt', 'w').write(lf)
print('registered datasets:', list(snippet.keys()))
PY

echo "== [4/4] 生成训练配置 train.yaml =="
LF=$(cat "$WORK/lf_dir.txt")
MODEL=$(cat "$WORK/model_dir.txt")
sed -e "s#__LF_DATA__#$LF/data#" -e "s#__MODEL_DIR__#$MODEL#" "$WORK/hv-train.yaml" > "$WORK/train.yaml"
echo "--- train.yaml ---"
cat "$WORK/train.yaml"

echo
echo "自检："
nvidia-smi -L 2>/dev/null || echo "  (无卡模式：看不到 GPU，正常，付费用时再切 GPU 模式)"
python3 -c "import torch;print('  torch', torch.__version__, '| cuda', torch.version.cuda)" 2>/dev/null || echo "  torch 未就绪"
echo
echo "下一步：bash /root/autodl-tmp/hv/hv-run.sh"