Added short benchmark description and moved imports to speed up help

Author: PierreMarza

src/thunder/benchmark.py

@@ -1,20 +1,7 @@
-import logging
-import os
-import shutil
 from typing import Callable
 
-import h5py
-import hydra
 from omegaconf import DictConfig
 
-from .datasets.utils import is_dataset_available
-from .models.utils import (
-    is_model_available,
-    load_custom_dataset_from_file,
-    load_custom_model_from_file,
-)
-from .utils.utils import print_task_hyperparams, save_config
-
 
 def benchmark(
     model: str | Callable,
@@ -51,7 +38,14 @@ def benchmark(
     from hydra import compose, initialize
     from omegaconf import OmegaConf
 
+    from .datasets.utils import is_dataset_available
+    from .models.utils import (
+        is_model_available,
+        load_custom_dataset_from_file,
+        load_custom_model_from_file,
+    )
     from .utils.config import get_config
+    from .utils.utils import print_task_hyperparams
 
     wandb_mode = "online" if online_wandb else "offline"
     adaptation_type = "lora" if lora else "frozen"
@@ -119,6 +113,11 @@ def run_benchmark(cfg: DictConfig, model_cls: Callable = None) -> None:
     :param cfg: config defining the job to run.
     """
 
+    import logging
+    import os
+    import shutil
+
+    import h5py
     import numpy as np
     import torch
     import wandb
@@ -136,7 +135,7 @@ def run_benchmark(cfg: DictConfig, model_cls: Callable = None) -> None:
     from .utils.constants import UtilsConstants
     from .utils.data import get_data, h5_to_np, load_embeddings
     from .utils.dice_loss import multiclass_dice_loss
-    from .utils.utils import set_seed
+    from .utils.utils import save_config, set_seed
 
     # Getting device
     device = torch.device("cuda" if torch.cuda.is_available() else "cpu")

src/thunder/datasets/data_splits.py

@@ -1,9 +1,5 @@
-import logging
 import os
-import random
-from collections import defaultdict
 from collections.abc import Callable
-from pathlib import Path
 from typing import List, Union
 
 
@@ -86,6 +82,8 @@ def generate_splits_for_dataset(dataset_name: str) -> None:
     Args:
         dataset_name (str): The name of the dataset to generate splits for.
     """
+    from pathlib import Path
+
     from omegaconf import OmegaConf
 
     from ..utils.constants import DatasetConstants
@@ -277,6 +275,9 @@ def create_few_shot_training_data(
     :param data_splits: data splits dictionary.
     :param nb_sets_per_nb_shot: number of data sets to create for each number of shots.
     """
+    import random
+    from collections import defaultdict
+
     # Creating lebel2 image dict
     train_images, train_labels = (
         data_splits["train"]["images"],
@@ -313,6 +314,8 @@ def generate_data_splits(
     :param dataset_yaml: path to the dataset yaml config file.
     :param split_function: function to run to generate the data splits.
     """
+    import logging
+
     from omegaconf import OmegaConf
 
     dataset_folder = os.path.join(base_folder, dataset_name)

src/thunder/datasets/download.py

@@ -1,9 +1,6 @@
-import logging
 import os
-from pathlib import Path
 from typing import List, Union
 
-from .data_splits import generate_splits
 from .dataset import *
 
 
@@ -39,6 +36,8 @@ def download_datasets(datasets: Union[List[str], str], make_splits: bool = False
         datasets (List[str] or str): A dataset name string or a List of dataset names to download or one of the following aliases: `all`, `classification`, `segmentation`.
         make_splits (bool): Whether to generate data splits for the datasets. Defaults to False.
     """
+    from .data_splits import generate_splits
+
     if "THUNDER_BASE_DATA_FOLDER" not in os.environ:
         raise EnvironmentError(
             "Please set base data directory of thunder using `export THUNDER_BASE_DATA_FOLDER=/base/data/directory`"
@@ -105,6 +104,9 @@ def download_datasets(datasets: Union[List[str], str], make_splits: bool = False
 
 
 def download_dataset(dataset: str):
+    import logging
+    from pathlib import Path
+
     root_folder = os.path.join(
         os.environ["THUNDER_BASE_DATA_FOLDER"], f"datasets/{dataset}"
     )

src/thunder/main.py

@@ -86,6 +86,7 @@ def benchmark(
     ] = False,
     kwargs: Annotated[List[str], typer.Argument(help="Additional arguments")] = None,
 ):
+    """Benchmark a model on a dataset for a task."""
     from . import benchmark
 
     if "THUNDER_BASE_DATA_FOLDER" not in os.environ:

src/thunder/models/download.py

@@ -1,10 +1,6 @@
 import logging
-import os
-from pathlib import Path
 from typing import List, Union
 
-from ..utils.utils import wget_download
-
 # Configure logging
 logging.basicConfig(
     level=logging.INFO, format="%(asctime)s - %(levelname)s - %(message)s"
@@ -185,8 +181,13 @@ def download_models(models: Union[List[str], str]) -> None:
 
 
 def download_model(model: str) -> None:
+    import os
+    from pathlib import Path
+
     from huggingface_hub import hf_hub_download
 
+    from ..utils.utils import wget_download
+
     if "THUNDER_BASE_DATA_FOLDER" not in os.environ:
         raise EnvironmentError(
             "Please set base data directory of thunder using `export THUNDER_BASE_DATA_FOLDER=/base/data/directory`"