Implement dataloader with meaningful blocks separation

[dfsbora]

When splitting the data for training, we want to avoid splitting a single block (e.g. a proof environment in latex) in different parts if they could fit all together in a segment.

[dfsbora]

implemented, but testing and running some experiments before pushing