Update dataset and problem code with updates to schemas

There were updates to support augmentation tasks, added support for LUPI datasets and problems, and some other smaller things. We should make sure we support/updated to all of those.