[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

@@ -378,7 +378,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
@@ -507,7 +509,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
 
@@ -516,7 +520,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

@@ -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.

litgpt/config.py

@@ -3,7 +3,7 @@
 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

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,6 +5,7 @@
 Based on the nanoGPT implementation: https://github.com/karpathy/nanoGPT and
 https://github.com/EleutherAI/gpt-neox/tree/main/megatron/model.
 """
+
 from functools import partial
 from typing import Any, List, Optional, Tuple, Union
 

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(*_, **__):