Global imports (Lightning-AI#1170)

litgpt/data/alpaca.py

@@ -7,9 +7,10 @@
 from typing import Optional, Union
 
 import torch
-from torch.utils.data import random_split, DataLoader
 from lightning_utilities.core.imports import RequirementCache
-from litgpt.data import SFTDataset, get_sft_collate_fn, DataModule
+from torch.utils.data import DataLoader, random_split
+
+from litgpt.data import DataModule, SFTDataset, get_sft_collate_fn
 from litgpt.prompts import PromptStyle
 from litgpt.tokenizer import Tokenizer
 

litgpt/data/alpaca_2k.py

@@ -3,8 +3,9 @@
 
 from dataclasses import dataclass, field
 from pathlib import Path
-from litgpt.data.alpaca import Alpaca
+
 from litgpt.data import SFTDataset
+from litgpt.data.alpaca import Alpaca
 
 
 @dataclass

litgpt/data/alpaca_gpt4.py

@@ -3,8 +3,8 @@
 
 from dataclasses import dataclass, field
 from pathlib import Path
-from litgpt.data.alpaca import Alpaca
 
+from litgpt.data.alpaca import Alpaca
 
 _URL = "https://raw.githubusercontent.com/Instruction-Tuning-with-GPT-4/GPT-4-LLM/main/data/alpaca_gpt4_data.json"
 

litgpt/data/base.py

@@ -1,13 +1,13 @@
 # Copyright Lightning AI. Licensed under the Apache License 2.0, see LICENSE file.
 from abc import abstractmethod
 from functools import partial
-from typing import List, Dict, Union, Optional, Callable, Any
+from typing import Any, Callable, Dict, List, Optional, Union
 
 import torch
+from lightning import LightningDataModule
 from torch import Tensor
 from torch.utils.data import Dataset
 
-from lightning import LightningDataModule
 from litgpt import Tokenizer
 from litgpt.prompts import PromptStyle
 

litgpt/data/deita.py

@@ -1,9 +1,8 @@
 # Copyright Lightning AI. Licensed under the Apache License 2.0, see LICENSE file.
 """Implementation derived from https://github.com/tloen/alpaca-lora"""
-from pathlib import Path
 from dataclasses import dataclass, field
-
-from typing import Optional, List, Union
+from pathlib import Path
+from typing import List, Optional, Union
 
 import torch
 from torch.utils.data import DataLoader

litgpt/data/dolly.py

@@ -9,7 +9,7 @@
 from torch.utils.data import random_split
 
 from litgpt import PromptStyle
-from litgpt.data import SFTDataset, Alpaca
+from litgpt.data import Alpaca, SFTDataset
 
 _URL: str = "https://huggingface.co/datasets/databricks/databricks-dolly-15k/resolve/main/databricks-dolly-15k.jsonl"
 

litgpt/data/flan.py

@@ -3,13 +3,13 @@
 import json
 from dataclasses import dataclass, field
 from pathlib import Path
-from typing import Optional, Dict, List, Set, Union
+from typing import Dict, List, Optional, Set, Union
 
 import torch
 from torch.utils.data import DataLoader
 
 from litgpt import PromptStyle
-from litgpt.data import SFTDataset, get_sft_collate_fn, DataModule
+from litgpt.data import DataModule, SFTDataset, get_sft_collate_fn
 from litgpt.data.alpaca import download_if_missing
 from litgpt.tokenizer import Tokenizer
 

litgpt/data/json_data.py

@@ -3,13 +3,13 @@
 import json
 from dataclasses import dataclass, field
 from pathlib import Path
-from typing import Optional, Union, Tuple, Any
+from typing import Any, Optional, Tuple, Union
 
 import torch
-from torch.utils.data import random_split, DataLoader
+from torch.utils.data import DataLoader, random_split
 
 from litgpt import PromptStyle
-from litgpt.data import SFTDataset, get_sft_collate_fn, DataModule
+from litgpt.data import DataModule, SFTDataset, get_sft_collate_fn
 from litgpt.tokenizer import Tokenizer
 
 

litgpt/data/lima.py

@@ -2,11 +2,10 @@
 """Implementation derived from https://github.com/tloen/alpaca-lora"""
 import os
 from dataclasses import dataclass, field
-
-from typing import Optional, List, Union
+from typing import List, Optional, Union
 
 import torch
-from torch.utils.data import random_split, DataLoader
+from torch.utils.data import DataLoader, random_split
 
 from litgpt import PromptStyle
 from litgpt.data import DataModule, SFTDataset, get_sft_collate_fn

litgpt/data/lit_data.py

@@ -2,7 +2,7 @@
 import os
 from dataclasses import dataclass, field
 from pathlib import Path
-from typing import Union, Optional, Tuple
+from typing import Optional, Tuple, Union
 
 from torch.utils.data import DataLoader
 

litgpt/data/longform.py

@@ -9,11 +9,10 @@
 from torch.utils.data import DataLoader
 
 from litgpt import PromptStyle
-from litgpt.data import SFTDataset, get_sft_collate_fn, DataModule
+from litgpt.data import DataModule, SFTDataset, get_sft_collate_fn
 from litgpt.data.alpaca import download_if_missing
 from litgpt.tokenizer import Tokenizer
 
-
 _URL = "https://raw.githubusercontent.com/akoksal/LongForm/main/dataset"
 
 

litgpt/data/openwebtext.py

@@ -3,7 +3,7 @@
 from dataclasses import dataclass, field
 from functools import partial
 from pathlib import Path
-from typing import Union, Optional
+from typing import Optional, Union
 
 from torch.utils.data import DataLoader
 

litgpt/data/prepare_slimpajama.py

@@ -6,8 +6,8 @@
 from pathlib import Path
 
 from litgpt import Tokenizer
-from litgpt.utils import CLI
 from litgpt.data.prepare_starcoder import DataChunkRecipe
+from litgpt.utils import CLI
 
 
 class SlimPajamaDataRecipe(DataChunkRecipe):

litgpt/data/tinyllama.py

@@ -1,7 +1,7 @@
 # Copyright Lightning AI. Licensed under the Apache License 2.0, see LICENSE file.
 from dataclasses import dataclass, field
 from pathlib import Path
-from typing import Union, Optional
+from typing import Optional, Union
 
 from torch.utils.data import DataLoader
 

litgpt/pretrain.py

@@ -30,9 +30,9 @@
     copy_config_files,
     num_parameters,
     parse_devices,
+    reset_parameters,
     save_config,
     save_hyperparameters,
-    reset_parameters,
 )
 
 
