how processing in batch works in datasets 
State: closed
Labels: name: dataset request

Hi,
I need to process my datasets before it is passed to dataloader in batch, 
here is my codes 

```
class AbstractTask(ABC):
    task_name: str = NotImplemented
    preprocessor: Callable = NotImplemented
    split_to_data_split: Mapping[str, str] = NotImplemented
    tokenizer: Callable = NotImplemented
    max_source_length: str = NotImplemented
    max_target_length: str = NotImplemented
    # TODO: should not be a task item, but cannot see other ways.
    tpu_num_cores: int = None

    # The arguments set are for all tasks and needs to be kept common.
    def __init__(self, config):
        self.max_source_length = config['max_source_length']
        self.max_target_length = config['max_target_length']
        self.tokenizer = config['tokenizer']
        self.tpu_num_cores = config['tpu_num_cores']

    def _encode(self, batch) -> Dict[str, torch.Tensor]:
        batch_encoding = self.tokenizer.prepare_seq2seq_batch(
            [x["src_texts"] for x in batch],
            tgt_texts=[x["tgt_texts"] for x in batch],
            max_length=self.max_source_length,
            max_target_length=self.max_target_length,
            padding="max_length" if self.tpu_num_cores is not None else "longest",  # TPU hack
            return_tensors="pt"
        )
        return batch_encoding.data


    def data_split(self, split):
        return self.split_to_data_split[split]

    def get_dataset(self, split, n_obs=None):
        split = self.data_split(split)
        if n_obs is not None:
            split = split+"[:{}]".format(n_obs)
        dataset = load_dataset(self.task_name, split=split)
        dataset = dataset.map(self.preprocessor, remove_columns=dataset.column_names)
        dataset = dataset.map(lambda batch: self._encode(batch), batched=True)
        dataset.set_format(type="torch", columns=['input_ids', 'token_type_ids', 'attention_mask', 'label'])
        return dataset

```

I call it like 

`AutoTask.get(task, train_dataset_config).get_dataset(split="train", n_obs=data_args.n_train) 
`

This gives the following error, to me because the data inside the   dataset = dataset.map(lambda batch: self._encode(batch), batched=True) is not processed in batch, could you tell me how I can process dataset in batch inside my function? thanks 

  File "finetune_multitask_trainer.py", line 192, in main
    if training_args.do_train else None
  File "finetune_multitask_trainer.py", line 191, in <dictcomp>
    split="train", n_obs=data_args.n_train) for task in data_args.task}
  File "/remote/idiap.svm/user.active/rkarimi/dev/internship/seq2seq/tasks.py", line 56, in get_dataset
    dataset = dataset.map(lambda batch: self._encode(batch), batched=True)
  File "/idiap/user/rkarimi/libs/anaconda3/envs/internship/lib/python3.7/site-packages/datasets/arrow_dataset.py", line 1236, in map
    update_data = does_function_return_dict(test_inputs, test_indices)
  File "/idiap/user/rkarimi/libs/anaconda3/envs/internship/lib/python3.7/site-packages/datasets/arrow_dataset.py", line 1207, in does_function_return_dict
    function(*fn_args, indices, **fn_kwargs) if with_indices else function(*fn_args, **fn_kwargs)
  File "/remote/idiap.svm/user.active/rkarimi/dev/internship/seq2seq/tasks.py", line 56, in <lambda>
    dataset = dataset.map(lambda batch: self._encode(batch), batched=True)
  File "/remote/idiap.svm/user.active/rkarimi/dev/internship/seq2seq/tasks.py", line 37, in _encode
    [x["src_texts"] for x in batch],
  File "/remote/idiap.svm/user.active/rkarimi/dev/internship/seq2seq/tasks.py", line 37, in <listcomp>
    [x["src_texts"] for x in batch],
TypeError: string indices must be integers