ankahi / scripts /00_env_setup.sh

Add files using upload-large-folder tool

6980f6d verified 6 months ago

4.18 kB

#!/usr/bin/env bash
# scripts/00_env_setup.sh (CUDA 12.1 variant)
# One-shot setup for Ankahi training environment on H100.
# Uses cu121 wheels for better compatibility with newer torchao/unsloth.
# Usage: bash scripts/00_env_setup.sh
set -euo pipefail

echo "============================================================"
echo "  ANKAHI — Environment Setup (CUDA 12.1)"
echo "============================================================"

# ── 1. System checks ─────────────────────────────────────────────
echo "[1/8] Checking system..."
python_version=$(python3 --version 2>&1)
echo "  Python: $python_version"

nvidia-smi --query-gpu=name,memory.total,driver_version --format=csv,noheader || {
    echo "ERROR: nvidia-smi failed — is CUDA available?"
    exit 1
}

nvcc_version=$(nvcc --version 2>&1 | grep "release" | awk '{print $5}' | tr -d ',')
echo "  CUDA: $nvcc_version"

# ── 2. Upgrade pip ───────────────────────────────────────────────
echo "[2/8] Upgrading pip..."
python3 -m pip install --upgrade pip setuptools wheel --quiet

# ── 3. PyTorch (CUDA 12.4 wheels) ───────────────────────────────
echo "[3/8] Installing PyTorch 2.6.0 with CUDA 12.4..."
pip install torch==2.6.0 torchvision==0.21.0 torchaudio==2.6.0 \
    --index-url https://download.pytorch.org/whl/cu124 --quiet

# Verify
python3 -c "import torch; assert torch.cuda.is_available(), 'CUDA not available!'; print(f'  torch {torch.__version__}, CUDA {torch.version.cuda}, device: {torch.cuda.get_device_name(0)}')"

# ── 4. Flash-Attention ────────────────────────────────────────────
echo "[4/8] Installing flash-attn..."
pip install flash-attn>=2.6.0 --no-build-isolation --quiet

# ── 5. Core HuggingFace stack ─────────────────────────────────────
echo "[5/8] Installing HuggingFace stack..."
pip install \
    "transformers==4.57.1" \
    "torchao==0.9.0" \
    "datasets>=3.4.1,<4.4.0" \
    "peft>=0.12.0" \
    "trl>=0.18.2,<=0.24.0" \
    "accelerate>=0.30.0" \
    "bitsandbytes>=0.43.0" \
    wandb \
    --quiet

# ── 6. Unsloth ────────────────────────────────────────────────────
echo "[6/8] Installing Unsloth (Gemma 4 support)..."
pip install "unsloth[colab-new] @ git+https://github.com/unslothai/unsloth.git" 

python3 -c "from unsloth import FastVisionModel; print('  Unsloth FastVisionModel OK')"

# ── 7. Audio, TTS, and data deps ─────────────────────────────────
echo "[7/8] Installing audio/TTS/data dependencies..."
# numpy pinned to 1.26.4 for f5-tts compatibility
pip install \
    "numpy<=1.26.4" \
    librosa soundfile ffmpeg-python audiomentations \
    TTS \
    sacrebleu \
    "sentence-transformers>=3.0.0" \
    pandas Pillow pyarrow requests tqdm pyyaml python-dotenv \
    anthropic "google-generativeai>=0.7.0" \
    pytest pytest-cov black ruff \
    --quiet

# AI4Bharat Indic-TTS (clone if not present)
if [ ! -d "vendor/Indic-TTS" ]; then
    echo "  Cloning AI4Bharat Indic-TTS..."
    mkdir -p vendor
    git clone --depth 1 https://github.com/AI4Bharat/Indic-TTS vendor/Indic-TTS --quiet
    pip install -e vendor/Indic-TTS --quiet || echo "  (Indic-TTS install needs manual steps — see vendor/Indic-TTS/README)"
fi

# ── 8. Install ankahi package itself (editable) ───────────────────
echo "[8/8] Installing ankahi in editable mode..."
pip install -e . --quiet

