Environment info

• transformers version: 4.16.2
• Platform: Linux-5.4.144+-x86_64-with-Ubuntu-18.04-bionic
• Python version: 3.7.12
• PyTorch version (GPU?): 1.10.0+cu111 (False)
• Tensorflow version (GPU?): 2.7.0 (False)
• Flax version (CPU?/GPU?/TPU?): not installed (NA)
• Jax version: not installed
• JaxLib version: not installed
• Using GPU in script?: NA
• Using distributed or parallel set-up in script?: NA

Who can help

• Pipelines: @Narsil
• Documentation: @sgugger

To reproduce

Steps to reproduce the behavior:

Follow the commands in this google colab: https://colab.research.google.com/drive/1BkvrzISZbt_N0S6ntgG7_fN3LdfJE2DH?usp=sharing

Attempting to import KeyDataset

from transformers.pipelines.base import KeyDataset

Produces the following error:

ImportError: cannot import name 'KeyDataset' from 'transformers.pipelines.base'

Expected behavior

Successfully import KeyDataset to perform the example in the documentation for batch processing pipeline predictions (https://huggingface.co/docs/transformers/v4.16.2/en/main_classes/pipelines#pipeline-batching):

from transformers import pipeline
from transformers.pipelines.base import KeyDataset
import datasets

dataset = datasets.load_dataset("imdb", name="plain_text", split="unsupervised")
pipe = pipeline("text-classification", device=0)
for out in pipe(KeyDataset(dataset, "text"), batch_size=8, truncation="only_first"):
    print(out)