# Load .env
cd /workspace
sed -i '/^KVA_STORAGE=/d' .env
echo 'KVA_STORAGE=/workspace/kva_store' >> .env
export $(grep -v '^#' /workspace/.env | xargs)
echo "export $(grep -v '^#' /workspace/.env | xargs)" >> ~/.bashrc

# Github & Huggingface tokens
git config --global credential.helper store
echo "https://${GITHUB_USER}:${GITHUB_TOKEN}@github.com" > ~/.git-credentials
echo "https://${HF_USER}:${HF_TOKEN}@huggingface.co" >> ~/.git-credentials

pip install --upgrade pip
pip install huggingface_hub
python -c "from huggingface_hub.hf_api import HfFolder; import os; HfFolder.save_token(os.environ['HF_TOKEN'])"

apt-get update
apt-get install git-lfs

# Unsloth
pip install --upgrade --force-reinstall --no-cache-dir torch==2.3.1 triton \
  --index-url https://download.pytorch.org/whl/cu121
pip install "unsloth[cu121-ampere-torch230] @ git+https://github.com/unslothai/unsloth.git"

cd /workspace
mkdir /workspace/code

# KVA
cd /workspace/code
git clone https://github.com/nielsrolf/kva.git
cd kva
pip install -e .

# KVA data
cd /workspace
git clone https://github.com/nielsrolf/imo-experiment-data.git
mv imo-experiment-data kva_store

# experisana
cd /workspace/code
git clone https://github.com/nielsrolf/asana-worker.git
cd asana-worker
pip install -e .
cp /workspace/code/influencing-model-organisms/unsloth/.env /workspace/code/asana-worker
apt-get update
apt-get install screen -y

# Main repo
cd /workspace/code
git clone https://github.com/julianstastny/influencing-model-organisms.git
cp /workspace/.env /workspace/code/influencing-model-organisms/unsloth/.env
pip install -r /workspace/code/influencing-model-organisms/unsloth/requirements.txt
mkdir /workspace/models

# Dataset
cd /workspace
git lfs install
git clone https://huggingface.co/datasets/longtermrisk/influencing-model-organisms

# Github
git config --global user.email "$GIT_EMAIL"
git config --global user.name "$GIT_NAME"

# Start kva-ui in a screen session
screen -dmS kva-ui-session bash -c "kva-ui; exec sh"

# Start experisana worker in a screen session
screen -dmS experisana-worker-session bash -c "experisana worker; exec sh"

# Optional: Wait a few seconds to ensure services have started
sleep 5

# Optional: List running screen sessions to verify
screen -list