delay filter init; remove *args

lm_eval/api/filter.py

@@ -1,6 +1,6 @@
 from abc import ABC, abstractmethod
 from dataclasses import dataclass
-from typing import List
+from typing import Callable, Iterable, List, Union
 
 from lm_eval.api.instance import Instance
 
@@ -14,13 +14,13 @@ class Filter(ABC):
 
     """
 
-    def __init__(self, *args, **kwargs) -> None:
+    def __init__(self, **kwargs) -> None:
         """
         Can define custom behavior here, if an individual instantiation of a Filter class should have state.
         """
 
     @abstractmethod
-    def apply(self, resps, docs):
+    def apply(self, resps: Union[List, Iterable], docs: List[dict]) -> Iterable:
         """
         Defines the operation to perform on a list of the `inst.resps` properties of `Instance` objects.
         Should return the list of (filtered) response lists *in the same order as they were input*, e.g.
@@ -40,15 +40,16 @@ class FilterEnsemble:
     """
 
     name: str
-    filters: List[Filter]
+    filters: List[Callable[[], Filter]]
 
     def apply(self, instances: List[Instance]) -> None:
         resps, docs = zip(*((inst.resps, inst.doc) for inst in instances))
         resps, docs = list(resps), list(docs)
 
         for f in self.filters:
             # apply filters in sequence
-            resps = f.apply(resps, docs)
+            f_init = f()
+            resps = f_init.apply(resps, docs)
 
         # add the end results after filtering to filtered_requests of their respective source instances.
         # has key `self.name`: each FilterEnsemble applied in a given run should use a different name.

lm_eval/filters/__init__.py

@@ -1,4 +1,5 @@
-from typing import List
+from typing import List, Union
+from functools import partial
 
 from lm_eval.api.filter import FilterEnsemble
 from . import selection
@@ -22,7 +23,7 @@
 }
 
 
-def get_filter(filter_name):
+def get_filter(filter_name: str) -> Union[type, str]:
     if filter_name in FILTER_REGISTRY:
         return FILTER_REGISTRY[filter_name]
     else:
@@ -37,11 +38,8 @@ def build_filter_ensemble(
     """
     filters = []
     for function, kwargs in components:
-        if kwargs is None:
-            f = get_filter(function)()
-        else:
-            # create a filter given its name in the registry
-            f = get_filter(function)(**kwargs)  # TODO: pass kwargs to filters properly
+        # create a filter given its name in the registry
+        f = partial(get_filter(function), **kwargs)
         # add the filter as a pipeline step
         filters.append(f)
 

lm_eval/filters/selection.py

@@ -17,12 +17,14 @@ def apply(self, resps, docs):
 
 
 class TakeKFilter(Filter):
-    def __init__(self, *args, **kwargs) -> None:
+    def __init__(self, **kwargs) -> None:
         self.k = kwargs.pop("k")
 
-        super().__init__(*args, **kwargs)
+        super().__init__(**kwargs)
 
     def apply(self, resps, docs):
+        # need resp to be subscriptable to check below
+        resps = list(resps)
         # check we have at least k responses per doc, else we can't take the first k
         assert (
             len(resps[0]) >= self.k