chore: lint

Author: winglian

src/axolotl/core/trainers/base.py

@@ -464,9 +464,9 @@ def get_train_dataloader(self) -> DataLoader:
                 "pin_memory": self.args.dataloader_pin_memory,
             }
             if self.args.dataloader_prefetch_factor:
-                dataloader_params["prefetch_factor"] = (
-                    self.args.dataloader_prefetch_factor
-                )
+                dataloader_params[
+                    "prefetch_factor"
+                ] = self.args.dataloader_prefetch_factor
 
             sampler = self._get_train_sampler()
             if isinstance(sampler, BatchSampler):
@@ -511,9 +511,9 @@ def get_eval_dataloader(self, eval_dataset: Optional[Dataset] = None) -> DataLoa
                 "pin_memory": self.args.dataloader_pin_memory,
             }
             if self.args.dataloader_prefetch_factor:
-                dataloader_params["prefetch_factor"] = (
-                    self.args.dataloader_prefetch_factor
-                )
+                dataloader_params[
+                    "prefetch_factor"
+                ] = self.args.dataloader_prefetch_factor
 
             if isinstance(eval_sampler, BatchSampler):
                 dataloader_params["batch_sampler"] = eval_sampler

src/axolotl/core/training_args.py

@@ -34,7 +34,9 @@ class AxolotlTrainingMixins:
     )
     sample_packing_sequentially: bool = field(
         default=False,
-        metadata={"help": "Use next-fit sample packing that preserves the order of samples coming from the sampler. Use in combination with curriculum_sampling for fully sequential packing."},
+        metadata={
+            "help": "Use next-fit sample packing that preserves the order of samples coming from the sampler. Use in combination with curriculum_sampling for fully sequential packing."
+        },
     )
     multipack_real_batches: bool = field(
         default=False,

src/axolotl/utils/samplers/multipack.py

@@ -124,7 +124,7 @@ def allocate_sequentially(lengths: np.ndarray, rank: int, c: int, n: int):
 
     # First, do sequential packing into bins
     all_bins = []
-    current_bin = [0 for i in range(0)] # numba hint
+    current_bin = [0 for i in range(0)]  # numba hint
     remaining_capacity = c
 
     for idx, size in enumerate(lengths):
@@ -190,8 +190,9 @@ def __init__(
         self.len_across_ranks = None
 
         if self.sequential and not isinstance(sampler, SequentialSampler):
-            LOG.warn("using sequential sample packing with non-sequential sampler, did you want to also enable curriculum_sampling?")
-
+            LOG.warn(
+                "using sequential sample packing with non-sequential sampler, did you want to also enable curriculum_sampling?"
+            )
 
     def set_epoch(self, epoch: int):
         self.epoch = epoch

tests/test_packed_batch_sampler.py

@@ -35,7 +35,9 @@ class TestBatchedSamplerPacking:
     )
     @pytest.mark.parametrize("max_seq_length", [4096, 512])
     @pytest.mark.parametrize("sequential", [True, False])
-    def test_packing(self, batch_size, num_workers, tokenizer, max_seq_length, sequential):
+    def test_packing(
+        self, batch_size, num_workers, tokenizer, max_seq_length, sequential
+    ):
         import axolotl.monkeypatch.data.batch_dataset_fetcher  # pylint: disable=unused-import  # noqa: F401
 
         dataset = load_dataset(
@@ -71,7 +73,7 @@ def test_packing(self, batch_size, num_workers, tokenizer, max_seq_length, seque
             batch_max_len=max_seq_length,
             group_size=100000,
             bin_size=200,
-            sequential=sequential
+            sequential=sequential,
         )
 
         loader = DataLoader(