diff --git a/hyperlora/instruction_finetune.py b/hyperlora/instruction_finetune.py
index 6cd072a..c638a88 100644
--- a/hyperlora/instruction_finetune.py
+++ b/hyperlora/instruction_finetune.py
@@ -1,10 +1,18 @@
 import transformers
-from peft import (
-    LoraConfig,
-)
 from datasets import load_dataset
-from modeling_icae_multi_span import ICAE, ModelArguments, DataArguments, TrainingArguments
-from training_utils import instruct_ft_tokenize_function, DataCollatorForDynamicPadding, train_model
+from modeling_icae_multi_span import (
+    ICAE,
+    DataArguments,
+    ModelArguments,
+    TrainingArguments,
+)
+from peft import LoraConfig
+from training_utils import (
+    DataCollatorForDynamicPadding,
+    instruct_ft_tokenize_function,
+    train_model,
+)
+
 
 def main():
     parser = transformers.HfArgumentParser((ModelArguments, DataArguments, TrainingArguments))
@@ -12,22 +20,22 @@ def main():
 
     print(model_args)
     print(data_args)
-    
-    training_args.gradient_checkpointing_kwargs = {"use_reentrant": False}  # manually add this argument in the code
+
+    training_args.gradient_checkpointing_kwargs = {
+        "use_reentrant": False
+    }  # manually add this argument in the code
 
     lora_config = LoraConfig(
-        r=model_args.lora_r,
-        lora_alpha=32,
-        lora_dropout=0.05,
-        bias="none",
-        task_type="CAUSAL_LM"
+        r=model_args.lora_r, lora_alpha=32, lora_dropout=0.05, bias="none", task_type="CAUSAL_LM"
     )
-    
+
     # check model_args.mem_size and min_tokens_for_lm
-    assert (training_args.fixed_mem_size & (training_args.fixed_mem_size - 1)) == 0, "training_args.fixed_mem_size must be a power of 2"    
-    
+    assert (
+        training_args.fixed_mem_size & (training_args.fixed_mem_size - 1)
+    ) == 0, "training_args.fixed_mem_size must be a power of 2"
+
     memory_size = training_args.fixed_mem_size
-    
+
     train_file = "/path/to/train/file"
     eval_file = "/path/to/dev/file"
 
@@ -41,9 +49,14 @@ def main():
     model = ICAE(model_args, training_args, lora_config)
     MEM_TOKENS = list(range(model.vocab_size, model.vocab_size + memory_size))
 
-    train_dataset = train_dataset.map(instruct_ft_tokenize_function, batched=True, fn_kwargs={"model": model, "mem": MEM_TOKENS})
-    eval_dataset = eval_dataset.map(instruct_ft_tokenize_function, batched=True, fn_kwargs={"model": model, "mem": MEM_TOKENS})
+    train_dataset = train_dataset.map(
+        instruct_ft_tokenize_function, batched=True, fn_kwargs={"model": model, "mem": MEM_TOKENS}
+    )
+    eval_dataset = eval_dataset.map(
+        instruct_ft_tokenize_function, batched=True, fn_kwargs={"model": model, "mem": MEM_TOKENS}
+    )
 
     train_model(model, train_dataset, eval_dataset, training_args)
 
-main()
\ No newline at end of file
+
+main()
diff --git a/icae_v2/fine_tuned_inference.py b/icae_v2/fine_tuned_inference.py
index c675475..4ec8f47 100644
--- a/icae_v2/fine_tuned_inference.py
+++ b/icae_v2/fine_tuned_inference.py
@@ -1,11 +1,17 @@
 import json
-import torch
-from tqdm import tqdm
-from transformers import HfArgumentParser, AutoModelForCausalLM
-from peft import LoraConfig
-from modeling_icae_multi_span import ICAE, ModelArguments, DataArguments, TrainingArguments
 import sys
+
+import torch
+from modeling_icae_multi_span import (
+    ICAE,
+    DataArguments,
+    ModelArguments,
+    TrainingArguments,
+)
+from peft import LoraConfig
 from safetensors.torch import load_file
+from tqdm import tqdm
+from transformers import AutoModelForCausalLM, HfArgumentParser
 
 # Set the computation device
 device = "cuda"
@@ -16,11 +22,7 @@ model_args, data_args, training_args = parser.parse_args_into_dataclasses()
 
 # Define Lora configuration
 lora_config = LoraConfig(
-    r=512,
-    lora_alpha=32,
-    lora_dropout=model_args.lora_dropout,
-    bias="none",
-    task_type="CAUSAL_LM"
+    r=512, lora_alpha=32, lora_dropout=model_args.lora_dropout, bias="none", task_type="CAUSAL_LM"
 )
 
 # Initialize model and send it to CUDA device
@@ -29,14 +31,14 @@ model = ICAE(model_args, training_args, lora_config)
 # Load the fine-tuned checkpoint
 print(f"Loading trained checkpoint from {training_args.output_dir}")
 state_dict = load_file(training_args.output_dir)
