#!/bin/bash
# Install unsloth for Qwen3-8B LoRA fine-tuning on GTX 1080 (Pascal, CUDA 12.x)
# Uses GPU 1 (GTX 1080, 8GB VRAM) — leaves GPU 0 free for inference
# Run once. Safe to re-run.

set -e

echo "=== Unsloth install for Grace fine-tuning (GPU 1, GTX 1080) ==="
echo "CUDA: $(nvcc --version 2>/dev/null | grep release || echo 'checking docker...')"
echo ""

# Use a dedicated venv to avoid polluting system Python
VENV_DIR="$HOME/unsloth-env"

if [ ! -d "$VENV_DIR" ]; then
    echo "[1/5] Creating venv at $VENV_DIR..."
    python3 -m venv "$VENV_DIR"
else
    echo "[1/5] Venv already exists at $VENV_DIR"
fi

source "$VENV_DIR/bin/activate"

echo "[2/5] Installing PyTorch with CUDA 12.1 (Pascal-compatible)..."
pip install --quiet torch torchvision torchaudio --index-url https://download.pytorch.org/whl/cu121

echo "[3/5] Installing unsloth (Qwen3 + LoRA support)..."
pip install --quiet "unsloth[colab-new] @ git+https://github.com/unslothai/unsloth.git"

echo "[4/5] Installing training utilities..."
pip install --quiet \
    transformers \
    datasets \
    trl \
    peft \
    accelerate \
    bitsandbytes \
    wandb \
    scipy

echo "[5/5] Verifying install..."
python3 -c "
import torch
print(f'PyTorch: {torch.__version__}')
print(f'CUDA available: {torch.cuda.is_available()}')
if torch.cuda.is_available():
    print(f'GPU 0: {torch.cuda.get_device_name(0)}')
    print(f'GPU 1: {torch.cuda.get_device_name(1)}')
import unsloth
print(f'Unsloth: OK')
"

echo ""
echo "=== Install complete ==="
echo "Activate with: source $VENV_DIR/bin/activate"
echo "Run fine-tuning with: ~/scripts/finetune-lora.py"