Automated formatting (#1126)

Author: carmocca

README.md

@@ -18,7 +18,7 @@
 
  
 
-⚡ LitGPT is a hackable [implementation](litgpt/model.py) of state-of-the-art open-source large language models released under the **Apache 2.0 license**. 
+⚡ LitGPT is a hackable [implementation](litgpt/model.py) of state-of-the-art open-source large language models released under the **Apache 2.0 license**.
 
  
 ## LitGPT supports
@@ -141,7 +141,7 @@ For added convenience, you can also manually override config file setting via th
 
 
 ```bash
-litgpt finetune lora 
+litgpt finetune lora
   --config https://raw.githubusercontent.com/Lightning-AI/litgpt/main/config_hub/finetune/llama-2-7b/lora.yaml \
   --lora_r 4
 ```
@@ -150,7 +150,7 @@ You can browse the available configuration files [here](https://github.com/Light
 
  
 
-> [!TIP] 
+> [!TIP]
 > **Run large models on smaller consumer devices:**
 > We support 4-bit quantization (as in QLoRA), (bnb.nf4, bnb.nf4-dq, bnb.fp4, bnb.fp4-dq) and 8-bit quantization (bnb.int8) for inference by following [this guide](tutorials/quantize.md).
 
@@ -314,7 +314,7 @@ We welcome all individual contributors, regardless of their level of experience
 
  
 
-> [!TIP] 
+> [!TIP]
 > Unsure about contributing? Check out our [How to Contribute to LitGPT](https://lightning.ai/pages/community/tutorial/how-to-contribute-to-litgpt/) guide.
 
 If you have general questions about building with LitGPT, please [join our Discord](https://discord.gg/VptPCZkGNa).

litgpt/__main__.py

@@ -29,7 +29,7 @@
 
 
 def _new_parser(**kwargs: Any) -> "ArgumentParser":
-    from jsonargparse import ArgumentParser, ActionConfigFile
+    from jsonargparse import ActionConfigFile, ArgumentParser
 
     parser = ArgumentParser(**kwargs)
     parser.add_argument(
@@ -80,7 +80,7 @@ def main() -> None:
         "merge_lora": {"help": "Merges the LoRA weights with the base model.", "fn": merge_lora_fn},
     }
 
-    from jsonargparse import set_docstring_parse_options, set_config_read_mode
+    from jsonargparse import set_config_read_mode, set_docstring_parse_options
 
     set_docstring_parse_options(attribute_docstrings=True)
     set_config_read_mode(urls_enabled=True)

litgpt/chat/base.py

@@ -9,9 +9,9 @@
 import torch
 from lightning.fabric.plugins import BitsandbytesPrecision
 
-from litgpt.generate.base import next_token
 from litgpt import GPT, Config, PromptStyle, Tokenizer
-from litgpt.prompts import load_prompt_style, has_prompt_style
+from litgpt.generate.base import next_token
+from litgpt.prompts import has_prompt_style, load_prompt_style
 from litgpt.scripts.merge_lora import merge_lora
 from litgpt.utils import CLI, check_valid_checkpoint_dir, get_default_supported_precision, load_checkpoint
 
@@ -159,7 +159,9 @@ def main(
     model = fabric.setup_module(model)
 
     tokenizer = Tokenizer(checkpoint_dir)
-    prompt_style = load_prompt_style(checkpoint_dir) if has_prompt_style(checkpoint_dir) else PromptStyle.from_config(config)
+    prompt_style = (
+        load_prompt_style(checkpoint_dir) if has_prompt_style(checkpoint_dir) else PromptStyle.from_config(config)
+    )
     stop_tokens = prompt_style.stop_tokens(tokenizer)
 
     print(f"Now chatting with {config.name}.\nTo exit, press 'Enter' on an empty prompt.\n")

litgpt/config.py

@@ -1,12 +1,12 @@
 # Copyright Lightning AI. Licensed under the Apache License 2.0, see LICENSE file.
 
-import yaml
 from copy import deepcopy
 from dataclasses import dataclass, field
 from pathlib import Path
 from typing import Any, Literal, Optional, Type, Union
 
 import torch
+import yaml
 from typing_extensions import Self
 
 import litgpt.model

litgpt/data/alpaca.py

@@ -50,10 +50,7 @@ def __post_init__(self) -> None:
             self.prompt_style = PromptStyle.from_name(self.prompt_style)
 
     def connect(
-        self,
-        tokenizer: Optional[Tokenizer] = None,
-        batch_size: int = 1,
-        max_seq_length: Optional[int] = None
+        self, tokenizer: Optional[Tokenizer] = None, batch_size: int = 1, max_seq_length: Optional[int] = None
     ) -> None:
         self.tokenizer = tokenizer
         self.batch_size = batch_size
@@ -71,7 +68,7 @@ def setup(self, stage: str = "") -> None:
         train_data, test_data = random_split(
             data,
             [1.0 - self.val_split_fraction, self.val_split_fraction],
-            generator=torch.Generator().manual_seed(self.seed)
+            generator=torch.Generator().manual_seed(self.seed),
         )
         train_data, test_data = list(train_data), list(test_data)
 
@@ -99,7 +96,7 @@ def train_dataloader(self) -> DataLoader:
             shuffle=True,
             generator=torch.Generator().manual_seed(self.seed),
             num_workers=self.num_workers,
-            collate_fn=get_sft_collate_fn(max_seq_length=self.max_seq_length, ignore_index=self.ignore_index)
+            collate_fn=get_sft_collate_fn(max_seq_length=self.max_seq_length, ignore_index=self.ignore_index),
         )
 
     def val_dataloader(self) -> DataLoader:
@@ -108,7 +105,7 @@ def val_dataloader(self) -> DataLoader:
             batch_size=self.batch_size,
             shuffle=False,
             num_workers=self.num_workers,
-            collate_fn=get_sft_collate_fn(max_seq_length=self.max_seq_length, ignore_index=self.ignore_index)
+            collate_fn=get_sft_collate_fn(max_seq_length=self.max_seq_length, ignore_index=self.ignore_index),
         )
 
 

litgpt/data/alpaca_2k.py

@@ -50,4 +50,3 @@ def setup(self, stage: str = "") -> None:
             mask_prompt=self.mask_prompt,
             ignore_index=self.ignore_index,
         )
-

litgpt/data/base.py

@@ -17,10 +17,7 @@ class DataModule(LightningDataModule):
 
     @abstractmethod
     def connect(
-        self,
-        tokenizer: Optional[Tokenizer] = None,
-        batch_size: int = 1,
-        max_seq_length: Optional[int] = None
+        self, tokenizer: Optional[Tokenizer] = None, batch_size: int = 1, max_seq_length: Optional[int] = None
     ) -> None:
         """All settings that can't be determined at the time of instantiation need to be passed through here
         before any dataloaders can be accessed.
@@ -53,6 +50,7 @@ class SFTDataset(Dataset):
         labels: Same as input_ids, unless ``mask_prompt=True`` in which case the 'prompt' part is replaced with
             the ``ignore_index``.
     """
+
     def __init__(
         self,
         data: List[Dict[str, str]],
@@ -61,7 +59,7 @@ def __init__(
         max_seq_length: int = -1,
         mask_prompt: bool = True,
         ignore_index: int = -100,
-        transform: Optional[Callable[[Any], Any]] = None
+        transform: Optional[Callable[[Any], Any]] = None,
     ) -> None:
         self.data = data
         self.tokenizer = tokenizer
@@ -84,9 +82,7 @@ def __getitem__(self, idx: int) -> Dict[str, Tensor]:
         prompt_and_response = prompt + example["output"]
         encoded_prompt = self.tokenizer.encode(prompt, max_length=self.max_seq_length)
         encoded_prompt_and_response = self.tokenizer.encode(
-            prompt_and_response,
-            eos=True,
-            max_length=self.max_seq_length,
+            prompt_and_response, eos=True, max_length=self.max_seq_length
         )
 
         # The labels are the full prompt with response, but with the prompt masked out

litgpt/data/deita.py

@@ -45,10 +45,7 @@ def __post_init__(self) -> None:
             self.prompt_style = PromptStyle.from_name(self.prompt_style)
 
     def connect(
-        self,
-        tokenizer: Optional[Tokenizer] = None,
-        batch_size: int = 1,
-        max_seq_length: Optional[int] = None
+        self, tokenizer: Optional[Tokenizer] = None, batch_size: int = 1, max_seq_length: Optional[int] = None
     ) -> None:
         self.tokenizer = tokenizer
         self.batch_size = batch_size
@@ -99,7 +96,7 @@ def val_dataloader(self) -> DataLoader:
             batch_size=self.batch_size,
             shuffle=False,
             num_workers=self.num_workers,
-            collate_fn=get_sft_collate_fn(max_seq_length=self.max_seq_length, ignore_index=self.ignore_index)
+            collate_fn=get_sft_collate_fn(max_seq_length=self.max_seq_length, ignore_index=self.ignore_index),
         )
 
 

litgpt/data/dolly.py

@@ -50,7 +50,7 @@ def setup(self, stage: str = "") -> None:
         train_data, test_data = random_split(
             data,
             [1.0 - self.val_split_fraction, self.val_split_fraction],
-            generator=torch.Generator().manual_seed(self.seed)
+            generator=torch.Generator().manual_seed(self.seed),
         )
         train_data, test_data = list(train_data), list(test_data)
 

litgpt/data/flan.py

@@ -59,10 +59,7 @@ def __post_init__(self):
             self.subsets = list(supported_subsets)
 
     def connect(
-        self,
-        tokenizer: Optional[Tokenizer] = None,
-        batch_size: int = 1,
-        max_seq_length: Optional[int] = None
+        self, tokenizer: Optional[Tokenizer] = None, batch_size: int = 1, max_seq_length: Optional[int] = None
     ) -> None:
         self.tokenizer = tokenizer
         self.batch_size = batch_size
@@ -103,7 +100,7 @@ def _dataloader(self, split: str) -> DataLoader:
             shuffle=(split == "train"),
             generator=torch.Generator().manual_seed(self.seed),
             num_workers=self.num_workers,
-            collate_fn=get_sft_collate_fn(max_seq_length=self.max_seq_length, ignore_index=self.ignore_index)
+            collate_fn=get_sft_collate_fn(max_seq_length=self.max_seq_length, ignore_index=self.ignore_index),
         )