| # 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 |