Black

Author: mseeger

keys_values/data/longbench_v2.py

@@ -415,7 +415,9 @@ def filter_and_transform(
     num_used = 0
     num_total = 0
     if max_seq_length is not None:
-        print(f"\nProcessing dataset, filtering out records with > {max_seq_length} tokens")
+        print(
+            f"\nProcessing dataset, filtering out records with > {max_seq_length} tokens"
+        )
     else:
         print(f"\nProcessing dataset")
     if seq_lengths is None:

keys_values/data/sequence_classification.py

@@ -110,7 +110,10 @@ def __getitem__(self, idx: int) -> Dict[str, Any]:
         prompt = self.prompt_style.apply(prompt=example["instruction"], **example)
         max_length = -1 if self.max_seq_length is None else self.max_seq_length
         encoded_prompt = self.tokenizer.encode(
-            prompt, bos=False, eos=True, max_length=max_length,
+            prompt,
+            bos=False,
+            eos=True,
+            max_length=max_length,
         )
         token_counts = {"raw_plus_prompt_template": len(encoded_prompt)}
         raw_count = example.get("num_tokens_instruction")

keys_values/kvcache/gradient/inference_replay.py

@@ -349,7 +349,13 @@ def __init__(
     ):
         extra_kwargs = dict()
         # If args are not in `base_kwargs`, take them from `replay_log`
-        for name in ("tokenizer", "end_initial_regex", "max_initial_fraction", "include_end_string", "pad_id"):
+        for name in (
+            "tokenizer",
+            "end_initial_regex",
+            "max_initial_fraction",
+            "include_end_string",
+            "pad_id",
+        ):
             if name not in base_kwargs:
                 extra_kwargs[name] = getattr(replay_log, name)
         SmartInitialLastRecentlyInsertedKVCache.__init__(