make kwargs **kwargs again

src/explainers/base_explainer.py

@@ -1,5 +1,5 @@
 from abc import ABC, abstractmethod
-from typing import List, Optional, Union, Any
+from typing import Any, List, Optional, Union
 
 import torch
 

src/explainers/functional.py

@@ -2,7 +2,6 @@
 
 import torch
 
-from src.explainers.base_explainer import BaseExplainer
 from src.explainers.wrappers.captum_influence import CaptumSimilarity
 
 
@@ -34,15 +33,15 @@ def explain_fn_from_explainer(
     init_kwargs: Optional[Dict] = {},
     explain_kwargs: Optional[Dict] = {},
 ) -> torch.Tensor:
-        explainer = explainer_cls(
-            model=model,
-            model_id=model_id,
-            cache_dir=cache_dir,
-            train_dataset=train_dataset,
-            device=device,
-            explainer_kwargs=init_kwargs,
-        )
-        return explainer.explain(test=test_tensor, **explain_kwargs)
+    explainer = explainer_cls(
+        model=model,
+        model_id=model_id,
+        cache_dir=cache_dir,
+        train_dataset=train_dataset,
+        device=device,
+        **init_kwargs,
+    )
+    return explainer.explain(test=test_tensor, **explain_kwargs)
 
 
 def explainer_self_influence_interface(
@@ -60,7 +59,7 @@ def explainer_self_influence_interface(
         cache_dir=cache_dir,
         train_dataset=train_dataset,
         device=device,
-        explainer_kwargs=init_kwargs,
+        **init_kwargs,
     )
     return explainer.self_influence()
 

src/explainers/wrappers/captum_influence.py

@@ -1,4 +1,4 @@
-from typing import List, Optional, Union, Any, Dict
+from typing import Any, List, Optional, Union
 
 import torch
 from captum.influence import SimilarityInfluence
@@ -21,11 +21,14 @@ def __init__(
         train_dataset: torch.utils.data.Dataset,
         device: Union[str, torch.device],
         explainer_cls: type,
-        explain_kwargs: Dict[str, Any],
-        **kwargs,
+        **explain_kwargs: Any,
     ):
         super().__init__(
-            model=model, model_id=model_id, cache_dir=cache_dir, train_dataset=train_dataset, device=device,
+            model=model,
+            model_id=model_id,
+            cache_dir=cache_dir,
+            train_dataset=train_dataset,
+            device=device,
         )
         self.explainer_cls = explainer_cls
         self.explain_kwargs = explain_kwargs
@@ -66,29 +69,31 @@ def __init__(
         cache_dir: str,
         train_dataset: torch.utils.data.Dataset,
         device: Union[str, torch.device],
-        explainer_kwargs: Dict[str, Any],
+        **explainer_kwargs: Any,
     ):
         # extract and validate layer from kwargs
-        self._layer: Union[List[str], str] = None
+        self._layer: Optional[Union[List[str], str]] = None
         self.layer = explainer_kwargs.get("layers", [])
 
-        explainer_kwargs = {
-            "module": model,
-            "influence_src_dataset": train_dataset,
-            "activation_dir": cache_dir,
-            "model_id": model_id,
-            "similarity_direction": "max",
-            **explainer_kwargs,
-        }
+        # TODO: validate SimilarityInfluence kwargs
+        explainer_kwargs.update(
+            {
+                "module": model,
+                "influence_src_dataset": train_dataset,
+                "activation_dir": cache_dir,
+                "model_id": model_id,
+                "similarity_direction": "max",
+                **explainer_kwargs,
+            }
+        )
 
         super().__init__(
             model=model,
-            model_id=model_id,
             cache_dir=cache_dir,
             train_dataset=train_dataset,
             device=device,
             explainer_cls=SimilarityInfluence,
-            explain_kwargs=explainer_kwargs,
+            **explainer_kwargs,
         )
 
     @property

src/utils/validation.py

@@ -1,10 +1,8 @@
 import torch
 
 """
-This is a Python module that contains helper functions for validating input arguments.
-The plan is to collect them here and then create a universal validation decorator @validate_args 
-to check all the input arguments against the expected types specified e.g.
-as class attributes.
+This module contains utility functions for validation. The plan is to
+move the validation logic into a validation decorator at a later point.
 """
 
 

tests/explainers/test_explainers.py

@@ -69,7 +69,7 @@ def test_explain_stateful(test_id, model, dataset, explanations, test_tensor, te
         cache_dir=os.path.join("./cache", "test_id"),
         train_dataset=dataset,
         device="cpu",
-        explainer_kwargs=method_kwargs,
+        **method_kwargs,
     )
     explanations = explainer.explain(test_tensor)
     assert torch.allclose(explanations, explanations_exp), "Training data attributions are not as expected"

tests/explainers/test_self_influence.py

@@ -39,7 +39,12 @@ def test_self_influence(test_id, init_kwargs, request):
     )
 
     explainer_obj = CaptumSimilarity(
-        model=model, model_id="1", cache_dir="temp_captum2", train_dataset=rand_dataset, device="cpu", explainer_kwargs=init_kwargs
+        model=model,
+        model_id="1",
+        cache_dir="temp_captum2",
+        train_dataset=rand_dataset,
+        device="cpu",
+        **init_kwargs,
     )
     self_influence_rank_stateful = explainer_obj.self_influence()