Make `select_checkpoint_from_menu`'s behavior more generic over filenames

Small changes and fixes

Update train script after clean up

Reorganization and cleanups

Initial commit