#!/bin/bash
# Check if "restart" argument is passed to force normal training
if [ "$1" = "restart" ]; then
    echo "Force restart: Running normal training ..."
    python -c "
import os
from toy_models.models.trainer import train_transformer_from_config
current_dir = os.getcwd()
train_transformer_from_config('config.toml', current_dir)
"
else
    # Check for checkpoints and run appropriate training
    python -c "
import os
from pathlib import Path
from toy_models.models.trainer import train_transformer_from_config, restart_from_checkpoint
current_dir = os.getcwd()
# Check if checkpoints directory exists and has .pt files
latest_checkpoint = Path('latest_checkpoint.pt')
if latest_checkpoint.exists():
    print(f'Found checkpoint: {latest_checkpoint}. Restarting from checkpoint...')
    restart_from_checkpoint(current_dir)
else:
    print('Starting training from beginning ...')
    train_transformer_from_config(current_dir)
"
fi