Clean up the pre-processing method in `dataset.py`

[vratins]

The current pre-processing method in src/dataset.py (_preprocess_one) could be potentially split into smaller methods for clarity (separating ligand, mate, protein, water handling)