Commit 2ae642d8 authored by Baber's avatar Baber
Browse files

refactor build_filter_ensemble to simplify filter creation

parent b6f38ac8
from functools import partial
from typing import Iterable, List, Optional, Union
from typing import Optional, Union
from lm_eval.api.filter import FilterEnsemble
from lm_eval.api.registry import get_filter
......@@ -8,18 +8,16 @@ from . import custom, extraction, selection, transformation
def build_filter_ensemble(
filter_name: str, components: list[tuple[str, Optional[dict]]]
filter_name: str,
components: list[tuple[str, Optional[dict[str, Union[str, int, float]]]]],
) -> FilterEnsemble:
"""
Create a filtering pipeline.
"""
filters = []
for function, kwargs in components:
if kwargs is None:
kwargs = {}
# 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)
return FilterEnsemble(name=filter_name, filters=filters)
# create filters given its name in the registry, and add each as a pipeline step
return FilterEnsemble(
name=filter_name,
filters=[
partial(get_filter(func), **(kwargs or {})) for func, kwargs in components
],
)
Markdown is supported
0% or .
You are about to add 0 people to the discussion. Proceed with caution.
Finish editing this message first!
Please register or to comment