-model.load_state_dict(state_dict, strict=False) # only load lora and memory token embeddings
+model.load_state_dict(state_dict, strict=False)  # only load lora and memory token embeddings
 
 model = model.to(device)
 
 # Read the data file
 file_path = "./dev_v2.jsonl"
 lines = None
-with open(file_path, "r") as f:
+with open(file_path) as f:
     lines = f.readlines()
 
 # Prepare the model for evaluation
@@ -49,24 +51,40 @@ with torch.no_grad():
         for line in tqdm(lines):
             # Tokenize input text
             data = json.loads(line)
-            tokenized_input = model.tokenizer(data['input'], truncation=True, max_length=5120, padding=False, return_attention_mask=False)
-            tokenized_prompt = model.tokenizer(data['prompt'], truncation=False, padding=False, return_attention_mask=False, add_special_tokens=False)
+            tokenized_input = model.tokenizer(
+                data["input"],
+                truncation=True,
+                max_length=5120,
+                padding=False,
+                return_attention_mask=False,
+            )
+            tokenized_prompt = model.tokenizer(
+                data["prompt"],
+                truncation=False,
+                padding=False,
+                return_attention_mask=False,
+                add_special_tokens=False,
+            )
             # Generate compressed outputs
-            input_ids = torch.LongTensor([tokenized_input['input_ids']]).to(device)
+            input_ids = torch.LongTensor([tokenized_input["input_ids"]]).to(device)
             memory_slots = model._compress(input_ids)
-            
+
             # decoder input has 3 parts: prefix, memory slots and suffix
             # the following code is for Mistral tokenizer for example: 733, 16289, 28793 are for the Mistral instruction tempmlate
-            prompt_left_ids =  torch.LongTensor([[1, 733, 16289, 28793]]).to(device)
-            prompt_right_ids = [model.ft_token_id] + tokenized_prompt['input_ids'] + [733, 28748, 16289, 28793]
+            prompt_left_ids = torch.LongTensor([[1, 733, 16289, 28793]]).to(device)
+            prompt_right_ids = (
+                [model.ft_token_id] + tokenized_prompt["input_ids"] + [733, 28748, 16289, 28793]
+            )
             prompt_right_ids = torch.LongTensor([prompt_right_ids]).to(device)
 
             prompt_left_embs = model.tokens_to_embeddings(prompt_left_ids)
             prompt_right_embs = model.tokens_to_embeddings(prompt_right_ids)
             memory_slots = memory_slots.to(prompt_right_embs)
-                        
+
             # Concatenate and clone input embeddings
-            decoder_input_embeddings = torch.cat((prompt_left_embs, memory_slots.unsqueeze(0), prompt_right_embs), dim=1)
+            decoder_input_embeddings = torch.cat(
+                (prompt_left_embs, memory_slots.unsqueeze(0), prompt_right_embs), dim=1
+            )
             output = decoder_input_embeddings.clone()
 
             generate_text = []
@@ -74,29 +92,36 @@ with torch.no_grad():
 
             # Generate text output
             for i in range(max_out_length):
-                with model.icae.disable_adapter():   # no independent decoder; use self.icae
-                    out = model.icae(inputs_embeds=output, past_key_values=past_key_values, use_cache=True)
+                with model.icae.disable_adapter():  # no independent decoder; use self.icae
+                    out = model.icae(
+                        inputs_embeds=output, past_key_values=past_key_values, use_cache=True
+                    )
                 # out = decoder(inputs_embeds=output, past_key_values=past_key_values, use_cache=True)
-                logit = out.logits[:, -1, :model.vocab_size-1]
+                logit = out.logits[:, -1, : model.vocab_size - 1]
                 past_key_values = out.past_key_values
 
                 next_token_id = torch.argmax(logit, dim=-1)
                 # print(next_token_id)
-                
-                if next_token_id.item() == 2:   # eos
+
+                if next_token_id.item() == 2:  # eos
                     break
 
-                output = model.icae.get_base_model().model.embed_tokens(next_token_id).unsqueeze(1).to(device)
+                output = (
+                    model.icae.get_base_model()
+                    .model.embed_tokens(next_token_id)
+                    .unsqueeze(1)
+                    .to(device)
+                )
                 generate_text.append(next_token_id.item())
 
             generated_text = model.tokenizer.decode(generate_text)
 
             # Structure output data
             output_ = {
-                "input": data['input'],
+                "input": data["input"],
                 "prompt": data["prompt"],
                 "output": generated_text,
-                "answer": data["answer"]
+                "answer": data["answer"],
             }
 
             f.write(json.dumps(output_) + "\n")
diff --git a/icae_v2/instruction_finetune.py b/icae_v2/instruction_finetune.py
index 6cd072a..c638a88 100644
--- a/icae_v2/instruction_finetune.py
+++ b/icae_v2/instruction_finetune.py
@@ -1,10 +1,18 @@
 import transformers
