Clean up hardcoded dataset specific check from datamodule

[ikanher]

In datamodules.py in the LLM path, we have a hardcoded check for "wikitext" dataset.

This should be removed and the functionality should be provided with e.g. --dataset-subset XXX CLI flag.