echo ""
echo "============================================================"
echo "  Setup complete. Run sanity check next:"
echo "  python scripts/07_overfit_sanity_check.py"
echo "============================================================"

	#!/usr/bin/env bash
	# scripts/00_env_setup.sh (CUDA 12.1 variant)
	# One-shot setup for Ankahi training environment on H100.
	# Uses cu121 wheels for better compatibility with newer torchao/unsloth.
	# Usage: bash scripts/00_env_setup.sh
	set -euo pipefail

	echo "============================================================"
	echo " ANKAHI — Environment Setup (CUDA 12.1)"
	echo "============================================================"

	# ── 1. System checks ─────────────────────────────────────────────
	echo "[1/8] Checking system..."
	python_version=$(python3 --version 2>&1)
	echo " Python: $python_version"

	nvidia-smi --query-gpu=name,memory.total,driver_version --format=csv,noheader \|\| {
	echo "ERROR: nvidia-smi failed — is CUDA available?"
	exit 1
	}

	nvcc_version=$(nvcc --version 2>&1 \| grep "release" \| awk '{print $5}' \| tr -d ',')
	echo " CUDA: $nvcc_version"

	# ── 2. Upgrade pip ───────────────────────────────────────────────
	echo "[2/8] Upgrading pip..."
	python3 -m pip install --upgrade pip setuptools wheel --quiet

	# ── 3. PyTorch (CUDA 12.4 wheels) ───────────────────────────────
	echo "[3/8] Installing PyTorch 2.6.0 with CUDA 12.4..."
	pip install torch==2.6.0 torchvision==0.21.0 torchaudio==2.6.0 \
	--index-url https://download.pytorch.org/whl/cu124 --quiet

	# Verify
	python3 -c "import torch; assert torch.cuda.is_available(), 'CUDA not available!'; print(f' torch {torch.__version__}, CUDA {torch.version.cuda}, device: {torch.cuda.get_device_name(0)}')"

	# ── 4. Flash-Attention ────────────────────────────────────────────
	echo "[4/8] Installing flash-attn..."
	pip install flash-attn>=2.6.0 --no-build-isolation --quiet

	# ── 5. Core HuggingFace stack ─────────────────────────────────────
	echo "[5/8] Installing HuggingFace stack..."
	pip install \
	"transformers==4.57.1" \
	"torchao==0.9.0" \
	"datasets>=3.4.1,<4.4.0" \
	"peft>=0.12.0" \
	"trl>=0.18.2,<=0.24.0" \
	"accelerate>=0.30.0" \
	"bitsandbytes>=0.43.0" \
	wandb \
	--quiet

	# ── 6. Unsloth ────────────────────────────────────────────────────
	echo "[6/8] Installing Unsloth (Gemma 4 support)..."
	pip install "unsloth[colab-new] @ git+https://github.com/unslothai/unsloth.git"

	python3 -c "from unsloth import FastVisionModel; print(' Unsloth FastVisionModel OK')"

	# ── 7. Audio, TTS, and data deps ─────────────────────────────────
	echo "[7/8] Installing audio/TTS/data dependencies..."
	# numpy pinned to 1.26.4 for f5-tts compatibility
	pip install \
	"numpy<=1.26.4" \
	librosa soundfile ffmpeg-python audiomentations \
	TTS \
	sacrebleu \
	"sentence-transformers>=3.0.0" \
	pandas Pillow pyarrow requests tqdm pyyaml python-dotenv \
	anthropic "google-generativeai>=0.7.0" \
	pytest pytest-cov black ruff \
	--quiet

	# AI4Bharat Indic-TTS (clone if not present)
	if [ ! -d "vendor/Indic-TTS" ]; then
	echo " Cloning AI4Bharat Indic-TTS..."
	mkdir -p vendor
	git clone --depth 1 https://github.com/AI4Bharat/Indic-TTS vendor/Indic-TTS --quiet
	pip install -e vendor/Indic-TTS --quiet \|\| echo " (Indic-TTS install needs manual steps — see vendor/Indic-TTS/README)"
	fi

	# ── 8. Install ankahi package itself (editable) ───────────────────
	echo "[8/8] Installing ankahi in editable mode..."
	pip install -e . --quiet

	echo ""
	echo "============================================================"
	echo " Setup complete. Run sanity check next:"
	echo " python scripts/07_overfit_sanity_check.py"
	echo "============================================================"