#!/bin/bash
# Installation script for chest2vec dependencies
# This script installs PyTorch and flash-attention with the correct versions

set -e

echo "Installing PyTorch packages from custom index..."
pip install torch==2.6.0 torchvision==0.21.0 torchaudio==2.6.0 --index-url https://download.pytorch.org/whl/cu126

echo "Installing flash-attention from GitHub release..."
pip install https://github.com/Dao-AILab/flash-attention/releases/download/v2.8.3/flash_attn-2.8.3+cu12torch2.6cxx11abiTRUE-cp310-cp310-linux_x86_64.whl

echo "Installing chest2vec package..."
pip install chest2vec

echo "Installation complete!"