Add pre-commit as a dev dependency and run it

Ran "uv pre-commit run --all-files" which reads from .pre-commit-config.yaml

Unfortunately pre-commit does not respect tool settings in pyproject.toml, so
right now there's conflicting informations in pyproject.toml and
.pre-commit-config.yaml and so different settings and tool versions used
depending on how we run tools.

Author: smarter

bergson/hessians/data_filtering_ekfac.ipynb

@@ -13,24 +13,17 @@
    "metadata": {},
    "outputs": [],
    "source": [
+    "import json\n",
     "import os\n",
-    "from typing import Literal\n",
     "\n",
     "import matplotlib.pyplot as plt\n",
     "import numpy as np\n",
     "import pandas as pd\n",
     "import torch\n",
-    "from datasets import load_dataset\n",
-    "from tqdm.notebook import tqdm\n",
-    "\n",
-    "from bergson.data import load_gradients\n",
     "from safetensors.torch import load_file\n",
-    "\n",
-    "import numpy as np\n",
-    "import matplotlib.pyplot as plt\n",
     "from scipy.stats import spearmanr\n",
     "\n",
-    "import json"
+    "from bergson.data import load_gradients"
    ]
   },
   {
@@ -373,8 +366,6 @@
    ],
    "source": [
     "import numpy as np\n",
-    "import matplotlib.pyplot as plt\n",
-    "from scipy.stats import spearmanr\n",
     "\n",
     "# Calculate Spearman correlation\n",
     "mask = ~(np.isnan(attributions_scores) | np.isnan(attributions_ekfac_scores))\n",
@@ -461,9 +452,9 @@
     ")\n",
     "\n",
     "\n",
-    "plt.plot(np.array(top_percentages) * len(index), intersections, label=f\"Query without ekfac\")\n",
-    "plt.plot(np.array(top_percentages) * len(index), intersections_ekfac, label=f\"Query with ekfac\")\n",
-    "plt.plot(np.array(top_percentages) * len(index), intersections_random, label=f\"Random baseline\")\n",
+    "plt.plot(np.array(top_percentages) * len(index), intersections, label=\"Query without ekfac\")\n",
+    "plt.plot(np.array(top_percentages) * len(index), intersections_ekfac, label=\"Query with ekfac\")\n",
+    "plt.plot(np.array(top_percentages) * len(index), intersections_random, label=\"Random baseline\")\n",
     "plt.xlabel(\"Number of elements removed\")\n",
     "plt.ylabel('Number of elements in the \"correct\" half')\n",
     "plt.title(\"EK-FAC, no attn, on train set\")\n",
@@ -638,7 +629,6 @@
     "# load the saved attributions\n",
     "import json\n",
     "\n",
-    "\n",
     "all_attributions = {}\n",
     "\n",
     "for path in all_query_paths:\n",
@@ -849,7 +839,7 @@
    ],
    "source": [
     "# plot intersection\n",
-    "plt.plot(np.array(top_percentages) * len(index), intersection_12, label=f\"Intersection\")\n",
+    "plt.plot(np.array(top_percentages) * len(index), intersection_12, label=\"Intersection\")\n",
     "plt.plot(\n",
     "    [0, len(index) // 2, len(index)],\n",
     "    [0, len(index) // 2, len(index)],\n",
@@ -1218,9 +1208,9 @@
     ")\n",
     "\n",
     "\n",
-    "plt.plot(np.array(top_percentages) * len(index), intersections, label=f\"Query without ekfac\")\n",
-    "plt.plot(np.array(top_percentages) * len(index), intersections_ekfac, label=f\"Query with ekfac\")\n",
-    "plt.plot(np.array(top_percentages) * len(index), intersections_random, label=f\"Random baseline\")\n",
+    "plt.plot(np.array(top_percentages) * len(index), intersections, label=\"Query without ekfac\")\n",
+    "plt.plot(np.array(top_percentages) * len(index), intersections_ekfac, label=\"Query with ekfac\")\n",
+    "plt.plot(np.array(top_percentages) * len(index), intersections_random, label=\"Random baseline\")\n",
     "plt.xlabel(\"Number of elements removed\")\n",
     "plt.ylabel('Number of elements in the \"correct\" half')\n",
     "plt.legend()\n",
@@ -1374,9 +1364,10 @@
     }
    ],
    "source": [
-    "import torch\n",
     "import os\n",
     "\n",
+    "import torch\n",
+    "\n",
     "# Set the debug flag - this is the correct way\n",
     "os.environ[\"TORCH_COMPILE_DEBUG\"] = \"1\"\n",
     "\n",

bergson/hessians/misaligned_datasets.ipynb

@@ -13,23 +13,16 @@
    "metadata": {},
    "outputs": [],
    "source": [
+    "import json\n",
     "import os\n",
-    "from typing import Literal\n",
-    "import joblib\n",
+    "\n",
     "import matplotlib.pyplot as plt\n",
     "import numpy as np\n",
     "import pandas as pd\n",
     "import torch\n",
-    "from datasets import load_dataset\n",
-    "from tqdm.notebook import tqdm\n",
-    "import json\n",
-    "from datasets import Dataset\n",
-    "from bergson.data import load_gradients\n",
-    "from safetensors.torch import load_file\n",
-    "from sklearn.metrics import roc_auc_score\n",
-    "import numpy as np\n",
-    "\n",
-    "from sklearn.metrics import roc_auc_score, precision_recall_curve, auc\n"
+    "from datasets import Dataset, load_dataset\n",
+    "from sklearn.metrics import auc, precision_recall_curve, roc_auc_score\n",
+    "\n"
    ]
   },
   {
@@ -439,14 +432,14 @@
     "print(f\"PR AUC: {pr_auc:.4f}\")\n",
     "\n",
     "# Additional metrics for analysis\n",
-    "print(f\"\\nDataset composition:\")\n",
+    "print(\"\\nDataset composition:\")\n",
     "print(f\"Correct examples: {len(sorted_correct_scores)}\")\n",
     "print(f\"Incorrect examples: {len(sorted_incorrect_scores)}\")\n",
     "print(f\"Subtle incorrect examples: {len(sorted_subtle_scores)}\")\n",
     "print(f\"Total examples: {len(all_scores)}\")\n",
     "print(f\"Problematic ratio: {(len(sorted_incorrect_scores) + len(sorted_subtle_scores)) / len(all_scores):.3f}\")\n",
     "\n",
-    "print(f\"\\nScore statistics:\")\n",
+    "print(\"\\nScore statistics:\")\n",
     "print(f\"Correct scores - Mean: {sorted_correct_scores.mean():.4f}, Std: {sorted_correct_scores.std():.4f}\")\n",
     "print(f\"Incorrect scores - Mean: {sorted_incorrect_scores.mean():.4f}, Std: {sorted_incorrect_scores.std():.4f}\")\n",
     "print(f\"Subtle scores - Mean: {sorted_subtle_scores.mean():.4f}, Std: {sorted_subtle_scores.std():.4f}\")\n"

pyproject.toml

@@ -56,6 +56,8 @@ line-length = 120
 
 [dependency-groups]
 dev = [
+    "pre-commit>=4.2.0",
+    "pre-commit-uv>=4.1.5",
     "pyright>=1.1.406",
     "pytest>=8.4.2",
 ]

tests/ekfac_tests/apply_ekfac_ground_truth.ipynb

@@ -23,35 +23,19 @@
    "metadata": {},
    "outputs": [],
    "source": [
-    "import gc\n",
     "import hashlib\n",
     "import json\n",
     "import os\n",
-    "import random\n",
-    "from contextlib import nullcontext\n",
-    "from typing import Literal, Optional\n",
+    "from typing import Literal\n",
     "\n",
-    "import numpy as np\n",
     "import torch\n",
-    "import torch.distributed as dist\n",
-    "import torch.nn.functional as F\n",
     "from datasets import Dataset\n",
-    "from jaxtyping import Float\n",
-    "from safetensors import safe_open\n",
     "from safetensors.torch import load_file, save_file\n",
     "from torch import Tensor\n",
     "\n",
-    "from tqdm.auto import tqdm\n",
-    "from transformers import PreTrainedModel\n",
-    "\n",
     "from bergson.collection import collect_gradients\n",
-    "from bergson.data import DataConfig, IndexConfig, create_index, load_gradients, pad_and_tensor\n",
-    "from bergson.distributed import distributed_computing, setup_data_pipeline\n",
-    "from bergson.gradients import (\n",
-    "    GradientProcessor,\n",
-    ")\n",
-    "from bergson.hessians.collector import EkfacCollector\n",
-    "from bergson.hessians.logger import get_logger"
+    "from bergson.data import DataConfig, IndexConfig, load_gradients\n",
+    "from bergson.distributed import distributed_computing, setup_data_pipeline"
    ]
   },
   {