Clean up training pipeline and add export benchmarks

Initial commit