-from peft import (
-    LoraConfig,
-)
 from datasets import load_dataset
-from modeling_icae_multi_span import ICAE, ModelArguments, DataArguments, TrainingArguments
-from training_utils import instruct_ft_tokenize_function, DataCollatorForDynamicPadding, train_model
+from modeling_icae_multi_span import (
+    ICAE,
+    DataArguments,
+    ModelArguments,
+    TrainingArguments,
+)
+from peft import LoraConfig
+from training_utils import (
+    DataCollatorForDynamicPadding,
+    instruct_ft_tokenize_function,
+    train_model,
+)
+
 
 def main():
     parser = transformers.HfArgumentParser((ModelArguments, DataArguments, TrainingArguments))
@@ -12,22 +20,22 @@ def main():
 
     print(model_args)
     print(data_args)
-    
-    training_args.gradient_checkpointing_kwargs = {"use_reentrant": False}  # manually add this argument in the code
+
+    training_args.gradient_checkpointing_kwargs = {
+        "use_reentrant": False
+    }  # manually add this argument in the code
 
     lora_config = LoraConfig(
-        r=model_args.lora_r,
-        lora_alpha=32,
-        lora_dropout=0.05,
-        bias="none",
-        task_type="CAUSAL_LM"
+        r=model_args.lora_r, lora_alpha=32, lora_dropout=0.05, bias="none", task_type="CAUSAL_LM"
     )
-    
+
     # check model_args.mem_size and min_tokens_for_lm
-    assert (training_args.fixed_mem_size & (training_args.fixed_mem_size - 1)) == 0, "training_args.fixed_mem_size must be a power of 2"    
-    
+    assert (
+        training_args.fixed_mem_size & (training_args.fixed_mem_size - 1)
+    ) == 0, "training_args.fixed_mem_size must be a power of 2"
+
     memory_size = training_args.fixed_mem_size
-    
+
     train_file = "/path/to/train/file"
     eval_file = "/path/to/dev/file"
 
@@ -41,9 +49,14 @@ def main():
     model = ICAE(model_args, training_args, lora_config)
     MEM_TOKENS = list(range(model.vocab_size, model.vocab_size + memory_size))
 
-    train_dataset = train_dataset.map(instruct_ft_tokenize_function, batched=True, fn_kwargs={"model": model, "mem": MEM_TOKENS})
-    eval_dataset = eval_dataset.map(instruct_ft_tokenize_function, batched=True, fn_kwargs={"model": model, "mem": MEM_TOKENS})
+    train_dataset = train_dataset.map(
+        instruct_ft_tokenize_function, batched=True, fn_kwargs={"model": model, "mem": MEM_TOKENS}
+    )
+    eval_dataset = eval_dataset.map(
+        instruct_ft_tokenize_function, batched=True, fn_kwargs={"model": model, "mem": MEM_TOKENS}
+    )
 
     train_model(model, train_dataset, eval_dataset, training_args)
 
-main()
\ No newline at end of file
+
+main()
diff --git a/icae_v2/modeling_icae_multi_span.py b/icae_v2/modeling_icae_multi_span.py
index 3da71d8..5fa01a9 100644
--- a/icae_v2/modeling_icae_multi_span.py
+++ b/icae_v2/modeling_icae_multi_span.py
@@ -1,42 +1,43 @@
 # ICAE that supports multi span concat
 
-import transformers
-from transformers import AutoModelForCausalLM, AutoTokenizer, AutoConfig
+import math
 import os
-import torch
-import torch.nn as nn
 import random
 from dataclasses import dataclass, field
 from typing import Optional
-from peft import (
-    get_peft_model,
-)
-from torch.nn.functional import gelu
-import math
+
+import torch
+import torch.nn as nn
+import transformers
+from peft import get_peft_model
 from safetensors.torch import load_file
+from torch.nn.functional import gelu
+from transformers import AutoConfig, AutoModelForCausalLM, AutoTokenizer
 
 device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
 
+
 @dataclass
 class ModelArguments:
     model_name_or_path: str = field(default="mistralai/Mistral-7B-v0.1")
-    lora_r: int = field(
-        default=128,
-        metadata={"help": "lora rank"}
-    )
-    lora_dropout: float = field(
-        default=0.05,
-        metadata={"help": "lora dropout"}
-    )
+    lora_r: int = field(default=128, metadata={"help": "lora rank"})
+    lora_dropout: float = field(default=0.05, metadata={"help": "lora dropout"})
     train: bool = field(
         default=True,
-        metadata={"help": "if true, the model ckpt will be initialized for training; else, it's for inference"}
+        metadata={
+            "help": "if true, the model ckpt will be initialized for training; else, it's for inference"
+        },
     )
 
