File size: 942 Bytes
fde051d | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 | #!/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 |