#!/bin/bash

# Set environment variables
export CUDA_VISIBLE_DEVICES=0
export TRANSFORMERS_CACHE="/vol/bitbucket/jc4121/gptneo-2.7Bloratunning/cache"
export HF_DATASETS_CACHE="/vol/bitbucket/jc4121/gptneo-2.7Bloratunning/data"
export HF_HOME="/vol/bitbucket/jc4121/gptneo-2.7Bloratunning/hf_home"

# Create directories
mkdir -p $TRANSFORMERS_CACHE
mkdir -p $HF_DATASETS_CACHE
mkdir -p $HF_HOME
mkdir -p /vol/bitbucket/jc4121/gptneo-2.7Bloratunning/lib

# Add custom lib path to PYTHONPATH
export PYTHONPATH=/vol/bitbucket/jc4121/gptneo-2.7Bloratunning/lib:$PYTHONPATH

# Default artist
ARTIST="Taylor Swift"

# Check if an artist name was provided
if [ $# -ge 1 ]; then
    ARTIST="$1"
fi

# Run generation
echo "Generating lyrics for artist: $ARTIST"
python generate.py \
    --artist "$ARTIST" \
    --max_length 512 \
    --temperature 0.7 \
    --top_p 0.9 \
    --top_k 50 \
    --num_return_sequences 1