diff --git a/data/raw_datasets/context_number_xlarge/generate_data.py b/data/raw_datasets/context_number_xlarge/generate_data.py
index 85c7698..616145c 100644
--- a/data/raw_datasets/context_number_xlarge/generate_data.py
+++ b/data/raw_datasets/context_number_xlarge/generate_data.py
@@ -1,11 +1,11 @@
-import json
-import random
-from typing import List, Dict
-import os
 import itertools
+import json
+import os
+import random
+from typing import Dict, List
 
 
-def save_jsonl(data: List[Dict], filepath: str) -> None:
+def save_jsonl(data: list[dict], filepath: str) -> None:
     """Save data to a JSONL file."""
     parent_dir = os.path.dirname(filepath)
     if parent_dir:  # Only create directories if there's a parent path
@@ -16,7 +16,7 @@ def save_jsonl(data: List[Dict], filepath: str) -> None:
             f.write("\n")
 
 
-def get_random_combinations(numbers: List[int], n: int, k: int) -> List[List[int]]:
+def get_random_combinations(numbers: list[int], n: int, k: int) -> list[list[int]]:
     """Get n random combinations of k numbers from the list."""
     # using itertools.combinations hangs with large numbers
     # so we explicitly generate the n indices
diff --git a/data/raw_datasets/context_number_xxlarge/generate_data.py b/data/raw_datasets/context_number_xxlarge/generate_data.py
index defc5bb..9b66444 100644
--- a/data/raw_datasets/context_number_xxlarge/generate_data.py
+++ b/data/raw_datasets/context_number_xxlarge/generate_data.py
@@ -1,11 +1,11 @@
-import json
-import random
-from typing import List, Dict
-import os
 import itertools
+import json
+import os
+import random
+from typing import Dict, List
 
 
-def save_jsonl(data: List[Dict], filepath: str) -> None:
+def save_jsonl(data: list[dict], filepath: str) -> None:
     """Save data to a JSONL file."""
     parent_dir = os.path.dirname(filepath)
     if parent_dir:  # Only create directories if there's a parent path
@@ -16,7 +16,7 @@ def save_jsonl(data: List[Dict], filepath: str) -> None:
             f.write("\n")
 
 
-def get_random_combinations(numbers: List[int], n: int, k: int) -> List[List[int]]:
+def get_random_combinations(numbers: list[int], n: int, k: int) -> list[list[int]]:
     """Get n random combinations of k numbers from the list."""
     # using itertools.combinations hangs with large numbers
     # so we explicitly generate the n indices
diff --git a/data/raw_datasets/context_numbers_large/generate_data.py b/data/raw_datasets/context_numbers_large/generate_data.py
index 6d57ad8..688d528 100644
--- a/data/raw_datasets/context_numbers_large/generate_data.py
+++ b/data/raw_datasets/context_numbers_large/generate_data.py
@@ -1,11 +1,11 @@
-import json
-import random
-from typing import List, Dict
-import os
 import itertools
+import json
+import os
+import random
+from typing import Dict, List
 
 
-def save_jsonl(data: List[Dict], filepath: str) -> None:
+def save_jsonl(data: list[dict], filepath: str) -> None:
     """Save data to a JSONL file."""
     parent_dir = os.path.dirname(filepath)
     if parent_dir:  # Only create directories if there's a parent path
@@ -16,7 +16,7 @@ def save_jsonl(data: List[Dict], filepath: str) -> None:
             f.write("\n")
 
 
-def get_random_combinations(numbers: List[int], n: int, k: int) -> List[List[int]]:
+def get_random_combinations(numbers: list[int], n: int, k: int) -> list[list[int]]:
     """Get n random combinations of k numbers from the list."""
     # using itertools.combinations hangs with large numbers
     # so we explicitly generate the n indices
diff --git a/data/raw_datasets/context_numbers_medium/generate_data.py b/data/raw_datasets/context_numbers_medium/generate_data.py
index 514e917..f739cf4 100644
--- a/data/raw_datasets/context_numbers_medium/generate_data.py
+++ b/data/raw_datasets/context_numbers_medium/generate_data.py
@@ -1,11 +1,11 @@
-import json
-import random
-from typing import List, Dict
-import os
 import itertools
+import json
+import os
+import random
+from typing import Dict, List
 
 
-def save_jsonl(data: List[Dict], filepath: str) -> None:
+def save_jsonl(data: list[dict], filepath: str) -> None:
     """Save data to a JSONL file."""
     parent_dir = os.path.dirname(filepath)
     if parent_dir:  # Only create directories if there's a parent path
@@ -16,7 +16,7 @@ def save_jsonl(data: List[Dict], filepath: str) -> None:
             f.write("\n")
 
 
-def get_random_combinations(numbers: List[int], n: int, k: int) -> List[List[int]]:
+def get_random_combinations(numbers: list[int], n: int, k: int) -> list[list[int]]:
     """Get n random combinations of k numbers from the list."""
     # using itertools.combinations hangs with large numbers
     # so we explicitly generate the n indices
diff --git a/data/raw_datasets/context_numbers_small/generate_data.py b/data/raw_datasets/context_numbers_small/generate_data.py
index 75add1c..813a30d 100644
--- a/data/raw_datasets/context_numbers_small/generate_data.py
+++ b/data/raw_datasets/context_numbers_small/generate_data.py
@@ -1,10 +1,10 @@
 import json
