#!/bin/bash

## Setting up the virtual environment
VENV_NAME=venv-lola-inference

# Create python virtual environment
python -m venv ./$VENV_NAME
# Activate the virtual environment
source ./$VENV_NAME/bin/activate

# Note: Change the below installation as per your hardware setup
pip install torch==2.5.0 torchvision==0.20.0 torchaudio==2.5.0 --index-url https://download.pytorch.org/whl/cu121
pip install transformers[torch]==4.47.0

## Downloading the latest LOLA classes
LOLA_CLASS_DIR="lola_hf_model"
if [ -d "$LOLA_CLASS_DIR" ]; then
    echo "Directory $LOLA_CLASS_DIR already exists."
else
    echo "Creating directory $DIR_NAME."
    mkdir -p "$LOLA_CLASS_DIR"
fi

cd $LOLA_CLASS_DIR

[ -f "modeling_lola_gpt2.py" ] || wget https://huggingface.co/dice-research/lola_v1/resolve/main/modeling_lola_gpt2.py
[ -f "configuration_lola_gpt2.py" ] || wget https://huggingface.co/dice-research/lola_v1/resolve/main/configuration_lola_gpt2.py

echo "Setup complete."