#!/bin/bash
set -e

# 在容器启动后，首先从 requirements.txt 安装所有依赖包
# pip install -r requirements.txt

# bash install.sh -y
python3 -m pip install --upgrade pip
pip install uv -i https://pypi.tuna.tsinghua.edu.cn/simple

# 切换到项目目录
cd /ycz/Minimind

# 检查并修复虚拟环境
if [ ! -f .venv/bin/python ] || [ ! -x .venv/bin/python ]; then
    echo "Virtual environment is broken or missing, recreating with uv..."
    rm -rf .venv
    uv venv .venv
fi

# 不要手动激活虚拟环境，让uv自动管理
# . ./.venv/bin/activate

# 使用uv同步依赖
uv sync

# 安装完成后，执行主训练脚本
# "$@" 会将 experiment.yaml 中 entrypoint 定义的参数传递给 python 脚本
CUDA_VISIBLE_DEVICES=0 uv run python -m accelerate.commands.launch \
    --num_processes=1 \
    --mixed_precision=bf16 \
    --main_process_port=29500 \
    train_pretrain_accelerate.py "$@"