@@ -153,7 +153,7 @@ def main(
     t0 = time.perf_counter()
     with fabric.init_module(empty_init=True):
         model = GPT(config)
-    
+
     initialize_weights(fabric, model, n_layer=config.n_layer, n_embd=config.n_embd)
 
     if train.tie_embeddings:

tests/conftest.py

@@ -2,22 +2,13 @@
 
 import os
 import shutil
-import sys
 from pathlib import Path
 from typing import List
 
 import pytest
 import torch
 from lightning.fabric.utilities.testing import _runif_reasons
 
-wd = Path(__file__).parent.parent.absolute()
-
-
-@pytest.fixture(autouse=True)
-def add_wd_to_path():
-    # this adds support for running tests without the package installed
-    sys.path.append(str(wd))
-
 
 @pytest.fixture()
 def fake_checkpoint_dir(tmp_path):

tests/data/test_alpaca.py

@@ -1,10 +1,9 @@
 # Copyright Lightning AI. Licensed under the Apache License 2.0, see LICENSE file.
+from litgpt.data import Alpaca
+from litgpt.prompts import Alpaca as AlpacaPromptStyle
 
 
 def test_alpaca(mock_tokenizer, alpaca_path):
-    from litgpt.data import Alpaca
-    from litgpt.prompts import Alpaca as AlpacaPromptStyle
-
     alpaca = Alpaca(val_split_fraction=0.5, download_dir=alpaca_path.parent, file_name=alpaca_path.name, num_workers=0)
     assert isinstance(alpaca.prompt_style, AlpacaPromptStyle)
     alpaca.connect(mock_tokenizer, batch_size=2, max_seq_length=10)

tests/data/test_base.py

@@ -1,17 +1,17 @@
 # Copyright Lightning AI. Licensed under the Apache License 2.0, see LICENSE file.
-from unittest.mock import Mock
 
 import pytest
 import torch
 
+from litgpt.data import SFTDataset
+from litgpt.data import get_sft_collate_fn
+from litgpt.prompts import PromptStyle
+
 
 @pytest.mark.parametrize("mask_prompt", [True, False])
 @pytest.mark.parametrize("ignore_index", [-1, -100])
 @pytest.mark.parametrize("max_seq_length", [1000, 5])
 def test_sft_dataset(max_seq_length, ignore_index, mask_prompt, mock_tokenizer):
-    from litgpt.data import SFTDataset
-    from litgpt.prompts import PromptStyle
-
     class Style(PromptStyle):
         def apply(self, prompt, **kwargs):
             return f"In: {prompt} Out:"
@@ -42,8 +42,6 @@ def apply(self, prompt, **kwargs):
 @pytest.mark.parametrize("ignore_index", [-1, -100])
 @pytest.mark.parametrize("pad_id", [0, 100])
 def test_sft_collate_fn_padding(pad_id, ignore_index):
-    from litgpt.data import get_sft_collate_fn
-
     collate = get_sft_collate_fn(pad_id=pad_id, ignore_index=ignore_index)
     samples = [
         {"input_ids": torch.tensor([1, 2, 3]), "labels": torch.tensor([10, 20, 30])},
@@ -58,8 +56,6 @@ def test_sft_collate_fn_padding(pad_id, ignore_index):
 
 
 def test_sft_collate_fn_truncation():
-    from litgpt.data import get_sft_collate_fn
-
     collate = get_sft_collate_fn(max_seq_length=2)
     samples = [
         {"input_ids": torch.tensor([1, 2, 3]), "labels": torch.tensor([10, 20, 30])},

tests/data/test_deita.py

@@ -1,10 +1,12 @@
 # Copyright Lightning AI. Licensed under the Apache License 2.0, see LICENSE file.
 from unittest import mock
 
+from litgpt.data import Deita, SFTDataset
+from litgpt.data.deita import format_dataset
+from litgpt.prompts import Alpaca as AlpacaPromptStyle
 
-def test_format_dataset():
-    from litgpt.data.deita import format_dataset
 
+def test_format_dataset():
     data = [
         {
             "prompt": "prompt1",
@@ -43,9 +45,6 @@ def test_format_dataset():
 @mock.patch("litgpt.data.deita.format_dataset")
 @mock.patch("datasets.load_dataset")
 def test_deita(_, format_dataset_mock, mock_tokenizer, tmp_path):
-    from litgpt.data import Deita, SFTDataset
-    from litgpt.prompts import Alpaca as AlpacaPromptStyle
-
     format_dataset_mock.return_value = [
         {"instruction": "inst1", "output": "out1"},
         {"instruction": "inst2", "output": "out2"},

tests/data/test_dolly.py

@@ -1,10 +1,10 @@
 # Copyright Lightning AI. Licensed under the Apache License 2.0, see LICENSE file.
 
+from litgpt.data import Dolly
+from litgpt.prompts import Alpaca as AlpacaPromptStyle
 
-def test_dolly(mock_tokenizer, dolly_path):
-    from litgpt.data import Dolly
-    from litgpt.prompts import Alpaca as AlpacaPromptStyle
 
+def test_dolly(mock_tokenizer, dolly_path):
     alpaca = Dolly(val_split_fraction=0.5, download_dir=dolly_path.parent, file_name=dolly_path.name, num_workers=0)
     assert isinstance(alpaca.prompt_style, AlpacaPromptStyle)
     alpaca.connect(mock_tokenizer, batch_size=2, max_seq_length=10)

tests/data/test_json.py

@@ -1,13 +1,14 @@
 # Copyright Lightning AI. Licensed under the Apache License 2.0, see LICENSE file.
 import json
+
 import pytest
 
+from litgpt.data import JSON
+from litgpt.prompts import PromptStyle
+
 
 @pytest.mark.parametrize("as_jsonl", [False, True])
 def test_json(as_jsonl, tmp_path, mock_tokenizer):
-    from litgpt.data import JSON
-    from litgpt.prompts import PromptStyle
-
     class Style(PromptStyle):
         def apply(self, prompt, **kwargs):
             return f"X: {prompt} {kwargs['input']} Y:"
@@ -62,8 +63,6 @@ def apply(self, prompt, **kwargs):
 
 
 def test_json_input_validation(tmp_path):
-    from litgpt.data import JSON
-
     with pytest.raises(FileNotFoundError, match="The `json_path` must be a file or a directory"):
         JSON(tmp_path / "not exist")
 
@@ -85,8 +84,6 @@ def test_json_input_validation(tmp_path):
 
 @pytest.mark.parametrize("as_jsonl", [False, True])
 def test_json_with_splits(as_jsonl, tmp_path, mock_tokenizer):
-    from litgpt.data import JSON
-
     mock_train_data = [
         {"instruction": "Add", "input": "2+2", "output": "4"},
         {"instruction": "Subtract", "input": "5-3", "output": "2"},

tests/data/test_longform.py

@@ -1,10 +1,9 @@
 # Copyright Lightning AI. Licensed under the Apache License 2.0, see LICENSE file.
+from litgpt.data import LongForm
+from litgpt.prompts import Longform as LongFormPromptStyle
 
 
 def test_longform(mock_tokenizer, longform_path):
-    from litgpt.data import LongForm
-    from litgpt.prompts import Longform as LongFormPromptStyle
-
     alpaca = LongForm(download_dir=longform_path, num_workers=0)
     assert isinstance(alpaca.prompt_style, LongFormPromptStyle)
     alpaca.connect(mock_tokenizer, batch_size=2, max_seq_length=10)

tests/data/test_openwebtext.py

@@ -4,16 +4,16 @@
 from unittest.mock import ANY, call
 
 import pytest
+from litdata.streaming import StreamingDataLoader, StreamingDataset
 from torch.utils.data import DataLoader
 
+from litgpt.data import OpenWebText
+
 
 @pytest.mark.skipif(sys.platform == "win32", reason="Not in the mood to add Windows support right now.")
 @mock.patch("litdata.optimize")
 @mock.patch("datasets.load_dataset")
-def test_openwebtext(_, optimize_mock, tmp_path, monkeypatch, mock_tokenizer):
-    from litgpt.data import OpenWebText
-    from litdata.streaming import StreamingDataLoader, StreamingDataset
-
+def test_openwebtext(_, optimize_mock, tmp_path, mock_tokenizer):
     data = OpenWebText(data_path=(tmp_path / "openwebtext"))
     assert data.seq_length == 2048
     assert data.batch_size == 1

tests/data/test_tinyllama.py

@@ -1,13 +1,13 @@
 # Copyright Lightning AI. Licensed under the Apache License 2.0, see LICENSE file.
 
 import pytest
+from litdata.streaming import CombinedStreamingDataset, StreamingDataLoader, StreamingDataset
 from torch.utils.data import DataLoader
 
+from litgpt.data import TinyLlama
 
-def test_tinyllama(tmp_path, monkeypatch):
-    from litgpt.data import TinyLlama
-    from litdata.streaming import StreamingDataLoader, StreamingDataset, CombinedStreamingDataset
 
+def test_tinyllama(tmp_path):
     data = TinyLlama(data_path=(tmp_path / "data"))
     assert data.seq_length == 2048
     assert data.batch_size == 1