Datasets 1.12 dataset.filter TypeError: get_indices_from_mask_function() got an unexpected keyword argument
Opened by timothyjlaurent (closed)

## Describe the bug
Upgrading to 1.12 caused `dataset.filter` call to fail with 

> get_indices_from_mask_function() got an unexpected keyword argument valid_rel_labels


## Steps to reproduce the bug
```pythondef 

filter_good_rows(
    ex: Dict,
    valid_rel_labels: Set[str],
    valid_ner_labels: Set[str],
    tokenizer: PreTrainedTokenizerFast,
) -> bool:
    """Get the good rows"""
    encoding = get_encoding_for_text(text=ex["text"], tokenizer=tokenizer)
    ex["encoding"] = encoding
    for relation in ex["relations"]:
        if not is_valid_relation(relation, valid_rel_labels):
            return False
    for span in ex["spans"]:
        if not is_valid_span(span, valid_ner_labels, encoding):
            return False
    return True
    
def get_dataset():    
    loader_path = str(Path(__file__).parent / "prodigy_dataset_builder.py")
    ds = load_dataset(
        loader_path,
        name="prodigy-dataset",
        data_files=sorted(file_paths),
        cache_dir=cache_dir,
    )["train"]

    valid_ner_labels = set(vocab.ner_category)
    valid_relations = set(vocab.relation_types.keys())
    ds = ds.filter(
        filter_good_rows,
        fn_kwargs=dict(
            valid_rel_labels=valid_relations,
            valid_ner_labels=valid_ner_labels,
            tokenizer=vocab.tokenizer,
        ),
        keep_in_memory=True,
        num_proc=num_proc,
    )

```

`ds` is a `DatasetDict` produced by a jsonl dataset.
This runs fine on 1.11 but fails on 1.12

**Stack Trace**



## Expected results

I expect 1.12 datasets filter to filter the dataset without raising as it does on 1.11

## Actual results
```
tf_ner_rel_lib/dataset.py:695: in load_prodigy_arrow_datasets_from_jsonl
    ds = ds.filter(
../../../../.pyenv/versions/tf_ner_rel_lib/lib/python3.8/site-packages/datasets/arrow_dataset.py:185: in wrapper
    out: Union["Dataset", "DatasetDict"] = func(self, *args, **kwargs)
../../../../.pyenv/versions/tf_ner_rel_lib/lib/python3.8/site-packages/datasets/fingerprint.py:398: in wrapper
    out = func(self, *args, **kwargs)
../../../../.pyenv/versions/tf_ner_rel_lib/lib/python3.8/site-packages/datasets/arrow_dataset.py:2169: in filter
    indices = self.map(
../../../../.pyenv/versions/tf_ner_rel_lib/lib/python3.8/site-packages/datasets/arrow_dataset.py:1686: in map
    return self._map_single(
../../../../.pyenv/versions/tf_ner_rel_lib/lib/python3.8/site-packages/datasets/arrow_dataset.py:185: in wrapper
    out: Union["Dataset", "DatasetDict"] = func(self, *args, **kwargs)
../../../../.pyenv/versions/tf_ner_rel_lib/lib/python3.8/site-packages/datasets/fingerprint.py:398: in wrapper
    out = func(self, *args, **kwargs)
../../../../.pyenv/versions/tf_ner_rel_lib/lib/python3.8/site-packages/datasets/arrow_dataset.py:2048: in _map_single
    batch = apply_function_on_filtered_inputs(
_ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ 

inputs = {'_input_hash': [2108817714, 1477695082, -1021597032, 2130671338, -1260483858, -1203431639, ...], '_task_hash': [18070...ons', 'relations', 'relations', ...], 'answer': ['accept', 'accept', 'accept', 'accept', 'accept', 'accept', ...], ...}
indices = [0, 1, 2, 3, 4, 5, ...], check_same_num_examples = False, offset = 0

    def apply_function_on_filtered_inputs(inputs, indices, check_same_num_examples=False, offset=0):
        """Utility to apply the function on a selection of columns."""
        nonlocal update_data
        fn_args = [inputs] if input_columns is None else [inputs[col] for col in input_columns]
        if offset == 0:
            effective_indices = indices
        else:
            effective_indices = [i + offset for i in indices] if isinstance(indices, list) else indices + offset
        processed_inputs = (
>           function(*fn_args, effective_indices, **fn_kwargs) if with_indices else function(*fn_args, **fn_kwargs)
        )
E       TypeError: get_indices_from_mask_function() got an unexpected keyword argument 'valid_rel_labels'

../../../../.pyenv/versions/tf_ner_rel_lib/lib/python3.8/site-packages/datasets/arrow_dataset.py:1939: TypeError
```

## Environment info
<!-- You can run the command `datasets-cli env` and copy-and-paste its output below. -->
- `datasets` version: 1.12.1
- Platform: Mac
- Python version: 3.8.9
- PyArrow version: pyarrow==5.0.0



Comments:
Thanks for reporting, I'm looking into it :)

Fixed by #2950.