ruff

tests/lobster/data/test__calm_datamodule.py

@@ -1,11 +1,12 @@
 import unittest.mock
 
 import pytest
-from lobster.data import CalmLightningDataModule
-from lobster.datasets import CalmDataset
 from pandas import DataFrame
 from torch import Size
 
+from lobster.data import CalmLightningDataModule
+from lobster.datasets import CalmDataset
+
 
 @pytest.fixture
 def dm(tmp_path):

tests/lobster/data/test__ume_datamodule.py

@@ -1,11 +1,12 @@
 from pathlib import Path
 
 import pytest
-from lobster.data import UmeLightningDataModule
 from torch import Tensor
 from torch.utils.data import DataLoader
 from transformers import BatchEncoding
 
+from lobster.data import UmeLightningDataModule
+
 
 @pytest.fixture
 def dm(tmp_path):

tests/lobster/datasets/test__amplify_dataset.py

@@ -1,6 +1,7 @@
 import unittest.mock
 
 from datasets import Dataset
+
 from lobster.datasets import AMPLIFYIterableDataset
 
 

tests/lobster/datasets/test__calm_dataset.py

@@ -2,9 +2,10 @@
 from pathlib import Path
 
 from datasets import Dataset
-from lobster.datasets import CalmDataset, CalmIterableDataset
 from pandas import DataFrame
 
+from lobster.datasets import CalmDataset, CalmIterableDataset
+
 
 class TestCalmDataset:
     """Unit tests for CalmDataset."""

tests/lobster/datasets/test__concat_iterable_dataset.py

@@ -1,7 +1,8 @@
 import pytest
-from lobster.datasets import ConcatIterableDataset
 from torch.utils.data import IterableDataset
 
+from lobster.datasets import ConcatIterableDataset
+
 
 class IterableStringDataset(IterableDataset):
     def __init__(self, data):

tests/lobster/datasets/test__latent_generator_3d_coord_dataset.py

@@ -1,18 +1,20 @@
 import unittest.mock
 
 from datasets import Dataset
+
 from lobster.datasets import LatentGeneratorPinderIterableDataset
 
+
 class TestLatentGeneratorPinderIterableDataset:
     @unittest.mock.patch("lobster.datasets._huggingface_iterable_dataset.load_dataset")
     def test__iter__(self, mock_load_dataset):
         mock_load_dataset.return_value = Dataset.from_list(
             [
                 {
                     "__index_level_0__": "7u9z__G1_Q92736--7u9z__D1_P68106",
-                    "lg_token_string": [ "ft ec ec hp ek bt bt ek . da da ek da da ec da hx ec", 
-                                        "ec ec ec da da da hp hp hp . bt da ig ig ig ig da gv hp", 
-                                        "gj dp gj c gj ec hx cd cz cg . ec ec gj if da ft fe ft", 
+                    "lg_token_string": [ "ft ec ec hp ek bt bt ek . da da ek da da ec da hx ec",
+                                        "ec ec ec da da da hp hp hp . bt da ig ig ig ig da gv hp",
+                                        "gj dp gj c gj ec hx cd cz cg . ec ec gj if da ft fe ft",
                                         "fe l hx hx da da hx hx hx hx hx . hx hx da ec ec da ft"],
                 }
             ]
@@ -27,4 +29,4 @@ def test__iter__(self, mock_load_dataset):
 
 
 if __name__ == "__main__":
-    TestLatentGeneratorPinderIterableDataset().test__iter__()
+    TestLatentGeneratorPinderIterableDataset().test__iter__()

tests/lobster/datasets/test__m3_20m_dataset.py

@@ -3,9 +3,10 @@
 from tempfile import NamedTemporaryFile
 
 from datasets import Dataset
-from lobster.datasets import M320MDataset, M320MIterableDataset
 from pandas import DataFrame
 
+from lobster.datasets import M320MDataset, M320MIterableDataset
+
 
 class TestM320MDataset:
     """Unit tests for M320MDataset."""

tests/lobster/tokenization/test__latent_generator_tokenizer.py

@@ -1,8 +1,9 @@
+from transformers import PreTrainedTokenizerFast
+
 from lobster.tokenization._latent_generator_3d_coord_tokenizer import (
     LatentGenerator3DCoordTokenizerFast,
     _make_latent_generator_3d_coord_tokenizer,
 )
-from transformers import PreTrainedTokenizerFast
 
 
 def test__make_latent_generator_3d_coord_tokenizer():