+
 @dataclass
 class DataArguments:
     data_path: str = field(default=None, metadata={"help": "Path to the training data."})
-    debug_data: bool = field(default=False, metadata={"help": "Enable debug dataset to quickly verify the training process"})
+    debug_data: bool = field(
+        default=False,
+        metadata={"help": "Enable debug dataset to quickly verify the training process"},
+    )
+
 
 @dataclass
 class TrainingArguments(transformers.TrainingArguments):
@@ -44,7 +45,9 @@ class TrainingArguments(transformers.TrainingArguments):
     optim: str = field(default="adamw_torch")
     model_max_length: int = field(
         default=28000,
-        metadata={"help": "Maximum sequence length. Sequences will be right padded (and possibly truncated)."},
+        metadata={
+            "help": "Maximum sequence length. Sequences will be right padded (and possibly truncated)."
+        },
     )
     fixed_mem_size: int = field(
         default=128,
@@ -68,13 +71,16 @@ class TrainingArguments(transformers.TrainingArguments):
     )
     add_special_token_for_lm: bool = field(
         default=False,
-        metadata={"help": "Add a special token for the prompt of language modeling; default: False"},
+        metadata={
+            "help": "Add a special token for the prompt of language modeling; default: False"
+        },
     )
     restore_from: str = field(
         default="",
-        metadata={"help": "The checkpoint that should be restored from for fine-tuning"}
+        metadata={"help": "The checkpoint that should be restored from for fine-tuning"},
     )
 
+
 def print_trainable_parameters(model):
     trainable_parameters = 0
     all_param = 0
@@ -82,7 +88,9 @@ def print_trainable_parameters(model):
         all_param += param.numel()
         if param.requires_grad:
             trainable_parameters += param.numel()
-    print(f"trainable params: {trainable_parameters} || all params: {all_param} || trainable%: {100 * trainable_parameters / all_param}")
+    print(
+        f"trainable params: {trainable_parameters} || all params: {all_param} || trainable%: {100 * trainable_parameters / all_param}"
+    )
     # for name, param in model.named_parameters():
     #     if param.requires_grad:
     #         print(name, param.shape)
@@ -99,46 +107,61 @@ class ICAE(torch.nn.Module):
         self.model_args = model_args
         self.training_args = training_args
         self.model_name = model_args.model_name_or_path
-        self.icae = AutoModelForCausalLM.from_pretrained(self.model_name, torch_dtype=torch.float16 if training_args.bf16 is False else torch.bfloat16, use_flash_attention_2=True, resume_download=True)
-        
-        self.training = self.model_args.train    
-        
-        if self.training:    # indepedent model for gradient checkpointing
-            self.decoder = AutoModelForCausalLM.from_pretrained(self.model_name, torch_dtype=torch.float16 if training_args.bf16 is False else torch.bfloat16, use_flash_attention_2=True, resume_download=True)
+        self.icae = AutoModelForCausalLM.from_pretrained(
+            self.model_name,
+            torch_dtype=torch.float16 if training_args.bf16 is False else torch.bfloat16,
+            use_flash_attention_2=True,
+            resume_download=True,
+        )
 
-        self.vocab_size = self.icae.config.vocab_size + 1    # [PAD] token
+        self.training = self.model_args.train
+
+        if self.training:  # indepedent model for gradient checkpointing
+            self.decoder = AutoModelForCausalLM.from_pretrained(
+                self.model_name,
+                torch_dtype=torch.float16 if training_args.bf16 is False else torch.bfloat16,
+                use_flash_attention_2=True,
+                resume_download=True,
+            )
+
+        self.vocab_size = self.icae.config.vocab_size + 1  # [PAD] token
         self.pad_token_id = self.vocab_size - 1
         self.mean_compression_rate = training_args.mean_compression_rate
 
         # tunable
         self.mem_size = self.training_args.fixed_mem_size
-        self.vocab_size_with_mem = self.vocab_size + self.mem_size # so, the mem tokens are in the range [self.vocab_size, self.vocab_size + self.mem_size)
+        self.vocab_size_with_mem = (
+            self.vocab_size + self.mem_size
+        )  # so, the mem tokens are in the range [self.vocab_size, self.vocab_size + self.mem_size)
 
         # special tokens in addition to mem and length tokens
         self.ae_token_id = self.vocab_size_with_mem + 0
         self.lm_token_id = self.vocab_size_with_mem + 1
-        self.ft_token_id = self.vocab_size_with_mem + 2        
+        self.ft_token_id = self.vocab_size_with_mem + 2
+
+        self.icae.resize_token_embeddings(self.vocab_size_with_mem + 3)
 
-        self.icae.resize_token_embeddings(self.vocab_size_with_mem + 3) 
-        
         # special tokens for Llama-2/Mistral tokenizer
         self.bos_id = 1
         self.eos_id = 2
-        
+
         self.dim = self.icae.config.hidden_size
         self.icae = get_peft_model(self.icae, lora_config)
