pytorch-lightning

1
import sys
2

3
from lightning_utilities.core.imports import RequirementCache
4

5
__all__ = []
6

7
if not RequirementCache("litdata"):
8
    raise ModuleNotFoundError("Please, run `pip install litdata`")  # E111
9

10
else:
11
    import litdata
12

13
    # Enable resolution at least for lower data namespace
14
    sys.modules["lightning.data"] = litdata
15

16
    from litdata.processing.functions import map, optimize, walk
17
    from litdata.streaming.combined import CombinedStreamingDataset
18
    from litdata.streaming.dataloader import StreamingDataLoader
19
    from litdata.streaming.dataset import StreamingDataset
20

21
    __all__ += [
22
        "LightningDataset",
23
        "StreamingDataset",
24
        "CombinedStreamingDataset",
25
        "StreamingDataLoader",
26
        "LightningIterableDataset",
27
        "map",
28
        "optimize",
29
        "walk",
30
    ]
31

32
# TODO: Move this to litdata
33
if RequirementCache("lightning_sdk"):
34
    from lightning_sdk import Machine
35

36
    __all__ += ["Machine"]
37