Turn on import sorting (#168)

elk/__init__.py

@@ -1,3 +1,3 @@
-from .extraction import extract_hiddens, Extract
+from .extraction import Extract, extract_hiddens
 
 __all__ = ["extract_hiddens", "Extract"]

elk/calibration.py

@@ -1,8 +1,9 @@
+import warnings
 from dataclasses import dataclass, field
-from torch import Tensor
 from typing import NamedTuple
+
 import torch
-import warnings
+from torch import Tensor
 
 
 class CalibrationEstimate(NamedTuple):

elk/eigsh.py

@@ -1,7 +1,8 @@
-from torch import Tensor
 from typing import Literal, Optional
+
 import torch
 import torch.nn.functional as F
+from torch import Tensor
 
 
 def lanczos_eigsh(

elk/extraction/__init__.py

@@ -1,6 +1,6 @@
 from .balanced_sampler import BalancedSampler, FewShotSampler
-from .extraction import Extract, extract_hiddens, extract
-from .generator import _GeneratorConfig, _GeneratorBuilder
+from .extraction import Extract, extract, extract_hiddens
+from .generator import _GeneratorBuilder, _GeneratorConfig
 from .prompt_loading import PromptConfig, load_prompts
 
 __all__ = [

elk/extraction/balanced_sampler.py

@@ -1,12 +1,14 @@
-from ..math_util import stochastic_round_constrained
-from ..utils import infer_label_column
-from ..utils.typing import assert_type
 from collections import deque
-from datasets import IterableDataset, Features
 from itertools import cycle
 from random import Random
+from typing import Iterable, Iterator, Optional
+
+from datasets import Features, IterableDataset
 from torch.utils.data import IterableDataset as TorchIterableDataset
-from typing import Iterator, Optional, Iterable
+
+from ..math_util import stochastic_round_constrained
+from ..utils import infer_label_column
+from ..utils.typing import assert_type
 
 
 class BalancedSampler(TorchIterableDataset):

elk/extraction/extraction.py

@@ -6,9 +6,6 @@
 from typing import Iterable, Literal, Optional, Union
 
 import torch
-from simple_parsing import Serializable, field
-from transformers import AutoConfig, AutoModel, AutoTokenizer, PreTrainedModel
-
 from datasets import (
     Array3D,
     ClassLabel,
@@ -20,6 +17,9 @@
     Value,
     get_dataset_config_info,
 )
+from simple_parsing import Serializable, field
+from transformers import AutoConfig, AutoModel, AutoTokenizer, PreTrainedModel
+
 from elk.utils.typing import float32_to_int16
 
 from ..utils import (

elk/extraction/generator.py

@@ -1,5 +1,5 @@
 from dataclasses import dataclass, field
-from typing import Callable, Optional, Any, Dict
+from typing import Any, Callable, Dict, Optional
 
 import datasets
 from datasets.splits import NamedSplit

elk/extraction/prompt_loading.py

@@ -1,3 +1,15 @@
+from dataclasses import dataclass
+from random import Random
+from typing import Any, Iterator, Literal, Optional
+
+from datasets import (
+    Dataset,
+    Features,
+    load_dataset,
+)
+from datasets.distributed import split_dataset_by_node
+from simple_parsing.helpers import Serializable, field
+
 from ..promptsource import DatasetTemplates
 from ..utils import (
     assert_type,
@@ -7,16 +19,6 @@
     select_train_val_splits,
 )
 from .balanced_sampler import FewShotSampler
-from dataclasses import dataclass
-from datasets import (
-    load_dataset,
-    Dataset,
-    Features,
-)
-from datasets.distributed import split_dataset_by_node
-from random import Random
-from simple_parsing.helpers import field, Serializable
-from typing import Any, Iterator, Literal, Optional
 
 
 @dataclass

elk/files.py

@@ -1,13 +1,13 @@
 """Helper functions for dealing with files."""
 
-from pathlib import Path
 import json
 import os
 import random
+from pathlib import Path
 from typing import Optional
 
-from simple_parsing import Serializable
 import yaml
+from simple_parsing import Serializable
 
 
 def elk_reporter_dir() -> Path:

elk/logging.py

@@ -1,4 +1,5 @@
 import logging
+
 from .utils import select_train_val_splits
 
 

elk/math_util.py

@@ -1,7 +1,8 @@
-from torch import Tensor
 import math
 import random
+
 import torch
+from torch import Tensor
 
 
 @torch.jit.script

elk/parsing.py

@@ -1,4 +1,5 @@
 import re
+
 from .training.losses import LOSSES
 
 

elk/promptsource/templates.py

@@ -1,14 +1,14 @@
-from collections import Counter, defaultdict
-from jinja2 import BaseLoader, Environment, meta
-from pathlib import Path
-from shutil import rmtree
-from typing import Optional
 import logging
 import os
 import random
 import uuid
-import yaml
+from collections import Counter, defaultdict
+from pathlib import Path
+from shutil import rmtree
+from typing import Optional
 
+import yaml
+from jinja2 import BaseLoader, Environment, meta
 
 # Truncation of jinja template variables
 # 1710 = 300 words x 4.7 avg characters per word + 300 spaces

elk/run.py

@@ -5,10 +5,10 @@
 from pathlib import Path
 from typing import (
     TYPE_CHECKING,
-    Optional,
-    Union,
     Callable,
     Iterator,
+    Optional,
+    Union,
 )
 
 import numpy as np
@@ -21,7 +21,7 @@
 from elk.extraction.extraction import extract
 from elk.files import create_output_directory, save_config, save_meta
 from elk.training.preprocessing import normalize
-from elk.utils.csv import write_iterator_to_file, Log
+from elk.utils.csv import Log, write_iterator_to_file
 from elk.utils.data_utils import get_layers, select_train_val_splits
 from elk.utils.typing import assert_type, int16_to_float32
 

elk/training/__init__.py

@@ -2,7 +2,6 @@
 from .eigen_reporter import EigenReporter, EigenReporterConfig
 from .reporter import OptimConfig, Reporter, ReporterConfig
 
-
 __all__ = [
     "Reporter",
     "ReporterConfig",

elk/training/ccs_reporter.py

@@ -1,17 +1,19 @@
 """An ELK reporter network."""
 
-from ..parsing import parse_loss
-from ..utils.typing import assert_type
-from .losses import LOSSES
-from .reporter import Reporter, ReporterConfig
+import math
 from copy import deepcopy
 from dataclasses import dataclass, field
-from torch import Tensor
-from torch.nn.functional import binary_cross_entropy as bce
-from typing import cast, Literal, Optional
-import math
+from typing import Literal, Optional, cast
+
 import torch
 import torch.nn as nn
+from torch import Tensor
+from torch.nn.functional import binary_cross_entropy as bce
+
+from ..parsing import parse_loss
+from ..utils.typing import assert_type
+from .losses import LOSSES
+from .reporter import Reporter, ReporterConfig
 
 
 @dataclass

elk/training/classifier.py

@@ -1,11 +1,14 @@
 from dataclasses import dataclass
+from typing import Optional
+
+import torch
+from torch import Tensor
 from torch.nn.functional import (
     binary_cross_entropy_with_logits as bce_with_logits,
+)
+from torch.nn.functional import (
     cross_entropy,
 )
-from torch import Tensor
-from typing import Optional
-import torch
 
 
 @dataclass

elk/training/eigen_reporter.py

@@ -1,12 +1,14 @@
 """An ELK reporter network."""
 
-from ..math_util import cov_mean_fused
-from ..eigsh import lanczos_eigsh
-from .reporter import Reporter, ReporterConfig
 from dataclasses import dataclass
-from torch import nn, optim, Tensor
 from typing import Optional
+
 import torch
+from torch import Tensor, nn, optim
+
+from ..eigsh import lanczos_eigsh
+from ..math_util import cov_mean_fused
+from .reporter import Reporter, ReporterConfig
 
 
 @dataclass

elk/training/losses.py

@@ -1,10 +1,11 @@
 """Loss functions for training reporters."""
 
-from torch import Tensor
-import torch
 import warnings
 from inspect import signature
 
+import torch
+from torch import Tensor
+
 LOSSES = dict()  # Registry of loss functions
 
 

elk/training/preprocessing.py

@@ -1,6 +1,7 @@
 """Preprocessing functions for training."""
 
 from typing import Literal
+
 import torch
 
 

elk/training/reporter.py

@@ -1,16 +1,18 @@
 """An ELK reporter network."""
 
-from ..calibration import CalibrationError
-from .classifier import Classifier
 from abc import ABC, abstractmethod
 from dataclasses import dataclass
 from pathlib import Path
-from simple_parsing.helpers import Serializable
-from sklearn.metrics import roc_auc_score
-from torch import Tensor
 from typing import Literal, NamedTuple, Optional, Union
+
 import torch
 import torch.nn as nn
+from simple_parsing.helpers import Serializable
+from sklearn.metrics import roc_auc_score
+from torch import Tensor
+
+from ..calibration import CalibrationError
+from .classifier import Classifier
 
 
 class EvalResult(NamedTuple):

elk/training/train.py

@@ -5,7 +5,7 @@
 from dataclasses import dataclass
 from functools import partial
 from pathlib import Path
-from typing import Literal, Optional, Callable
+from typing import Callable, Literal, Optional
 
 import torch
 from simple_parsing import Serializable, field, subgroups
@@ -16,12 +16,12 @@
 from elk.run import Run
 from elk.utils.typing import assert_type
 
+from ..utils import select_usable_devices
 from .ccs_reporter import CcsReporter, CcsReporterConfig
 from .classifier import Classifier
 from .eigen_reporter import EigenReporter, EigenReporterConfig
 from .reporter import OptimConfig, Reporter, ReporterConfig
 from .train_log import ElicitLog
-from ..utils import select_usable_devices
 
 
 @dataclass

elk/utils/__init__.py

@@ -5,7 +5,6 @@
     infer_num_classes,
     select_train_val_splits,
 )
-
 from .gpu_utils import select_usable_devices
 from .tree_utils import pytree_map
 from .typing import assert_type, float32_to_int16, int16_to_float32

elk/utils/csv.py

@@ -1,6 +1,6 @@
 import csv
 from pathlib import Path
-from typing import Iterator, Callable, TextIO, TypeVar
+from typing import Callable, Iterator, TextIO, TypeVar
 
 from datasets import DatasetDict
 

elk/utils/data_utils.py

@@ -1,15 +1,17 @@
-from .typing import assert_type
-from ..promptsource.templates import Template
+import copy
+from random import Random
+from typing import Any, Iterable, List
+
 from datasets import (
     ClassLabel,
     DatasetDict,
     Features,
     Split,
     Value,
 )
-from random import Random
-from typing import Iterable, List, Any
-import copy
+
+from ..promptsource.templates import Template
+from .typing import assert_type
 
 
 def get_columns_all_equal(dataset: DatasetDict) -> list[str]:

elk/utils/gpu_utils.py

@@ -1,12 +1,14 @@
 """Utilities that use PyNVML to get GPU usage info, and select GPUs accordingly."""
 
-from .typing import assert_type
-from typing import Optional
 import os
+import time
+import warnings
+from typing import Optional
+
 import pynvml
 import torch
-import warnings
-import time
+
+from .typing import assert_type
 
 
 def select_usable_devices(