Replace flake with ruff (EleutherAI#154)

* add ruff * explicit reexport * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * add back __all__ * add docstring for ruff * add explicit reexports * add readme * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci --------- Co-authored-by: James Chua <[email protected]> Co-authored-by: pre-commit-ci[bot] <66853113+pre-commit-ci[bot]@users.noreply.github.com>
michaelbyun · Apr 5, 2023 · 90755b6 · 90755b6
1 parent 0f9ba11
commit 90755b6
Show file tree

Hide file tree

Showing 19 changed files with 66 additions and 30 deletions.
diff --git a/.pre-commit-config.yaml b/.pre-commit-config.yaml
@@ -14,11 +14,11 @@ repos:
  rev: 23.1.0
  hooks:
  - id: black
-- repo: https://github.com/pycqa/flake8
- rev: '6.0.0'
+- repo: https://github.com/charliermarsh/ruff-pre-commit
+ rev: 'v0.0.257'
  hooks:
- - id: flake8
- args: ["--ignore=E203,F401,W503", --max-line-length=88]
+  - id: ruff
+  args: [--fix, --exit-non-zero-on-fix]
 - repo: https://github.com/codespell-project/codespell
  rev: v2.2.4
  hooks:

diff --git a/README.md b/README.md
@@ -59,6 +59,13 @@ We use [pyright](https://github.com/microsoft/pyright), which is built into the
 pyright
 ```
 
+### Run the linter
+We use [ruff](https://beta.ruff.rs/docs/). It is installed as a pre-commit hook, so you don't have to run it manually.
+If you want to run it manually, you can do so with:
+```bash
+ruff . --fix
+```
+
 ### Contributing to this repository
 
 If you work on a new feature / fix or some other code task, make sure to create an issue and assign it to yourself (Maybe, even share it in the elk channel of Eleuther's Discord with a small note). In this way, others know you are working on the issue and people won't do the same thing twice 👍 Also others can contact you easily.
diff --git a/elk/__init__.py b/elk/__init__.py
@@ -1 +1,3 @@
 from .extraction import extract_hiddens, Extract
+
+__all__ = ["extract_hiddens", "Extract"]
diff --git a/elk/__main__.py b/elk/__main__.py
@@ -1,7 +1,6 @@
 """Main entry point for `elk`."""
 
 from dataclasses import dataclass
-from pathlib import Path
 from typing import Union
 
 from simple_parsing import ArgumentParser

diff --git a/elk/evaluation/evaluate.py b/elk/evaluation/evaluate.py
@@ -1,28 +1,18 @@
-import csv
-import os
 from dataclasses import dataclass
 from functools import partial
 from pathlib import Path
-from typing import Callable, Literal, Optional, cast
+from typing import Callable, Literal, Optional
 
 import torch
-import torch.multiprocessing as mp
 from simple_parsing.helpers import Serializable, field
-from torch import Tensor
-from tqdm.auto import tqdm
 
-from datasets import DatasetDict
 from elk.evaluation.evaluate_log import EvalLog
 from elk.extraction.extraction import Extract
 from elk.run import Run
 from elk.training import Reporter
 
-from ..files import elk_reporter_dir, memorably_named_dir
-from ..training.preprocessing import normalize
+from ..files import elk_reporter_dir
 from ..utils import (
- assert_type,
- int16_to_float32,
- select_train_val_splits,
  select_usable_devices,
 )
 

diff --git a/elk/extraction/__init__.py b/elk/extraction/__init__.py
@@ -2,3 +2,15 @@
 from .extraction import Extract, extract_hiddens, extract
 from .generator import _GeneratorConfig, _GeneratorBuilder
 from .prompt_loading import PromptConfig, load_prompts
+
+__all__ = [
+ "BalancedSampler",
+ "FewShotSampler",
+ "Extract",
+ "extract_hiddens",
+ "extract",
+ "_GeneratorConfig",
+ "_GeneratorBuilder",
+ "PromptConfig",
+ "load_prompts",
+]
diff --git a/elk/extraction/balanced_sampler.py b/elk/extraction/balanced_sampler.py
@@ -2,7 +2,6 @@
 from ..utils import infer_label_column
 from ..utils.typing import assert_type
 from collections import deque
-from dataclasses import dataclass
 from datasets import IterableDataset, Features
 from itertools import cycle
 from random import Random

diff --git a/elk/extraction/extraction.py b/elk/extraction/extraction.py
@@ -24,7 +24,6 @@
 
 from ..utils import (
  assert_type,
- infer_label_column,
  select_train_val_splits,
  select_usable_devices,
 )

diff --git a/elk/extraction/prompt_dataset.py b/elk/extraction/prompt_dataset.py
diff --git a/elk/extraction/prompt_loading.py b/elk/extraction/prompt_loading.py
@@ -9,13 +9,9 @@
 from .balanced_sampler import FewShotSampler
 from dataclasses import dataclass
 from datasets import (
- interleave_datasets,
  load_dataset,
- ClassLabel,
  Dataset,
  Features,
- IterableDataset,
- Sequence,
 )
 from datasets.distributed import split_dataset_by_node
 from random import Random

diff --git a/elk/promptsource/__init__.py b/elk/promptsource/__init__.py
@@ -1 +1,3 @@
 from .templates import DatasetTemplates, Template
+
+__all__ = ["DatasetTemplates", "Template"]
diff --git a/elk/training/__init__.py b/elk/training/__init__.py
@@ -1,3 +1,14 @@
 from .ccs_reporter import CcsReporter, CcsReporterConfig
 from .eigen_reporter import EigenReporter, EigenReporterConfig
 from .reporter import OptimConfig, Reporter, ReporterConfig
+
+
+__all__ = [
+ "Reporter",
+ "ReporterConfig",
+ "CcsReporter",
+ "CcsReporterConfig",
+ "EigenReporter",
+ "EigenReporterConfig",
+ "OptimConfig",
+]
diff --git a/elk/training/ccs_reporter.py b/elk/training/ccs_reporter.py
@@ -8,7 +8,7 @@
 from dataclasses import dataclass, field
 from torch import Tensor
 from torch.nn.functional import binary_cross_entropy as bce
-from typing import cast, Literal, NamedTuple, Optional
+from typing import cast, Literal, Optional
 import math
 import torch
 import torch.nn as nn

diff --git a/elk/training/classifier.py b/elk/training/classifier.py
@@ -6,7 +6,6 @@
 from torch import Tensor
 from typing import Optional
 import torch
-import warnings
 
 
 @dataclass

diff --git a/elk/utils/__init__.py b/elk/utils/__init__.py
@@ -9,3 +9,16 @@
 from .gpu_utils import select_usable_devices
 from .tree_utils import pytree_map
 from .typing import assert_type, float32_to_int16, int16_to_float32
+
+__all__ = [
+ "binarize",
+ "get_columns_all_equal",
+ "infer_label_column",
+ "infer_num_classes",
+ "float32_to_int16",
+ "int16_to_float32",
+ "select_train_val_splits",
+ "select_usable_devices",
+ "pytree_map",
+ "assert_type",
+]
diff --git a/elk/utils/data_utils.py b/elk/utils/data_utils.py
@@ -8,9 +8,7 @@
  Value,
 )
 from random import Random
-import torch
-from typing import Iterable, Optional, List, Any
-import numpy as np
+from typing import Iterable, List, Any
 import copy
 
 

diff --git a/pyproject.toml b/pyproject.toml
@@ -55,3 +55,14 @@ testpaths = ["tests"]
 
 [tool.setuptools.packages.find]
 include = ["elk*"]
+
+[tool.ruff]
+# Enable pycodestyle (`E`) and Pyflakes (`F`) codes
+# See https://beta.ruff.rs/docs/rules/ for more possible rules
+select = ["E", "F"]
+# Same as Black.
+line-length = 88
+# Avoid automatically removing unused imports in __init__.py files.
+# Such imports will be flagged with a dedicated message suggesting
+# that the import is either added to the module's __all__ symbol
+ignore-init-module-imports = true
diff --git a/tests/test_math.py b/tests/test_math.py
@@ -3,7 +3,6 @@
 from random import Random
 import math
 import numpy as np
-import pytest
 import torch
 
 

diff --git a/tests/test_smoke_elicit.py b/tests/test_smoke_elicit.py
@@ -3,7 +3,6 @@
 from elk import Extract
 from elk.extraction import PromptConfig
 from elk.training import CcsReporterConfig, EigenReporterConfig
-from elk.training import train
 from elk.training.train import Elicit