Fix bug concerning smart-lastrec (#102)

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

@@ -347,16 +347,23 @@ def __init__(
         replay_log: SmartInitialLastRecentlyInsertedKVCacheReplayLog,
         **base_kwargs,
     ):
+        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",
+        ):
+            if name not in base_kwargs:
+                extra_kwargs[name] = getattr(replay_log, name)
         SmartInitialLastRecentlyInsertedKVCache.__init__(
             self,
             config=config,
             buffers=buffers,
             block_idx=block_idx,
-            tokenizer=replay_log.tokenizer,
-            end_initial_regex=replay_log.end_initial_regex,
-            max_initial_fraction=replay_log.max_initial_fraction,
-            include_end_string=replay_log.include_end_string,
-            pad_id=replay_log.pad_id,
+            **extra_kwargs,
             **base_kwargs,
         )
         InferenceReplayCacheMixin.__init__(self)