fix: Bugs & improvements in explainers and explain scripts

.gitignore

@@ -79,4 +79,6 @@ logs/*
 scripts/*/bar_rank.png
 scripts/*/bar_rank.csv
 log/*
-logs_*/*
+logs_*/*
+local_eval_results/*
+scripts/bert_qnli_bench/eval_defs_notes.txt

config/eval/bert_qnli.yaml

@@ -12,8 +12,8 @@ hydra:
     n_jobs: 1
 
 bench: qnli_class_detection
-cache_dir: /data2/bareeva/Projects/quanda/cluster_output/eval_bench/qnli
-results_dir: /data2/bareeva/Projects/quanda/cluster_output/eval_results/qnli
+cache_dir: /data2/bareeva/Projects/quanda/qnli_bench
+results_dir: /data2/bareeva/Projects/quanda/cluster_output_new/eval_results/qnli
 device: cuda:0
 batch_size: 8
 

config/eval/cifar_resnet9.yaml

@@ -12,8 +12,8 @@ hydra:
     n_jobs: 1
 
 bench: cifar_class_detection
-cache_dir: /data/cluster/users/bareeva/quanda_output/eval_bench/cifar
-results_dir: /data/cluster/users/bareeva/quanda_output/eval_results/cifar
+cache_dir: /data/cluster/users/bareeva/quanda_output_new/eval_bench/cifar
+results_dir: /data/cluster/users/bareeva/quanda_output_new/eval_results/cifar
 
 #cache_dir: ./eval_bench/cifar
 #results_dir: ./eval_results/cifar

config/eval/explainer/dattri_arnoldi.yaml

@@ -6,4 +6,6 @@ kwargs:
     path: quanda.explainers.wrappers.dattri_losses.bert_classification_batched_loss
   proj_dim: 50
   max_iter: 100
-  batch_size: ${batch_size}
+  precompute_data_ratio: 0.0005
+  batch_size: 8
+  regularization: 0.01

config/eval/explainer/dattri_ekfac.yaml

@@ -5,4 +5,5 @@ kwargs:
     _target_: hydra.utils.get_method
     path: quanda.explainers.wrappers.dattri_losses.bert_classification_batched_loss
   damping: 0.0
-  batch_size: ${batch_size}
+  max_iter: 500
+  batch_size: 8

config/eval/explainer/dattri_gradcos.yaml

@@ -4,4 +4,4 @@ kwargs:
   loss_func:
     _target_: hydra.utils.get_method
     path: quanda.explainers.wrappers.dattri_losses.bert_classification_per_sample_loss
-  batch_size: ${batch_size}
+  batch_size: 4

config/eval/explainer/dattri_graddot.yaml

@@ -4,4 +4,4 @@ kwargs:
   loss_func:
     _target_: hydra.utils.get_method
     path: quanda.explainers.wrappers.dattri_losses.bert_classification_per_sample_loss
-  batch_size: ${batch_size}
+  batch_size: 4

config/eval/explainer/dattri_tracin.yaml

@@ -7,4 +7,4 @@ kwargs:
   learning_rate:
     2.0e-05
   normalized_grad: false
-  batch_size: ${batch_size}
+  batch_size: 4

config/eval/explainer/kronfluence.yaml

@@ -0,0 +1,8 @@
+name: kronfluence
+cls: quanda.explainers.wrappers.Kronfluence
+kwargs:
+  task_module:
+    _target_: quanda.explainers.wrappers.kronfluence_tasks.TextClassificationTask
+    tracked_modules: [classifier]
+  batch_size: 8
+  device: ${device}

config/eval/explainer/representer_points.yaml

@@ -6,3 +6,4 @@ kwargs:
   batch_size: 64
   load_from_disk: true
   show_progress: false
+  normalize: false

config/eval/explainer/trak.yaml

@@ -5,3 +5,4 @@ kwargs:
   proj_dim: 2048
   batch_size: 64
   load_from_disk: true
+  lambda_reg: 0.0

config/eval/mnist_lenet.yaml

@@ -12,8 +12,8 @@ hydra:
     n_jobs: 1
 
 bench: mnist_class_detection
-cache_dir: /data/cluster/users/bareeva/quanda_output/eval_bench/mnist
-results_dir: /data/cluster/users/bareeva/quanda_output/eval_results/mnist
+cache_dir: /data/cluster/users/bareeva/quanda_output_new/eval_bench/mnist
+results_dir: /data/cluster/users/bareeva/quanda_output_new/eval_results/mnist
 
 #cache_dir: ./eval_bench/mnist
 #results_dir: ./eval_results/mnist

quanda/benchmarks/base.py

@@ -4,7 +4,7 @@
 import json
 import os
 import warnings
-from abc import ABC
+from abc import ABC, abstractmethod
 from typing import Any, Callable, List, Optional, Union
 
 import datasets  # type: ignore
@@ -381,15 +381,15 @@ def train(
         train_dl = ds_handler.create_dataloader(
             dataset=obj.train_dataset,
             batch_size=batch_size,
-            shuffle=False,  # TODO: true
+            shuffle=True,
             num_workers=trainer.num_workers,
         )
         if obj.val_dataset is not None:
             val_ds_handler = get_dataset_handler(dataset=obj.val_dataset)
             val_dl = val_ds_handler.create_dataloader(
                 dataset=obj.val_dataset,
                 batch_size=batch_size,
-                shuffle=False,  # TODO: true
+                shuffle=False,
                 num_workers=trainer.num_workers,
             )
         else:
@@ -453,9 +453,8 @@ def train(
         obj.model.to(obj.device)
         obj.model.eval()
 
-        assert isinstance(obj.model, PyTorchModelHubMixin), (
-            "Model must inherit from PyTorchModelHubMixin."
-        )
+        if not isinstance(obj.model, PyTorchModelHubMixin):
+            raise TypeError("Model must inherit from PyTorchModelHubMixin.")
         if snapshot_dirs:
             obj.checkpoints = snapshot_dirs
         else:
@@ -492,9 +491,10 @@ def train_and_push_to_hub(
                 batch_size=batch_size,
                 load_meta_from_disk=load_meta_from_disk,
             )
-            assert isinstance(obj.model, PyTorchModelHubMixin), (
-                "Model must inherit from PyTorchModelHubMixin."
-            )
+            if not isinstance(obj.model, PyTorchModelHubMixin):
+                raise TypeError(
+                    "Model must inherit from PyTorchModelHubMixin."
+                )
 
             repo_id = config["ckpt"]
             num_checkpoints = int(config.get("num_checkpoints", 1))
@@ -512,6 +512,10 @@ def train_and_push_to_hub(
                         repo_id=repo_id,
                         revision=revision,
                     )
+                upload_folder(
+                    folder_path=obj.checkpoints[-1],
+                    repo_id=repo_id,
+                )
 
         pid_suffix = getattr(obj, "_pid_suffix", "")
         metadata_dir = BenchConfigParser.get_metadata_dir(
@@ -755,6 +759,7 @@ def overall_objective(self, sanity_check_results: dict) -> float:
         """
         return sum(sanity_check_results.values()) / len(sanity_check_results)
 
+    @abstractmethod
     def evaluate(
         self,
         explainer_cls: type,
@@ -766,8 +771,37 @@ def evaluate(
         use_cached_expl: bool = False,
         use_hf_expl: bool = False,
     ):
-        """Run the evaluation using the benchmark."""
-        pass
+        """Run the evaluation using the benchmark.
+
+        Parameters
+        ----------
+        explainer_cls : type
+            Explainer subclass to instantiate for this evaluation.
+        expl_kwargs : Optional[dict]
+            Extra kwargs passed through to ``explainer_cls``.
+        batch_size : int
+            Batch size used when iterating the eval dataset.
+        max_eval_n : Optional[int]
+            Cap on the number of eval samples; ``None`` means all.
+        eval_seed : int
+            Seed used when sampling the eval subset.
+        cache_dir : Optional[str]
+            Directory where explanations are cached on disk. Required when
+            ``use_cached_expl`` or ``use_hf_expl`` is ``True``.
+        use_cached_expl : bool
+            Load precomputed explanations from ``cache_dir`` instead of
+            recomputing them.
+        use_hf_expl : bool
+            Download precomputed explanations from the HF Hub into
+            ``cache_dir`` before loading.
+
+        Returns
+        -------
+        dict
+            Metric scores produced by this benchmark's metric(s).
+
+        """
+        raise NotImplementedError
 
     def _resolve_precomputed_explanations(
         self,
@@ -808,10 +842,6 @@ def _resolve_precomputed_explanations(
             return ExplanationsCache.load(path=cache_dir, device=self.device)
         return None
 
-    def save_metadata(self):
-        """Save metadata to disk."""
-        raise NotImplementedError
-
     def _prepare_explainer(
         self,
         dataset: torch.utils.data.Dataset,
@@ -967,8 +997,11 @@ def explain_and_push_to_hub(
             max_eval_n=max_eval_n,
             eval_seed=eval_seed,
         )
-        assert obj._explanations_id is not None
-        assert obj._explanations_dir is not None
+        if obj._explanations_id is None or obj._explanations_dir is None:
+            raise RuntimeError(
+                "explain() must populate _explanations_id and "
+                "_explanations_dir before pushing to the Hub."
+            )
         create_repo(
             repo_id=obj._explanations_id,
             repo_type="dataset",
@@ -1026,7 +1059,11 @@ def _iter_explanations(
             if precomputed_explanations is not None:
                 explanations = precomputed_explanations[i].to(self.device)
             else:
-                assert explainer is not None
+                if explainer is None:
+                    raise RuntimeError(
+                        "explainer must be provided when "
+                        "precomputed_explanations is None."
+                    )
                 explanations = explainer.explain(
                     test_data=inputs, targets=targets
                 )

quanda/benchmarks/config_parser.py

@@ -247,7 +247,8 @@ def load_pretrained_base(
         module_cfg = model_cfg["module"]
         module_cls = pl_modules[module_cfg["name"]]
         model = module_cls.from_pretrained_base(  # type: ignore[attr-defined]
-            pretrained_model_name=pretrained_model_name
+            pretrained_model_name=pretrained_model_name,
+            num_labels=model_cfg.get("num_labels", 2),
         )
         model.to(device)
         return model

quanda/benchmarks/downstream_eval/mislabeling_detection.py

@@ -81,10 +81,11 @@ def sanity_check(self, batch_size: int = 32) -> dict:
         """
         results = super().sanity_check(batch_size)
 
-        assert isinstance(self.train_dataset, LabelFlippingDataset), (
-            "Training dataset in Mislabeling Metric should have flipped "
-            "labels."
-        )
+        if not isinstance(self.train_dataset, LabelFlippingDataset):
+            raise TypeError(
+                "Training dataset in Mislabeling Metric should have flipped "
+                "labels."
+            )
         train_dl = torch.utils.data.DataLoader(
             self.train_dataset.filtered(self.train_dataset.transform_indices),
             batch_size=batch_size,

quanda/benchmarks/downstream_eval/shortcut_detection.py

@@ -108,9 +108,10 @@ def _extra_kwargs_from_config(
                 "SampleTransformationDataset as the training dataset."
             )
 
-        assert train_dataset.metadata.cls_idx is not None, (
-            "The training dataset must have a class index in its metadata."
-        )
+        if train_dataset.metadata.cls_idx is None:
+            raise ValueError(
+                "The training dataset must have a class index in its metadata."
+            )
 
         eval_ds_config = config["eval_dataset"]
         eval_indices = eval_ds_config["filter_indices"]
@@ -151,8 +152,14 @@ def sanity_check(self, batch_size: int = 32) -> dict:
         """
         results = super().sanity_check(batch_size)
 
-        assert isinstance(self.train_dataset, SampleTransformationDataset)
-        assert isinstance(self.eval_dataset, SampleTransformationDataset)
+        if not isinstance(self.train_dataset, SampleTransformationDataset):
+            raise TypeError(
+                "train_dataset must be a SampleTransformationDataset."
+            )
+        if not isinstance(self.eval_dataset, SampleTransformationDataset):
+            raise TypeError(
+                "eval_dataset must be a SampleTransformationDataset."
+            )
 
         train_dl = torch.utils.data.DataLoader(
             self.train_dataset.filtered(self.train_dataset.transform_indices),
@@ -277,7 +284,10 @@ def evaluate(
             )
         )
 
-        assert isinstance(self.train_dataset, SampleTransformationDataset)
+        if not isinstance(self.train_dataset, SampleTransformationDataset):
+            raise TypeError(
+                "train_dataset must be a SampleTransformationDataset."
+            )
         metric = ShortcutDetectionMetric(
             model=self.model,
             checkpoints=self.checkpoints,

quanda/benchmarks/ground_truth/linear_datamodeling.py

@@ -244,7 +244,8 @@ def train(  # type: ignore[override]
             batch_size=batch_size,
             load_meta_from_disk=load_meta_from_disk,
         )
-        assert isinstance(obj, LinearDatamodeling)
+        if not isinstance(obj, LinearDatamodeling):
+            raise TypeError("Expected a LinearDatamodeling instance.")
 
         if skip_subsets or cls._lds_skip_subsets:
             return obj
@@ -304,7 +305,8 @@ def train_subset(
             offline=True,
             device=device,
         )
-        assert isinstance(obj, LinearDatamodeling)
+        if not isinstance(obj, LinearDatamodeling):
+            raise TypeError("Expected a LinearDatamodeling instance.")
 
         pretrained_base = BenchConfigParser.load_pretrained_base(
             model_cfg=config["model"], device=device
@@ -443,7 +445,8 @@ def train_and_push_to_hub(
         finally:
             cls._push_subsets_during_train = False
             cls._lds_skip_subsets = False
-        assert isinstance(obj, LinearDatamodeling)
+        if not isinstance(obj, LinearDatamodeling):
+            raise TypeError("Expected a LinearDatamodeling instance.")
         return obj
 
     def sanity_check(self, batch_size: int = 32) -> dict:

quanda/benchmarks/resources/configs/1ea9143-bert_qnli_LDS.yaml

@@ -63,7 +63,7 @@ logger:
 id: 1ea9143-bert_qnli_LDS
 explanations_group: bert_qnli_base
 bench: LDS
-bench_save_dir: bench_out
+bench_save_dir: bench_out #/data/cluster/users/bareeva/quanda_output/train2
 log_dir: hydra_logs
 repo_id: quanda-bench-test
 cache_dir: tmp