Add support to train the model on other text classification datasets, which will make it suitable for more use cases.