[pre-commit.ci] auto fixes from pre-commit.com hooks

for more information, see https://pre-commit.ci

elk/__main__.py

@@ -22,7 +22,6 @@ def sweep(args):
  Example: elk sweep --models google/t5-v1_1-base microsoft/deberta-v2-xxlarge-mnli --datasets imdb amazon_polarity --max-examples 100
  """
 
-
  # extract and train
  names = []
  for model in args.models:
@@ -41,25 +40,19 @@ def sweep(args):
  args.hidden_states = names
  evaluate(args)
 
+
 # TODO: Move function to a better place...
 def get_sweep_parser():
  parser = ArgumentParser(add_help=False)
  add_sweep_args(parser)
  return parser
 
+
 # TODO: Move function to a better place...
 def add_sweep_args(parser):
+ parser.add_argument("--models", nargs="+", type=str, help="Model to sweep over.")
  parser.add_argument(
- "--models",
- nargs="+",
- type=str,
- help="Model to sweep over."
- )
- parser.add_argument(
- "--datasets",
- nargs="+",
- type=str,
- help="Dataset to sweep over."
+ "--datasets", nargs="+", type=str, help="Dataset to sweep over."
  )
 
 
@@ -70,7 +63,6 @@ def sweep(args):
  Example: elk sweep --models google/t5-v1_1-base microsoft/deberta-v2-xxlarge-mnli --datasets imdb amazon_polarity --max-examples 100
  """
 
-
  # extract and train
  names = []
  for model in args.models:
@@ -89,25 +81,19 @@ def sweep(args):
  args.hidden_states = names
  evaluate(args)
 
+
 # TODO: Move function to a better place...
 def get_sweep_parser():
  parser = ArgumentParser(add_help=False)
  add_sweep_args(parser)
  return parser
 
+
 # TODO: Move function to a better place...
 def add_sweep_args(parser):
+ parser.add_argument("--models", nargs="+", type=str, help="Model to sweep over.")
  parser.add_argument(
- "--models",
- nargs="+",
- type=str,
- help="Model to sweep over."
- )
- parser.add_argument(
- "--datasets",
- nargs="+",
- type=str,
- help="Dataset to sweep over."
+ "--datasets", nargs="+", type=str, help="Dataset to sweep over."
  )
 
 

elk/evaluate/evaluate.py

@@ -15,8 +15,12 @@ def evaluate(args):
 
  _, hiddens = normalize(hiddens, hiddens, args.normalization)
 
- reporter_root_path = elk_cache_dir() / args.name / "reporters" / args.reporter_name
- reporters = torch.load(reporter_root_path / "reporters.pt", map_location=args.device)
+ reporter_root_path = (
+ elk_cache_dir() / args.name / "reporters" / args.reporter_name
+ )
+ reporters = torch.load(
+ reporter_root_path / "reporters.pt", map_location=args.device
+ )
 
  L = hiddens.shape[1]
 
@@ -59,4 +63,4 @@ def evaluate(args):
  for i, stats in enumerate(statistics):
  writer.writerow([L - i] + [s for s in stats])
 
- print("Evaluation done.")
+ print("Evaluation done.")

elk/evaluate/parser.py

@@ -13,14 +13,12 @@ def add_eval_args(parser):
  parser.add_argument(
  "name",
  type=str,
- help="Name of the experiment containing" 
- "the reporters you want to evaluate.",
+ help="Name of the experiment containing" "the reporters you want to evaluate.",
  )
  parser.add_argument(
  "reporter_name",
  type=str,
- help="Name of the reporter subfolder"
- "to save the trained reporters to.",
+ help="Name of the reporter subfolder" "to save the trained reporters to.",
  )
  parser.add_argument(
  "--hidden-states",

elk/training/parser.py

@@ -13,9 +13,8 @@ def add_train_args(parser: ArgumentParser):
  parser.add_argument(
  "--reporter-name",
  type=str,
- help="Name of the reporter subfolder"
- "to save the trained reporters to.",
- default=None
+ help="Name of the reporter subfolder" "to save the trained reporters to.",
+ default=None,
  )
  parser.add_argument(
  "--device",

elk/training/train.py

@@ -38,7 +38,7 @@ def train(args):
  rank = dist.get_rank() if dist.is_initialized() else 0
  if dist.is_initialized() and not args.skip_baseline and rank == 0:
  print("Skipping LR baseline during distributed training.")
- 
+
  if not args.reporter_name:
  args.reporter_name = args_to_uuid(args)
  print("args.reporter_name", args.reporter_name)
@@ -161,7 +161,7 @@ def train(args):
 
  reporters.reverse()
  lr_models.reverse()
- 
+
  path = elk_cache_dir() / args.name / "reporters" / args.reporter_name
  path.mkdir(parents=True, exist_ok=True)