Re-export KeyDataset. (#15645)

* Re-export `KeyDataset`.

* Update the docs locations.
This commit is contained in:
Nicolas Patry 2022-02-15 17:49:38 +01:00 committed by GitHub
parent 28e6155d8a
commit cdf19c501d
No known key found for this signature in database
GPG key ID: 4AEE18F83AFDEB23
3 changed files with 7 additions and 4 deletions

View file

@ -78,7 +78,7 @@ GPU. If it doesn't don't hesitate to create an issue.
```python
import datasets
from transformers import pipeline
from transformers.pipelines.base import KeyDataset
from transformers.pipelines.pt_utils import KeyDataset
from tqdm.auto import tqdm
pipe = pipeline("automatic-speech-recognition", model="facebook/wav2vec2-base-960h", device=0)
@ -128,7 +128,7 @@ whenever the pipeline uses its streaming ability (so when passing lists or `Data
```python
from transformers import pipeline
from transformers.pipelines.base import KeyDataset
from transformers.pipelines.pt_utils import KeyDataset
import datasets
dataset = datasets.load_dataset("imdb", name="plain_text", split="unsupervised")

View file

@ -118,7 +118,7 @@ Next, load a dataset (see the 🤗 Datasets [Quick Start](https://huggingface.co
Now you can iterate over the dataset with the pipeline. `KeyDataset` retrieves the item in the dictionary returned by the dataset:
```py
>>> from transformers.pipelines.base import KeyDataset
>>> from transformers.pipelines.pt_utils import KeyDataset
>>> from tqdm.auto import tqdm
>>> for out in tqdm(speech_recognizer(KeyDataset(dataset, "file"))):
@ -316,4 +316,4 @@ One particularly cool 🤗 Transformers feature is the ability to save a model a
>>> tokenizer = AutoTokenizer.from_pretrained(pt_save_directory)
>>> tf_model = TFAutoModelForSequenceClassification.from_pretrained(pt_save_directory, from_pt=True)
```
```

View file

@ -49,6 +49,9 @@ if is_torch_available():
from torch.utils.data import DataLoader, Dataset
from ..models.auto.modeling_auto import AutoModel
# Re-export for backward compatibility
from .pt_utils import KeyDataset
else:
Dataset = None
KeyDataset = None