-        
+
         device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
 
         self.memory_token_embed = nn.Embedding(self.mem_size + 3, self.dim, padding_idx=None)
         self.loss_fct = nn.CrossEntropyLoss(ignore_index=-100)
         self.tokenizer = AutoTokenizer.from_pretrained(self.model_name, use_fast=False)
-        self.append_sequence = torch.arange(self.vocab_size, self.vocab_size + self.mem_size, dtype=torch.long, device=device).unsqueeze(0)    # mem tokens
-        
+        self.append_sequence = torch.arange(
+            self.vocab_size, self.vocab_size + self.mem_size, dtype=torch.long, device=device
+        ).unsqueeze(
+            0
+        )  # mem tokens
+
         if self.training:
             self.init()
 
-
     def init(self):
         print("Freezing the decoder...")
         freeze_model(self.decoder)
@@ -151,15 +174,15 @@ class ICAE(torch.nn.Module):
             print(f"Finished loading from {self.training_args.restore_from}")
         print("Enabling gradient checkpointing...")
         # self.icae.gradient_checkpointing_enable(gradient_checkpointing_kwargs={"use_reentrant": False})
-        self.decoder.gradient_checkpointing_enable(gradient_checkpointing_kwargs={"use_reentrant": False})
-                
-        
+        self.decoder.gradient_checkpointing_enable(
+            gradient_checkpointing_kwargs={"use_reentrant": False}
+        )
+
     def compute_num_segments(self, total_length):
         assert total_length > 0
         num_segments = math.ceil(total_length / (self.mem_size * self.mean_compression_rate))
         return num_segments
 
