#!/bin/bash
echo "Starting building the environment..."
uv sync --offline
echo "Activated virtual environment"
source .venv/bin/activate
echo "Copying files..."
cp debug/utils.py .venv/lib/python3.10/site-packages/transformers/generation/utils.py
cp debug/modeling_llama.py .venv/lib/python3.10/site-packages/transformers/models/llama/modeling_llama.py

if [ -d "$HOME/Llama-2-7b-hf" ]; then
    export LLM_PATH="$HOME/Llama-2-7b-hf"
else
    export LLM_PATH="meta-llama/Llama-2-7b-hf"
fi
echo "LLM_PATH: $LLM_PATH"

if [ -n "$LD_LIBRARY_PATH" ]; then
    echo "LD_LIBRARY_PATH is already set. Skipping CUDA setup..."
else
    echo "Setting up CUDA..."
fi
cuda_dirs=(/usr/local/cuda-*/)
if [ ${#cuda_dirs[@]} -eq 0 ]; then

    echo "Can not find any cuda folder in /usr/local"
    exit 1

    # Sort the cuda folders and get the latest one
    latest_cuda=$(for dir in "${cuda_dirs[@]}"; do echo "${dir%/}"; done | sort -V | tail -n 1)

    echo "The latest CUDA folder is: ${latest_cuda}"

    export CUDA_HOME="${latest_cuda}"
    export PATH="${latest_cuda}/bin:${PATH}"
    export LD_LIBRARY_PATH="${latest_cuda}/lib64/stubs:${latest_cuda}/lib64:${latest_cuda}/cudnn/lib:${LD_LIBRARY_PATH}"

fi

echo "Starting training..."
sh train_movielens.sh