Draft: Introduce automatic sanity check for training

Pull Request Title

Description

Wants to merge: vladb/test_train into main

Type of change

• Bug fix
• New feature
• Enhancement
• Documentation update
• Other (specify right below)

Merge request commits

• Introduce automatic sanity check for training

Related Issues

Screenshots or GIFs

Checklist

• I have tested the code with the changes manually.
• My code follows the project's style guidelines.
• I have documented my code for others to understand.
• I have updated documentation as needed (including README.md, code comments and doc strings).

Reviewer Guidelines

Additional Notes

@mentions

optimus/datasets/__init__.py

+import os
+
+from enum import Enum
+from functools import partial
+
 from .tinystories import TinyStoriesDataset
 from .wikitext103 import WikiText103Dataset
+from .line_ds import LineDataset
+
+
+class Datasets(Enum):
+
+    WIKITEXT103 = WikiText103Dataset
+    TINYSTORIES = TinyStoriesDataset
+
+    TEST_DS = partial(LineDataset, path=os.getenv("TEST_DATASET_PATH", default='tests/fixtures/sample.txt'))
+
+    def __call__(self, **kwargs):
+        return self.value(**kwargs)
\ No newline at end of file