-
     def forward(
         self,
         input_ids: torch.LongTensor = None,
@@ -171,73 +194,93 @@ class ICAE(torch.nn.Module):
         total_length = input_ids.size(1)
         num_segments = self.compute_num_segments(total_length)
         segment_length = math.ceil(total_length / num_segments)
-        
+
         prompt_answer_embs = self.icae.get_base_model().model.embed_tokens(prompt_answer_ids)
         max_compressed_length = num_segments * self.mem_size
         compress_outputs = torch.zeros((max_compressed_length, self.dim)).to(prompt_answer_embs)
-        
+
         for segment_idx in range(num_segments):
-            
+
             start_idx = segment_idx * segment_length
             end_idx = min((segment_idx + 1) * segment_length, total_length)
             segment_input_ids = input_ids[:, start_idx:end_idx]
             segment_input_ids = torch.cat([segment_input_ids, self.append_sequence], dim=1)
             mem_flag = segment_input_ids >= self.vocab_size
 
-            segment_input_embedding = self.icae.get_base_model().model.embed_tokens(segment_input_ids)
-            segment_input_embedding[mem_flag] = self.memory_token_embed(segment_input_ids[mem_flag] - self.vocab_size).to(segment_input_embedding)
+            segment_input_embedding = self.icae.get_base_model().model.embed_tokens(
+                segment_input_ids
+            )
+            segment_input_embedding[mem_flag] = self.memory_token_embed(
+                segment_input_ids[mem_flag] - self.vocab_size
+            ).to(segment_input_embedding)
 
             # compress the current segment
-            segment_compress_outputs = self.icae(inputs_embeds=segment_input_embedding, output_hidden_states=True)
+            segment_compress_outputs = self.icae(
+                inputs_embeds=segment_input_embedding, output_hidden_states=True
+            )
             segment_compress_outputs = segment_compress_outputs.hidden_states[-1]
 
             # collect memory tokens
-            compress_outputs[segment_idx*self.mem_size: self.mem_size*(segment_idx+1)] = segment_compress_outputs[mem_flag]
-            
+            compress_outputs[
+                segment_idx * self.mem_size : self.mem_size * (segment_idx + 1)
+            ] = segment_compress_outputs[mem_flag]
+
             del segment_input_ids, segment_input_embedding
             torch.cuda.empty_cache()
-            
+
         # decoder part
-        decoder_mem_flag = (prompt_answer_ids >= self.vocab_size) & (prompt_answer_ids < self.vocab_size + self.mem_size)   # only mem tokens
+        decoder_mem_flag = (prompt_answer_ids >= self.vocab_size) & (
+            prompt_answer_ids < self.vocab_size + self.mem_size
+        )  # only mem tokens
 
         prompt_answer_embs[decoder_mem_flag] = compress_outputs  # replace memory slots
         special_prompt = prompt_answer_ids >= self.vocab_size_with_mem
-        prompt_answer_embs[special_prompt] = self.memory_token_embed(prompt_answer_ids[special_prompt] - self.vocab_size).to(prompt_answer_embs)    # replace special token's embedding from self.memory_token_embed
-        
-        if self.training:   # has an independent se.f.decoder
-            decoder_outputs = self.decoder(inputs_embeds=prompt_answer_embs, output_hidden_states=True)
-        else:
-            with self.icae.disable_adapter():   # no independent decoder; use self.icae
-                decoder_outputs = self.icae(inputs_embeds=prompt_answer_embs, output_hidden_states=True)
+        prompt_answer_embs[special_prompt] = self.memory_token_embed(
+            prompt_answer_ids[special_prompt] - self.vocab_size
+        ).to(
+            prompt_answer_embs
+        )  # replace special token's embedding from self.memory_token_embed
 
+        if self.training:  # has an independent se.f.decoder
+            decoder_outputs = self.decoder(
+                inputs_embeds=prompt_answer_embs, output_hidden_states=True
+            )
+        else:
+            with self.icae.disable_adapter():  # no independent decoder; use self.icae
+                decoder_outputs = self.icae(
+                    inputs_embeds=prompt_answer_embs, output_hidden_states=True
+                )
 
         logits = decoder_outputs.logits
-        effective_logits = logits[:,:-1,:].reshape(-1, logits.size(-1))
-        target_ids = labels[:,1:].reshape(-1)
+        effective_logits = logits[:, :-1, :].reshape(-1, logits.size(-1))
+        target_ids = labels[:, 1:].reshape(-1)
         loss = self.loss_fct(effective_logits, target_ids)
         return {"loss": loss, "logits": logits}
-    
-    
-    def tokens_to_embeddings(self, token_ids):   # input_tokens can be either normal tokens and special tokens
+
+    def tokens_to_embeddings(
+        self, token_ids
+    ):  # input_tokens can be either normal tokens and special tokens
         embeddings = self.icae.get_base_model().model.embed_tokens(token_ids)
         special_flags = token_ids >= self.vocab_size
-        embeddings[special_flags] = self.memory_token_embed(token_ids[special_flags] - self.vocab_size).to(embeddings)    # replace special token's embedding from self.memory_token_embed
+        embeddings[special_flags] = self.memory_token_embed(
+            token_ids[special_flags] - self.vocab_size
+        ).to(
+            embeddings
+        )  # replace special token's embedding from self.memory_token_embed
         return embeddings
-        
-    
+
     def _compress(
-        self,
-        input_ids: torch.LongTensor = None
+        self, input_ids: torch.LongTensor = None
     ):  # for inference; compress a fixed length of input into memory slots
 
         batch_size = input_ids.size(0)
         total_length = input_ids.size(1)
         num_segments = self.compute_num_segments(total_length)
         segment_length = math.ceil(total_length / num_segments)
-        
+
         max_compressed_length = num_segments * self.mem_size
         compress_outputs = torch.zeros((max_compressed_length, self.dim))
-        
+
         for segment_idx in range(num_segments):
             start_idx = segment_idx * segment_length
             end_idx = min((segment_idx + 1) * segment_length, total_length)
@@ -245,17 +288,25 @@ class ICAE(torch.nn.Module):
             segment_input_ids = torch.cat([segment_input_ids, self.append_sequence], dim=1)
             mem_flag = segment_input_ids >= self.vocab_size
 
-            segment_input_embedding = self.icae.get_base_model().model.embed_tokens(segment_input_ids)
-            segment_input_embedding[mem_flag] = self.memory_token_embed(segment_input_ids[mem_flag] - self.vocab_size).to(segment_input_embedding)
+            segment_input_embedding = self.icae.get_base_model().model.embed_tokens(
+                segment_input_ids
+            )
+            segment_input_embedding[mem_flag] = self.memory_token_embed(
+                segment_input_ids[mem_flag] - self.vocab_size
+            ).to(segment_input_embedding)
 
             # compress the current segment
-            segment_compress_outputs = self.icae(inputs_embeds=segment_input_embedding, output_hidden_states=True)
+            segment_compress_outputs = self.icae(
+                inputs_embeds=segment_input_embedding, output_hidden_states=True
+            )
             segment_compress_outputs = segment_compress_outputs.hidden_states[-1]
 
             # collect memory tokens
-            compress_outputs[segment_idx*self.mem_size: self.mem_size*(segment_idx+1)] = segment_compress_outputs[mem_flag]
-            
+            compress_outputs[
+                segment_idx * self.mem_size : self.mem_size * (segment_idx + 1)
+            ] = segment_compress_outputs[mem_flag]
+
             del segment_input_ids, segment_input_embedding
             torch.cuda.empty_cache()
-        
-        return compress_outputs
\ No newline at end of file
+
+        return compress_outputs
diff --git a/icae_v2/pretrain.py b/icae_v2/pretrain.py
index ed78a7a..eaf3ecb 100644
--- a/icae_v2/pretrain.py
+++ b/icae_v2/pretrain.py
@@ -1,10 +1,18 @@
 import transformers
-from peft import (
-    LoraConfig,
-)
 from datasets import load_dataset
-from modeling_icae_multi_span import ICAE, ModelArguments, DataArguments, TrainingArguments
-from training_utils import pretrain_tokenize_function, DataCollatorForDynamicPadding, train_model
+from modeling_icae_multi_span import (
+    ICAE,
+    DataArguments,
+    ModelArguments,
+    TrainingArguments,
+)
+from peft import LoraConfig
+from training_utils import (
+    DataCollatorForDynamicPadding,
+    pretrain_tokenize_function,
+    train_model,
+)
+
 
 def main():
     parser = transformers.HfArgumentParser((ModelArguments, DataArguments, TrainingArguments))
@@ -12,22 +20,23 @@ def main():
 
     print(model_args)
     print(data_args)
-    
-    training_args.gradient_checkpointing_kwargs = {"use_reentrant": False}  # manually add this argument in the code
+
+    training_args.gradient_checkpointing_kwargs = {
+        "use_reentrant": False
+    }  # manually add this argument in the code
 
     lora_config = LoraConfig(
-        r=model_args.lora_r,
-        lora_alpha=32,
-        lora_dropout=0.05,
-        bias="none",
-        task_type="CAUSAL_LM"
+        r=model_args.lora_r, lora_alpha=32, lora_dropout=0.05, bias="none", task_type="CAUSAL_LM"
     )
-    
-    # check model_args.mem_size and min_tokens_for_lm
-    assert (training_args.fixed_mem_size & (training_args.fixed_mem_size - 1)) == 0, "training_args.fixed_mem_size must be a power of 2"    
-    assert training_args.leave_tokens_for_lm <= training_args.min_tokens_for_lm, "leave_tokens_for_lm should be fewer than min_tokens_for_lm"
 
-    
+    # check model_args.mem_size and min_tokens_for_lm
+    assert (
+        training_args.fixed_mem_size & (training_args.fixed_mem_size - 1)
+    ) == 0, "training_args.fixed_mem_size must be a power of 2"
+    assert (
+        training_args.leave_tokens_for_lm <= training_args.min_tokens_for_lm
+    ), "leave_tokens_for_lm should be fewer than min_tokens_for_lm"
+
     memory_size = training_args.fixed_mem_size
 
     train_file = "/path/to/train/file"
@@ -35,17 +44,27 @@ def main():
 
     print("Loading dataset...")
 
-    dataset = load_dataset("json", data_files={"train": train_file, "eval": eval_file}, streaming=True) # streaming can be removed if the dataset is not very large.
+    dataset = load_dataset(
+        "json", data_files={"train": train_file, "eval": eval_file}, streaming=True
+    )  # streaming can be removed if the dataset is not very large.
     train_dataset = dataset["train"]
     eval_dataset = dataset["eval"]
 
     model = ICAE(model_args, training_args, lora_config)
     MEM_TOKENS = list(range(model.vocab_size, model.vocab_size + memory_size))
 
-    train_dataset = train_dataset.map(pretrain_tokenize_function, batched=True, batch_size=64, fn_kwargs={"model": model, "mem": MEM_TOKENS, "lm_ratio": training_args.lm_ratio})
-    eval_dataset = eval_dataset.map(pretrain_tokenize_function, batched=True, fn_kwargs={"model": model, "mem": MEM_TOKENS})   # don't add lm in the dev set.
+    train_dataset = train_dataset.map(
+        pretrain_tokenize_function,
+        batched=True,
+        batch_size=64,
+        fn_kwargs={"model": model, "mem": MEM_TOKENS, "lm_ratio": training_args.lm_ratio},
+    )
+    eval_dataset = eval_dataset.map(
+        pretrain_tokenize_function, batched=True, fn_kwargs={"model": model, "mem": MEM_TOKENS}
+    )  # don't add lm in the dev set.
 
     data_collator = DataCollatorForDynamicPadding(model.pad_token_id)
     train_model(model, train_dataset, eval_dataset, training_args, data_collator)
 
-main()
\ No newline at end of file
+
+main()
diff --git a/icae_v2/pretrained_inference.py b/icae_v2/pretrained_inference.py
index 68271a2..e744321 100644
--- a/icae_v2/pretrained_inference.py
+++ b/icae_v2/pretrained_inference.py
@@ -1,11 +1,17 @@
 import json
+import sys
+
 import torch
+from modeling_icae_multi_span import (
+    ICAE,
+    DataArguments,
+    ModelArguments,
+    TrainingArguments,
+)
+from peft import LoraConfig
+from safetensors.torch import load_file
 from tqdm import tqdm
 from transformers import HfArgumentParser
-from peft import LoraConfig
-from modeling_icae_multi_span import ICAE, ModelArguments, DataArguments, TrainingArguments
-import sys
-from safetensors.torch import load_file
 
 # Set the computation device
 device = "cuda"
@@ -16,11 +22,7 @@ model_args, data_args, training_args = parser.parse_args_into_dataclasses()
 
 # Define Lora configuration
 lora_config = LoraConfig(
-    r=512,
-    lora_alpha=32,
-    lora_dropout=model_args.lora_dropout,
-    bias="none",
-    task_type="CAUSAL_LM"
+    r=512, lora_alpha=32, lora_dropout=model_args.lora_dropout, bias="none", task_type="CAUSAL_LM"
 )
 
 # Initialize model and send it to CUDA device
@@ -29,11 +31,13 @@ model = ICAE(model_args, training_args, lora_config)
 # Load the fine-tuned checkpoint
 print(f"Loading trained checkpoint from {training_args.output_dir}")
 state_dict = load_file(training_args.output_dir)
-model.load_state_dict(state_dict, strict=False) # only load lora and memory token embeddings
+model.load_state_dict(state_dict, strict=False)  # only load lora and memory token embeddings
 
 model = model.to(device)
 
-lines = ["I don't have a favorite condiment as I don't consume food or condiments. However, I can tell you that many people enjoy condiments like ketchup, mayonnaise, mustard, soy sauce, hot sauce, and ranch dressing, among others. The favorite condiment can vary greatly from person to person, depending on their taste preferences and cultural influences."]
+lines = [
+    "I don't have a favorite condiment as I don't consume food or condiments. However, I can tell you that many people enjoy condiments like ketchup, mayonnaise, mustard, soy sauce, hot sauce, and ranch dressing, among others. The favorite condiment can vary greatly from person to person, depending on their taste preferences and cultural influences."
+]
 
 # Prepare the model for evaluation
 model.eval()
@@ -42,21 +46,23 @@ with torch.no_grad():
 
         for line in tqdm(lines):
             # Tokenize input text
-            tokenized_text = model.tokenizer(line, truncation=True,
-                                          max_length=5120, padding=False,
-                                          return_attention_mask=False)
+            tokenized_text = model.tokenizer(
+                line, truncation=True, max_length=5120, padding=False, return_attention_mask=False
+            )
             # Generate compressed outputs
-            input_ids = torch.LongTensor([tokenized_text['input_ids']]).to(device)
+            input_ids = torch.LongTensor([tokenized_text["input_ids"]]).to(device)
             memory_slots = model._compress(input_ids)
-            
+
             # prompt_output = model.tokenizer(data['prompt'], add_special_tokens=False, padding=False)
             prompt_ids = torch.LongTensor([[model.ae_token_id]]).to(device)
 
             prompt_answer_embs = model.tokens_to_embeddings(prompt_ids)
             memory_slots = memory_slots.to(prompt_answer_embs)
-                        
+
             # Concatenate and clone input embeddings
-            decoder_input_embeddings = torch.cat((memory_slots.unsqueeze(0), prompt_answer_embs), dim=1)
+            decoder_input_embeddings = torch.cat(
+                (memory_slots.unsqueeze(0), prompt_answer_embs), dim=1
+            )
             output = decoder_input_embeddings.clone()
 
             generate_text = []
@@ -64,25 +70,30 @@ with torch.no_grad():
 
             # Generate text output
             for i in range(512):
-                with model.icae.disable_adapter():   # no independent decoder; use self.icae
-                    out = model.icae(inputs_embeds=output, past_key_values=past_key_values, use_cache=True)
-                logit = out.logits[:, -1, :model.vocab_size-1]
+                with model.icae.disable_adapter():  # no independent decoder; use self.icae
+                    out = model.icae(
+                        inputs_embeds=output, past_key_values=past_key_values, use_cache=True
+                    )
+                logit = out.logits[:, -1, : model.vocab_size - 1]
                 past_key_values = out.past_key_values
 
                 next_token_id = torch.argmax(logit, dim=-1)
                 # print(next_token_id)
-                
-                if next_token_id.item() == 2:   # eos
+
+                if next_token_id.item() == 2:  # eos
                     break
 
-                output = model.icae.get_base_model().model.embed_tokens(next_token_id).unsqueeze(1).to(device)
+                output = (
+                    model.icae.get_base_model()
+                    .model.embed_tokens(next_token_id)
+                    .unsqueeze(1)
+                    .to(device)
+                )
                 generate_text.append(next_token_id.item())
 
             generated_text = model.tokenizer.decode(generate_text)
 
             # Structure output data
-            output_ = {
-                "output": generated_text
-            }
+            output_ = {"output": generated_text}
 
             f.write(json.dumps(output_) + "\n")
diff --git a/icae_v2/training_utils.py b/icae_v2/training_utils.py
index bfa5d94..9b3ca9f 100644
--- a/icae_v2/training_utils.py
+++ b/icae_v2/training_utils.py
@@ -1,10 +1,10 @@
-from transformers import Trainer
+import math
 import os
-import torch
 import random
 
+import torch
+from transformers import Trainer
 from transformers.trainer_utils import get_last_checkpoint
-import math
 
 device = torch.device("cuda" if torch.cuda.is_available() else "cpu")