-import random
-from typing import List, Dict
 import os
+import random
+from typing import Dict, List
 
 
-def save_jsonl(data: List[Dict], filepath: str) -> None:
+def save_jsonl(data: list[dict], filepath: str) -> None:
     """Save data to a JSONL file."""
     parent_dir = os.path.dirname(filepath)
     if parent_dir:  # Only create directories if there's a parent path
diff --git a/hyperlora/intx_sft.py b/hyperlora/intx_sft.py
index 53183c5..f66523e 100644
--- a/hyperlora/intx_sft.py
+++ b/hyperlora/intx_sft.py
@@ -1,12 +1,12 @@
-from collections import defaultdict
-from copy import copy
-from functools import partial
-from importlib.resources import read_binary
 import logging
 import os
 import random
 import string
 import time
+from collections import defaultdict
+from copy import copy
+from functools import partial
+from importlib.resources import read_binary
 
 import numpy as np
 import torch
@@ -18,9 +18,10 @@ from data_utils import (
     tokenize_chat_messages,
     tokenize_ctx_text,
 )
-from datasets import load_dataset, disable_caching
+from datasets import disable_caching, load_dataset
 from model_loading import get_lora_config, get_model_and_tokenizer
 from modeling_utils import HyperLoRA, ModulatedPretrainedModel, get_hypernet_config
+from rouge_score import rouge_scorer
 from training_utils import TRAINING_TASK, train_model
 from transformers import (
     AutoModelForCausalLM,
@@ -30,7 +31,6 @@ from transformers import (
     HfArgumentParser,
     TrainingArguments,
 )
-from rouge_score import rouge_scorer
 from utils import (
     extract_cli_args,
     get_run_name,
@@ -44,10 +44,10 @@ from utils import (
 from configs import (
     ArgumentParser,
     CtxTrainingArguments,
+    DataArguments,
     ExperimentSetup,
     LoRAArguments,
     ModelArguments,
-    DataArguments,
 )
 
 logger = logging.getLogger()
diff --git a/hyperlora/training_utils.py b/hyperlora/training_utils.py
index 2700c06..0f82499 100644
--- a/hyperlora/training_utils.py
+++ b/hyperlora/training_utils.py
@@ -6,13 +6,12 @@ from enum import Enum
 import numpy as np
 from transformers import (
     GenerationConfig,
-    Trainer,
     Seq2SeqTrainer,
     Seq2SeqTrainingArguments,
+    Trainer,
 )
 from transformers.trainer_utils import get_last_checkpoint
 
-
 TRAINING_TASK = Enum("TRAINING_TASK", ["CAUSAL_LM", "COMPLETION"])
 
 logger = logging.getLogger()
diff --git a/hyperlora/utils.py b/hyperlora/utils.py
index 3ad07de..8c98da3 100644
--- a/hyperlora/utils.py
+++ b/hyperlora/utils.py
@@ -1,15 +1,14 @@
 import ast
+import logging
 import os
 import random
 import string
 import time
-import yaml
-import logging
 from contextlib import contextmanager
 from typing import Iterable, Optional
 
-
 import torch
+import yaml
 from peft import PeftConfig, PeftModel
 from peft.tuners.tuners_utils import BaseTunerLayer, check_target_module_exists
 from peft.utils import get_peft_model_state_dict
diff --git a/webui/app.py b/webui/app.py
index fbe86ee..9227985 100644
--- a/webui/app.py
+++ b/webui/app.py
@@ -1,9 +1,9 @@
-import os
 import json
+import os
 
 import torch
 import yaml
-from flask import Flask, render_template, request, abort, jsonify
+from flask import Flask, abort, jsonify, render_template, request
 from transformers import pipeline
 
 app = Flask(__name__)
@@ -27,7 +27,7 @@ def get_run_data(run_path):
     """
     results_file = os.path.join(run_path, "all_results.json")
     try:
-        with open(results_file, "r") as f:
+        with open(results_file) as f:
             data = json.load(f)
     except FileNotFoundError:
         return None
@@ -59,7 +59,7 @@ def get_generated_text_data(run_path):
         filename = f"{split}_generated_text.jsonl"
         filepath = os.path.join(run_path, filename)
         try:
-            with open(filepath, "r") as f:
+            with open(filepath) as f:
                 lines = f.readlines()
                 data = [json.loads(line) for line in lines]
             generated_data[split] = data
@@ -103,7 +103,7 @@ def visualize(run):
     # Load config.yaml from the run directory
     config_path = os.path.join(logdir, "config.yaml")
     try:
-        with open(config_path, "r") as f:
+        with open(config_path) as f:
             config = yaml.safe_load(f)
             model_name = config.get("model_name_or_path")
     except FileNotFoundError:
@@ -133,7 +133,7 @@ def load_model():
     logdir = os.path.join(TRAIN_OUTPUTS_DIR, run)
     config_path = os.path.join(logdir, "config.yaml")
     try:
-        with open(config_path, "r") as f:
+        with open(config_path) as f:
             config = yaml.safe_load(f)
             chat_model_name = config.get("model_name_or_path")
     except FileNotFoundError: