File size: 4,184 Bytes
6980f6d | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 | #!/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 "============================================================" |