Remove/refactor chain filtering logic in dataset.py

[vratins]

As of now we don't filter by chain when training the model, hence that part of the code is vestigial, can be removed for now -- we can add it back if any point in the future we decide to select by chain.

[vratins]

Closed with #43