refactor: deduplicate trainer and clean up load_data API

[marksverdhei]

Summary

Follow-up to #12 addressing review suggestions:

• Trainer deduplication: Extract _prepare_inputs(), _per_sample_kl(), _validate_batch(), and _zero_loss() from compute_loss and prediction_step, eliminating ~60 lines of duplicated logic
• load_data() API cleanup: Accept explicit dataset/dataset_split kwargs instead of requiring a DataConfig-like object, removing the type("_DC", ...) hack from cli.py

Test plan

• All existing tests pass (same results as main)
• ruff lint + format clean

🤖 Generated with Claude Code

[marksverdhei]

This is now conflicting with main after #16 (eval_dataset) and #20/#17 (pending perf PRs that also touch trainer.py) were merged/proposed. The refactoring approach is solid — deduplicating compute_loss/prediction_step into shared helpers and cleaning up the load_data API.

Needs a rebase onto current main to resolve conflicts, particularly in cli.py (eval_dataset routing changed in #16) and trainer.py.

[marksverdhei]

Closing — the trainer dedup is reasonable in isolation but (a) the project is pivoting away from PromptBakingTrainer toward SFT, so this polishes code that may be removed, and (b) the load_data API change introduces an awkward dual-mode signature (data_config OR explicit kwargs) with confusing merge logic (dataset_split = data_config.dataset_split if dataset == data_config.dataset else dataset_split) that's harder to reason about than the original. Will revisit cleanup once the pivot lands.