[pre-commit.ci] auto fixes from pre-commit.com hooks

for more information, see https://pre-commit.ci

Author: pre-commit-ci[bot]

litgpt/api.py

@@ -384,7 +384,9 @@ def distribute(
                 else:
                     kv_cache_size = fixed_kv_cache_size
                 model.set_kv_cache(
-                    batch_size=1, max_seq_length=kv_cache_size, device=fabric.device,
+                    batch_size=1,
+                    max_seq_length=kv_cache_size,
+                    device=fabric.device,
                 )
                 self.kv_cache_initialized = True
                 self.fixed_kv_cache_size = fixed_kv_cache_size
@@ -513,7 +515,9 @@ def generate(
             else:
                 device = self.preprocessor.device
             self.model.set_kv_cache(
-                batch_size=1, max_seq_length=max_returned_tokens, device=device,
+                batch_size=1,
+                max_seq_length=max_returned_tokens,
+                device=device,
             )
             self.kv_cache_initialized = True
 
@@ -522,7 +526,9 @@ def generate(
             tmp_device = self.model.mha.mask_cache.device
             self.model.clear_kv_cache()
             self.model.set_kv_cache(
-                batch_size=1, max_seq_length=max_returned_tokens, device=tmp_device,
+                batch_size=1,
+                max_seq_length=max_returned_tokens,
+                device=tmp_device,
             )
         else:
             for block in self.model.transformer.h:

litgpt/attention.py

@@ -136,9 +136,9 @@ def __call__(
         if use_mask:
             # Special case requires building a mask. `mask_cache` is only needed
             # then.
-            assert (
-                self.mask_cache is not None
-            ), "mask_cache must be given if sliding window attention is used, or if input_pos given and T > 1"
+            assert self.mask_cache is not None, (
+                "mask_cache must be given if sliding window attention is used, or if input_pos given and T > 1"
+            )
             if is_causal:
                 mask = self.mask_cache[:T, :T].view(1, 1, T, T)
                 is_causal = False
@@ -156,9 +156,7 @@ def __call__(
                 nh_k = self.config.n_query_groups
                 q_per_kv = nh_q // nh_k
                 if q_per_kv > 1:
-                    mask = mask.unsqueeze(2).expand(
-                        -1, -1, q_per_kv, -1, -1
-                    ).reshape(B, nh_q, T, -1)
+                    mask = mask.unsqueeze(2).expand(-1, -1, q_per_kv, -1, -1).reshape(B, nh_q, T, -1)
 
         # Efficient attention using Flash Attention CUDA kernels.
         # NOTE: efficient implementation is disabled if `mask` is not None or softcapping is enabled.
@@ -236,7 +234,8 @@ def scaled_dot_product_attention(
     # `torch.nn.functional.scaled_dot_product_attention`
     sqrt_scale = math.sqrt(scale)
     scores = _attention_compute_scores(
-        sqrt_scale * query, sqrt_scale * key,
+        sqrt_scale * query,
+        sqrt_scale * key,
     )
     scores = do_softcapping(scores, attention_logit_softcapping)
     if mask is None and is_causal:

litgpt/config.py

@@ -3,15 +3,14 @@
 from copy import deepcopy
 from dataclasses import dataclass, field
 from pathlib import Path
-from typing import Any, Callable, Literal, Optional, Type, Union, List
+from typing import Any, Callable, List, Literal, Optional, Type, Union
 
 import torch
 import yaml
 from typing_extensions import Self
 
 from litgpt.utils import find_multiple
 
-
 # See `Config.start_of_layer_hook`. A start of layer hook is called just before
 # a layer is computed. The call is `hook(x, block_idx, input_pos)`, where
 # `x` is the layer input, `block_idx` the number of the layer, and `input_pos`

litgpt/generate/base.py

@@ -171,7 +171,7 @@ def generate_fn(
 
     prompt_size = prompt.size(0)
     if prompt_size == 0:
-        raise ValueError(f"prompt must not be empty")
+        raise ValueError("prompt must not be empty")
     sample_kwargs = dict(
         temperature=temperature,
         top_k=top_k,

litgpt/model.py

@@ -5,9 +5,10 @@
 Based on the nanoGPT implementation: https://github.com/karpathy/nanoGPT and
 https://github.com/EleutherAI/gpt-neox/tree/main/megatron/model.
 """
+
+from dataclasses import replace
 from functools import partial
 from typing import Any, List, Optional, Tuple, Union
-from dataclasses import replace
 
 import torch
 import torch.nn as nn
@@ -189,7 +190,8 @@ def reset_parameters(self) -> None:
         self.mha.set_seq_length(self.max_seq_length, device=self.cos.device)
 
     def set_start_of_layer_hook(
-        self, hook: Optional[StartOfLayerHook],
+        self,
+        hook: Optional[StartOfLayerHook],
     ):
         """
         Sets a function `hook(x, block_idx, input_pos)`, which is called
@@ -452,7 +454,9 @@ def get_kv_cache_params(self) -> Optional[KVCacheParams]:
         batch_size = min(c.batch_size for c in caches)
         cache_length = min(c.cache_length for c in caches)
         params = replace(
-            params, batch_size=batch_size, cache_length=cache_length,
+            params,
+            batch_size=batch_size,
+            cache_length=cache_length,
         )
         return params
 

tests/generate/test_main.py

@@ -174,7 +174,8 @@ def test_main(fake_checkpoint_dir, monkeypatch, tensor_like):
     )
     assert (
         generate_mock.mock_calls
-        == [call(ANY, tensor_like, len_return_value, **sample_kwargs, eos_id=tokenizer_mock.return_value.eos_id)] * num_samples
+        == [call(ANY, tensor_like, len_return_value, **sample_kwargs, eos_id=tokenizer_mock.return_value.eos_id)]
+        * num_samples
     )
     expected_output = "foo bar baz\n" * num_samples
     # Allow for the config to be printed before the expected repeated strings.
@@ -209,9 +210,7 @@ def test_sample(temperature):
     )
     # Note: Both `sample` and `batched_sample` create only 1 sample, not 3.
     # It is like passing `logits[:, 1-:, :]`
-    token = batched_sample(
-        logits, kwargs=dict(temperature=temperature, top_p=0.8)
-    )
+    token = batched_sample(logits, kwargs=dict(temperature=temperature, top_p=0.8))
 
     assert token.shape == (2, 1)
     # sample is batch size 1 only for now - this should be [0, 1] once batched generation is supported

tests/test_batch.py

@@ -32,7 +32,9 @@ def create_llm(tmp_path, batch_size, max_seq_length, device) -> tuple[LLM, GPT]:
     )
     model: GPT = llm.model
     model.set_kv_cache(
-        batch_size=batch_size, max_seq_length=max_seq_length, device=device,
+        batch_size=batch_size,
+        max_seq_length=max_seq_length,
+        device=device,
     )
 
     return llm, model
@@ -89,7 +91,9 @@ def test_batched_equivalence(tmp_path):
     # Switch to batched generation
     model.clear_kv_cache()
     model.set_kv_cache(
-        batch_size=batch_size, max_seq_length=max_seq_length, device=device,
+        batch_size=batch_size,
+        max_seq_length=max_seq_length,
+        device=device,
     )
 
     toks_1: torch.Tensor = batched_next_token(

tests/test_chat.py

@@ -47,12 +47,8 @@ def test_generate(monkeypatch, generated, stop_tokens, expected):
     model.config.block_size = 100
     model.max_seq_length = 100
     # Mock methods called during generation
-    monkeypatch.setattr(
-        model, "kv_cache_max_prefill_length", lambda: 80
-    )
-    monkeypatch.setattr(
-        model, "kv_cache_max_tokens_forward", lambda: 20
-    )
+    monkeypatch.setattr(model, "kv_cache_max_prefill_length", lambda: 80)
+    monkeypatch.setattr(model, "kv_cache_max_tokens_forward", lambda: 20)
     it = iter(generated)
 
     def multinomial(*_, **__):

tests/test_model.py

@@ -4,7 +4,6 @@
 import random
 from copy import deepcopy
 from functools import partial
-from unittest import mock
 
 import pytest
 import torch