BAAI/Brainmu-Spike

1#!/usr/bin/env bash2set -euo pipefail3cd "$(dirname "$0")/.."4command -v uv >/dev/null || { echo "Install uv first: https://docs.astral.sh/uv/getting-started/installation/"; exit 2; }5ENV="${BRAINMU_ENV:-$PWD/.venv}"6[[ ! -e "$ENV" ]] || { echo "Environment already exists: $ENV (not overwritten)"; exit 2; }7uv venv --python 3.12 "$ENV"8PY="$ENV/bin/python"9uv pip install --python "$PY" torch==2.11.0 torchvision==0.26.0 --index-url https://download.pytorch.org/whl/cu13010uv pip install --python "$PY" -r requirements.txt packaging ninja wheel setuptools11# No build isolation: compile against the torch installed above.12MAX_JOBS="${MAX_JOBS:-4}" uv pip install --python "$PY" flash-attn==2.8.3.post1 --no-build-isolation13"$PY" -c 'import torch, torchvision, flash_attn; print("Environment ready:", torch.__version__, torchvision.__version__, flash_attn.__version__)'14echo "Run: bash scripts/test.sh"15