__init__.py 30.5 KB
Newer Older
1
2
3
4
# flake8: noqa
# There's no way to ignore "F401 '...' imported but unused" warnings in this
# module, but to preserve other warnings. So, don't check this module at all.

5
6
import io
import json
7
import os
8

9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
# coding=utf-8
# Copyright 2018 The HuggingFace Inc. team.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#     http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
import warnings
24
from typing import TYPE_CHECKING, Any, Dict, List, Optional, Tuple, Union
25
26

from ..configuration_utils import PretrainedConfig
27
28
29
30
from ..feature_extraction_utils import PreTrainedFeatureExtractor
from ..models.auto.configuration_auto import AutoConfig
from ..models.auto.feature_extraction_auto import FEATURE_EXTRACTOR_MAPPING, AutoFeatureExtractor
from ..models.auto.tokenization_auto import TOKENIZER_MAPPING, AutoTokenizer
31
from ..tokenization_utils import PreTrainedTokenizer
32
from ..tokenization_utils_fast import PreTrainedTokenizerFast
33
from ..utils import HUGGINGFACE_CO_RESOLVE_ENDPOINT, http_get, is_tf_available, is_torch_available, logging
34
from .audio_classification import AudioClassificationPipeline
35
from .automatic_speech_recognition import AutomaticSpeechRecognitionPipeline
36
37
38
39
40
41
42
43
from .base import (
    ArgumentHandler,
    CsvPipelineDataFormat,
    JsonPipelineDataFormat,
    PipedPipelineDataFormat,
    Pipeline,
    PipelineDataFormat,
    PipelineException,
44
    PipelineRegistry,
45
    get_default_model_and_revision,
46
    infer_framework_load_model,
47
48
49
50
)
from .conversational import Conversation, ConversationalPipeline
from .feature_extraction import FeatureExtractionPipeline
from .fill_mask import FillMaskPipeline
51
from .image_classification import ImageClassificationPipeline
52
from .image_segmentation import ImageSegmentationPipeline
53
from .object_detection import ObjectDetectionPipeline
54
55
56
57
58
from .question_answering import QuestionAnsweringArgumentHandler, QuestionAnsweringPipeline
from .table_question_answering import TableQuestionAnsweringArgumentHandler, TableQuestionAnsweringPipeline
from .text2text_generation import SummarizationPipeline, Text2TextGenerationPipeline, TranslationPipeline
from .text_classification import TextClassificationPipeline
from .text_generation import TextGenerationPipeline
59
60
61
62
63
64
from .token_classification import (
    AggregationStrategy,
    NerPipeline,
    TokenClassificationArgumentHandler,
    TokenClassificationPipeline,
)
65
from .visual_question_answering import VisualQuestionAnsweringPipeline
66
from .zero_shot_classification import ZeroShotClassificationArgumentHandler, ZeroShotClassificationPipeline
67
from .zero_shot_image_classification import ZeroShotImageClassificationPipeline
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84


if is_tf_available():
    import tensorflow as tf

    from ..models.auto.modeling_tf_auto import (
        TF_MODEL_FOR_QUESTION_ANSWERING_MAPPING,
        TF_MODEL_FOR_SEQ_TO_SEQ_CAUSAL_LM_MAPPING,
        TF_MODEL_FOR_SEQUENCE_CLASSIFICATION_MAPPING,
        TF_MODEL_FOR_TOKEN_CLASSIFICATION_MAPPING,
        TF_MODEL_WITH_LM_HEAD_MAPPING,
        TFAutoModel,
        TFAutoModelForCausalLM,
        TFAutoModelForMaskedLM,
        TFAutoModelForQuestionAnswering,
        TFAutoModelForSeq2SeqLM,
        TFAutoModelForSequenceClassification,
Kamal Raj's avatar
Kamal Raj committed
85
        TFAutoModelForTableQuestionAnswering,
86
87
88
89
90
91
92
93
94
95
96
97
98
        TFAutoModelForTokenClassification,
    )

if is_torch_available():
    import torch

    from ..models.auto.modeling_auto import (
        MODEL_FOR_MASKED_LM_MAPPING,
        MODEL_FOR_QUESTION_ANSWERING_MAPPING,
        MODEL_FOR_SEQ_TO_SEQ_CAUSAL_LM_MAPPING,
        MODEL_FOR_SEQUENCE_CLASSIFICATION_MAPPING,
        MODEL_FOR_TABLE_QUESTION_ANSWERING_MAPPING,
        MODEL_FOR_TOKEN_CLASSIFICATION_MAPPING,
99
        MODEL_FOR_VISUAL_QUESTION_ANSWERING_MAPPING,
100
        AutoModel,
101
        AutoModelForAudioClassification,
102
        AutoModelForCausalLM,
103
        AutoModelForCTC,
104
        AutoModelForImageClassification,
105
        AutoModelForImageSegmentation,
106
        AutoModelForMaskedLM,
107
        AutoModelForObjectDetection,
108
        AutoModelForQuestionAnswering,
109
        AutoModelForSemanticSegmentation,
110
111
        AutoModelForSeq2SeqLM,
        AutoModelForSequenceClassification,
112
        AutoModelForSpeechSeq2Seq,
113
114
        AutoModelForTableQuestionAnswering,
        AutoModelForTokenClassification,
115
        AutoModelForVisualQuestionAnswering,
116
117
118
119
120
121
122
123
124
    )
if TYPE_CHECKING:
    from ..modeling_tf_utils import TFPreTrainedModel
    from ..modeling_utils import PreTrainedModel

logger = logging.get_logger(__name__)


# Register all the supported tasks here
125
126
127
TASK_ALIASES = {
    "sentiment-analysis": "text-classification",
    "ner": "token-classification",
128
    "vqa": "visual-question-answering",
129
}
130
SUPPORTED_TASKS = {
131
132
133
134
    "audio-classification": {
        "impl": AudioClassificationPipeline,
        "tf": (),
        "pt": (AutoModelForAudioClassification,) if is_torch_available() else (),
135
        "default": {"model": {"pt": ("superb/wav2vec2-base-superb-ks", "372e048")}},
136
        "type": "audio",
137
    },
138
139
140
    "automatic-speech-recognition": {
        "impl": AutomaticSpeechRecognitionPipeline,
        "tf": (),
141
        "pt": (AutoModelForCTC, AutoModelForSpeechSeq2Seq) if is_torch_available() else (),
142
        "default": {"model": {"pt": ("facebook/wav2vec2-base-960h", "55bb623")}},
143
        "type": "multimodal",
144
    },
145
146
    "feature-extraction": {
        "impl": FeatureExtractionPipeline,
147
148
        "tf": (TFAutoModel,) if is_tf_available() else (),
        "pt": (AutoModel,) if is_torch_available() else (),
149
        "default": {"model": {"pt": ("distilbert-base-cased", "935ac13"), "tf": ("distilbert-base-cased", "935ac13")}},
150
        "type": "multimodal",
151
    },
152
    "text-classification": {
153
        "impl": TextClassificationPipeline,
154
155
        "tf": (TFAutoModelForSequenceClassification,) if is_tf_available() else (),
        "pt": (AutoModelForSequenceClassification,) if is_torch_available() else (),
156
157
        "default": {
            "model": {
158
159
                "pt": ("distilbert-base-uncased-finetuned-sst-2-english", "af0f99b"),
                "tf": ("distilbert-base-uncased-finetuned-sst-2-english", "af0f99b"),
160
161
            },
        },
162
        "type": "text",
163
    },
164
    "token-classification": {
165
        "impl": TokenClassificationPipeline,
166
167
        "tf": (TFAutoModelForTokenClassification,) if is_tf_available() else (),
        "pt": (AutoModelForTokenClassification,) if is_torch_available() else (),
168
169
        "default": {
            "model": {
170
171
                "pt": ("dbmdz/bert-large-cased-finetuned-conll03-english", "f2482bf"),
                "tf": ("dbmdz/bert-large-cased-finetuned-conll03-english", "f2482bf"),
172
173
            },
        },
174
        "type": "text",
175
176
177
    },
    "question-answering": {
        "impl": QuestionAnsweringPipeline,
178
179
        "tf": (TFAutoModelForQuestionAnswering,) if is_tf_available() else (),
        "pt": (AutoModelForQuestionAnswering,) if is_torch_available() else (),
180
        "default": {
181
182
183
184
            "model": {
                "pt": ("distilbert-base-cased-distilled-squad", "626af31"),
                "tf": ("distilbert-base-cased-distilled-squad", "626af31"),
            },
185
        },
186
        "type": "text",
187
188
189
    },
    "table-question-answering": {
        "impl": TableQuestionAnsweringPipeline,
190
        "pt": (AutoModelForTableQuestionAnswering,) if is_torch_available() else (),
Kamal Raj's avatar
Kamal Raj committed
191
        "tf": (TFAutoModelForTableQuestionAnswering,) if is_tf_available() else (),
192
193
        "default": {
            "model": {
194
195
                "pt": ("google/tapas-base-finetuned-wtq", "69ceee2"),
                "tf": ("google/tapas-base-finetuned-wtq", "69ceee2"),
196
197
            },
        },
198
        "type": "text",
199
    },
200
201
202
203
204
    "visual-question-answering": {
        "impl": VisualQuestionAnsweringPipeline,
        "pt": (AutoModelForVisualQuestionAnswering,) if is_torch_available() else (),
        "tf": (),
        "default": {
205
            "model": {"pt": ("dandelin/vilt-b32-finetuned-vqa", "4355f59")},
206
207
208
        },
        "type": "multimodal",
    },
209
210
    "fill-mask": {
        "impl": FillMaskPipeline,
211
212
        "tf": (TFAutoModelForMaskedLM,) if is_tf_available() else (),
        "pt": (AutoModelForMaskedLM,) if is_torch_available() else (),
213
        "default": {"model": {"pt": ("distilroberta-base", "ec58a5b"), "tf": ("distilroberta-base", "ec58a5b")}},
214
        "type": "text",
215
216
217
    },
    "summarization": {
        "impl": SummarizationPipeline,
218
219
        "tf": (TFAutoModelForSeq2SeqLM,) if is_tf_available() else (),
        "pt": (AutoModelForSeq2SeqLM,) if is_torch_available() else (),
220
        "default": {"model": {"pt": ("sshleifer/distilbart-cnn-12-6", "a4f8f3e"), "tf": ("t5-small", "d769bba")}},
221
        "type": "text",
222
223
224
225
    },
    # This task is a special case as it's parametrized by SRC, TGT languages.
    "translation": {
        "impl": TranslationPipeline,
226
227
        "tf": (TFAutoModelForSeq2SeqLM,) if is_tf_available() else (),
        "pt": (AutoModelForSeq2SeqLM,) if is_torch_available() else (),
228
        "default": {
229
230
231
            ("en", "fr"): {"model": {"pt": ("t5-base", "686f1db"), "tf": ("t5-base", "686f1db")}},
            ("en", "de"): {"model": {"pt": ("t5-base", "686f1db"), "tf": ("t5-base", "686f1db")}},
            ("en", "ro"): {"model": {"pt": ("t5-base", "686f1db"), "tf": ("t5-base", "686f1db")}},
232
        },
233
        "type": "text",
234
235
236
    },
    "text2text-generation": {
        "impl": Text2TextGenerationPipeline,
237
238
        "tf": (TFAutoModelForSeq2SeqLM,) if is_tf_available() else (),
        "pt": (AutoModelForSeq2SeqLM,) if is_torch_available() else (),
239
        "default": {"model": {"pt": ("t5-base", "686f1db"), "tf": ("t5-base", "686f1db")}},
240
        "type": "text",
241
242
243
    },
    "text-generation": {
        "impl": TextGenerationPipeline,
244
245
        "tf": (TFAutoModelForCausalLM,) if is_tf_available() else (),
        "pt": (AutoModelForCausalLM,) if is_torch_available() else (),
246
        "default": {"model": {"pt": ("gpt2", "6c0e608"), "tf": ("gpt2", "6c0e608")}},
247
        "type": "text",
248
249
250
    },
    "zero-shot-classification": {
        "impl": ZeroShotClassificationPipeline,
251
252
        "tf": (TFAutoModelForSequenceClassification,) if is_tf_available() else (),
        "pt": (AutoModelForSequenceClassification,) if is_torch_available() else (),
253
        "default": {
254
255
            "model": {"pt": ("facebook/bart-large-mnli", "c626438"), "tf": ("roberta-large-mnli", "130fb28")},
            "config": {"pt": ("facebook/bart-large-mnli", "c626438"), "tf": ("roberta-large-mnli", "130fb28")},
256
        },
257
        "type": "text",
258
    },
259
260
261
262
    "zero-shot-image-classification": {
        "impl": ZeroShotImageClassificationPipeline,
        "tf": (TFAutoModel,) if is_tf_available() else (),
        "pt": (AutoModel,) if is_torch_available() else (),
263
264
265
266
267
268
        "default": {
            "model": {
                "pt": ("openai/clip-vit-base-patch32", "f4881ba"),
                "tf": ("openai/clip-vit-base-patch32", "f4881ba"),
            }
        },
269
270
        "type": "multimodal",
    },
271
272
    "conversational": {
        "impl": ConversationalPipeline,
273
274
        "tf": (TFAutoModelForSeq2SeqLM, TFAutoModelForCausalLM) if is_tf_available() else (),
        "pt": (AutoModelForSeq2SeqLM, AutoModelForCausalLM) if is_torch_available() else (),
275
276
277
        "default": {
            "model": {"pt": ("microsoft/DialoGPT-medium", "8bada3b"), "tf": ("microsoft/DialoGPT-medium", "8bada3b")}
        },
278
        "type": "text",
279
    },
280
281
    "image-classification": {
        "impl": ImageClassificationPipeline,
282
283
        "tf": (),
        "pt": (AutoModelForImageClassification,) if is_torch_available() else (),
284
        "default": {"model": {"pt": ("google/vit-base-patch16-224", "5dca96d")}},
285
        "type": "image",
286
    },
287
288
289
    "image-segmentation": {
        "impl": ImageSegmentationPipeline,
        "tf": (),
290
        "pt": (AutoModelForImageSegmentation, AutoModelForSemanticSegmentation) if is_torch_available() else (),
291
        "default": {"model": {"pt": ("facebook/detr-resnet-50-panoptic", "fc15262")}},
292
        "type": "image",
293
    },
294
295
296
297
    "object-detection": {
        "impl": ObjectDetectionPipeline,
        "tf": (),
        "pt": (AutoModelForObjectDetection,) if is_torch_available() else (),
298
        "default": {"model": {"pt": ("facebook/detr-resnet-50", "2729413")}},
299
        "type": "image",
300
    },
301
302
}

303
304
305
306
307
308
309
310
311
312
NO_FEATURE_EXTRACTOR_TASKS = set()
NO_TOKENIZER_TASKS = set()
for task, values in SUPPORTED_TASKS.items():
    if values["type"] == "text":
        NO_FEATURE_EXTRACTOR_TASKS.add(task)
    elif values["type"] in {"audio", "image"}:
        NO_TOKENIZER_TASKS.add(task)
    elif values["type"] != "multimodal":
        raise ValueError(f"SUPPORTED_TASK {task} contains invalid type {values['type']}")

313
314
PIPELINE_REGISTRY = PipelineRegistry(supported_tasks=SUPPORTED_TASKS, task_aliases=TASK_ALIASES)

315

316
317
318
319
def get_supported_tasks() -> List[str]:
    """
    Returns a list of supported task strings.
    """
320
    return PIPELINE_REGISTRY.get_supported_tasks()
321
322


323
324
325
326
327
328
329
330
331
332
333
334
335
336
337
338
339
340
341
342
343
344
345
def get_task(model: str, use_auth_token: Optional[str] = None) -> str:
    tmp = io.BytesIO()
    headers = {}
    if use_auth_token:
        headers["Authorization"] = f"Bearer {use_auth_token}"

    try:
        http_get(f"https://huggingface.co/api/models/{model}", tmp, headers=headers)
        tmp.seek(0)
        body = tmp.read()
        data = json.loads(body)
    except Exception as e:
        raise RuntimeError(f"Instantiating a pipeline without a task set raised an error: {e}")
    if "pipeline_tag" not in data:
        raise RuntimeError(
            f"The model {model} does not seem to have a correct `pipeline_tag` set to infer the task automatically"
        )
    if data.get("library_name", "transformers") != "transformers":
        raise RuntimeError(f"This model is meant to be used with {data['library_name']} not with transformers")
    task = data["pipeline_tag"]
    return task


346
347
348
349
350
351
def check_task(task: str) -> Tuple[Dict, Any]:
    """
    Checks an incoming task string, to validate it's correct and return the default Pipeline and Model classes, and
    default models if they exist.

    Args:
352
        task (`str`):
353
354
            The task defining which pipeline will be returned. Currently accepted tasks are:

355
356
357
358
359
360
361
362
363
364
365
366
367
368
369
370
            - `"audio-classification"`
            - `"automatic-speech-recognition"`
            - `"conversational"`
            - `"feature-extraction"`
            - `"fill-mask"`
            - `"image-classification"`
            - `"question-answering"`
            - `"table-question-answering"`
            - `"text2text-generation"`
            - `"text-classification"` (alias `"sentiment-analysis"` available)
            - `"text-generation"`
            - `"token-classification"` (alias `"ner"` available)
            - `"translation"`
            - `"translation_xx_to_yy"`
            - `"summarization"`
            - `"zero-shot-classification"`
371
            - `"zero-shot-image-classification"`
372
373

    Returns:
Sylvain Gugger's avatar
Sylvain Gugger committed
374
375
        (task_defaults`dict`, task_options: (`tuple`, None)) The actual dictionary required to initialize the pipeline
        and some extra task options for parametrized tasks like "translation_XX_to_YY"
376
377
378


    """
379
    return PIPELINE_REGISTRY.check_task(task)
380
381
382


def pipeline(
383
    task: str = None,
384
385
    model: Optional = None,
    config: Optional[Union[str, PretrainedConfig]] = None,
386
    tokenizer: Optional[Union[str, PreTrainedTokenizer, PreTrainedTokenizerFast]] = None,
387
    feature_extractor: Optional[Union[str, PreTrainedFeatureExtractor]] = None,
388
389
390
    framework: Optional[str] = None,
    revision: Optional[str] = None,
    use_fast: bool = True,
391
    use_auth_token: Optional[Union[str, bool]] = None,
392
    model_kwargs: Dict[str, Any] = None,
393
    pipeline_class: Optional[Any] = None,
394
395
396
    **kwargs
) -> Pipeline:
    """
397
    Utility factory method to build a [`Pipeline`].
398
399
400

    Pipelines are made of:

401
402
        - A [tokenizer](tokenizer) in charge of mapping raw textual input to token.
        - A [model](model) to make predictions from the inputs.
403
404
405
        - Some (optional) post processing for enhancing model's output.

    Args:
406
        task (`str`):
407
408
            The task defining which pipeline will be returned. Currently accepted tasks are:

409
            - `"audio-classification"`: will return a [`AudioClassificationPipeline`].
Sylvain Gugger's avatar
Sylvain Gugger committed
410
            - `"automatic-speech-recognition"`: will return a [`AutomaticSpeechRecognitionPipeline`].
411
412
413
414
415
416
417
418
419
420
            - `"conversational"`: will return a [`ConversationalPipeline`].
            - `"feature-extraction"`: will return a [`FeatureExtractionPipeline`].
            - `"fill-mask"`: will return a [`FillMaskPipeline`]:.
            - `"image-classification"`: will return a [`ImageClassificationPipeline`].
            - `"question-answering"`: will return a [`QuestionAnsweringPipeline`].
            - `"table-question-answering"`: will return a [`TableQuestionAnsweringPipeline`].
            - `"text2text-generation"`: will return a [`Text2TextGenerationPipeline`].
            - `"text-classification"` (alias `"sentiment-analysis"` available): will return a
              [`TextClassificationPipeline`].
            - `"text-generation"`: will return a [`TextGenerationPipeline`]:.
Sylvain Gugger's avatar
Sylvain Gugger committed
421
            - `"token-classification"` (alias `"ner"` available): will return a [`TokenClassificationPipeline`].
422
423
424
425
426
427
            - `"translation"`: will return a [`TranslationPipeline`].
            - `"translation_xx_to_yy"`: will return a [`TranslationPipeline`].
            - `"summarization"`: will return a [`SummarizationPipeline`].
            - `"zero-shot-classification"`: will return a [`ZeroShotClassificationPipeline`].

        model (`str` or [`PreTrainedModel`] or [`TFPreTrainedModel`], *optional*):
428
            The model that will be used by the pipeline to make predictions. This can be a model identifier or an
Sylvain Gugger's avatar
Sylvain Gugger committed
429
430
            actual instance of a pretrained model inheriting from [`PreTrainedModel`] (for PyTorch) or
            [`TFPreTrainedModel`] (for TensorFlow).
431

432
433
            If not provided, the default for the `task` will be loaded.
        config (`str` or [`PretrainedConfig`], *optional*):
434
            The configuration that will be used by the pipeline to instantiate the model. This can be a model
Sylvain Gugger's avatar
Sylvain Gugger committed
435
            identifier or an actual pretrained model configuration inheriting from [`PretrainedConfig`].
436
437

            If not provided, the default configuration file for the requested model will be used. That means that if
Sylvain Gugger's avatar
Sylvain Gugger committed
438
439
            `model` is given, its default configuration will be used. However, if `model` is not supplied, this
            `task`'s default model's config is used instead.
440
        tokenizer (`str` or [`PreTrainedTokenizer`], *optional*):
441
            The tokenizer that will be used by the pipeline to encode data for the model. This can be a model
442
            identifier or an actual pretrained tokenizer inheriting from [`PreTrainedTokenizer`].
443

Sylvain Gugger's avatar
Sylvain Gugger committed
444
445
446
447
            If not provided, the default tokenizer for the given `model` will be loaded (if it is a string). If `model`
            is not specified or not a string, then the default tokenizer for `config` is loaded (if it is a string).
            However, if `config` is also not given or not a string, then the default tokenizer for the given `task`
            will be loaded.
448
        feature_extractor (`str` or [`PreTrainedFeatureExtractor`], *optional*):
449
            The feature extractor that will be used by the pipeline to encode data for the model. This can be a model
Sylvain Gugger's avatar
Sylvain Gugger committed
450
            identifier or an actual pretrained feature extractor inheriting from [`PreTrainedFeatureExtractor`].
451
452
453
454

            Feature extractors are used for non-NLP models, such as Speech or Vision models as well as multi-modal
            models. Multi-modal models will also require a tokenizer to be passed.

Sylvain Gugger's avatar
Sylvain Gugger committed
455
456
457
458
            If not provided, the default feature extractor for the given `model` will be loaded (if it is a string). If
            `model` is not specified or not a string, then the default feature extractor for `config` is loaded (if it
            is a string). However, if `config` is also not given or not a string, then the default feature extractor
            for the given `task` will be loaded.
459
        framework (`str`, *optional*):
Sylvain Gugger's avatar
Sylvain Gugger committed
460
461
            The framework to use, either `"pt"` for PyTorch or `"tf"` for TensorFlow. The specified framework must be
            installed.
462
463

            If no framework is specified, will default to the one currently installed. If no framework is specified and
Sylvain Gugger's avatar
Sylvain Gugger committed
464
465
            both frameworks are installed, will default to the framework of the `model`, or to PyTorch if no model is
            provided.
Stas Bekman's avatar
Stas Bekman committed
466
        revision (`str`, *optional*, defaults to `"main"`):
467
468
            When passing a task name or a string model identifier: The specific model version to use. It can be a
            branch name, a tag name, or a commit id, since we use a git-based system for storing models and other
469
470
471
472
            artifacts on huggingface.co, so `revision` can be any identifier allowed by git.
        use_fast (`bool`, *optional*, defaults to `True`):
            Whether or not to use a Fast tokenizer if possible (a [`PreTrainedTokenizerFast`]).
        use_auth_token (`str` or *bool*, *optional*):
Sylvain Gugger's avatar
Sylvain Gugger committed
473
            The token to use as HTTP bearer authorization for remote files. If `True`, will use the token generated
Stas Bekman's avatar
Stas Bekman committed
474
            when running `transformers-cli login` (stored in `~/.huggingface`).
Funtowicz Morgan's avatar
Funtowicz Morgan committed
475
        model_kwargs:
Sylvain Gugger's avatar
Sylvain Gugger committed
476
477
            Additional dictionary of keyword arguments passed along to the model's `from_pretrained(...,
            **model_kwargs)` function.
478
479
480
481
482
        kwargs:
            Additional keyword arguments passed along to the specific pipeline init (see the documentation for the
            corresponding pipeline class for possible values).

    Returns:
483
        [`Pipeline`]: A suitable pipeline for the task.
484

485
    Examples:
486

487
488
    ```python
    >>> from transformers import pipeline, AutoModelForTokenClassification, AutoTokenizer
489

490
    >>> # Sentiment analysis pipeline
Sylvain Gugger's avatar
Sylvain Gugger committed
491
    >>> pipeline("sentiment-analysis")
492

493
    >>> # Question answering pipeline, specifying the checkpoint identifier
Sylvain Gugger's avatar
Sylvain Gugger committed
494
    >>> pipeline("question-answering", model="distilbert-base-cased-distilled-squad", tokenizer="bert-base-cased")
495

496
497
498
    >>> # Named entity recognition pipeline, passing in a specific model and tokenizer
    >>> model = AutoModelForTokenClassification.from_pretrained("dbmdz/bert-large-cased-finetuned-conll03-english")
    >>> tokenizer = AutoTokenizer.from_pretrained("bert-base-cased")
Sylvain Gugger's avatar
Sylvain Gugger committed
499
    >>> pipeline("ner", model=model, tokenizer=tokenizer)
500
    ```"""
501
502
    if model_kwargs is None:
        model_kwargs = {}
503
504
505

    if task is None and model is None:
        raise RuntimeError(
506
507
            "Impossible to instantiate a pipeline without either a task or a model "
            "being specified. "
508
509
510
            "Please provide a task class or a model"
        )

511
512
    if model is None and tokenizer is not None:
        raise RuntimeError(
Sylvain Gugger's avatar
Sylvain Gugger committed
513
514
515
            "Impossible to instantiate a pipeline with tokenizer specified but not the model as the provided tokenizer"
            " may not be compatible with the default model. Please provide a PreTrainedModel class or a"
            " path/identifier to a pretrained model when providing tokenizer."
516
517
518
        )
    if model is None and feature_extractor is not None:
        raise RuntimeError(
Sylvain Gugger's avatar
Sylvain Gugger committed
519
520
521
            "Impossible to instantiate a pipeline with feature_extractor specified but not the model as the provided"
            " feature_extractor may not be compatible with the default model. Please provide a PreTrainedModel class"
            " or a path/identifier to a pretrained model when providing feature_extractor."
522
        )
523

524
525
526
527
528
529
530
531
    if task is None and model is not None:
        if not isinstance(model, str):
            raise RuntimeError(
                "Inferring the task automatically requires to check the hub with a model_id defined as a `str`."
                f"{model} is not a valid model_id."
            )
        task = get_task(model, use_auth_token)

532
533
    # Retrieve the task
    targeted_task, task_options = check_task(task)
534
535
    if pipeline_class is None:
        pipeline_class = targeted_task["impl"]
536
537
538
539

    # Use default model/config/tokenizer for the task if no model is provided
    if model is None:
        # At that point framework might still be undetermined
540
541
542
543
544
545
546
        model, default_revision = get_default_model_and_revision(targeted_task, framework, task_options)
        revision = revision if revision is not None else default_revision
        logger.warning(
            f"No model was supplied, defaulted to {model} and revision"
            f" {revision} ({HUGGINGFACE_CO_RESOLVE_ENDPOINT}/{model}).\n"
            "Using a pipeline without specifying a model name and revision in production is not recommended."
        )
547

548
549
550
    # Retrieve use_auth_token and add it to model_kwargs to be used in .from_pretrained
    model_kwargs["use_auth_token"] = model_kwargs.get("use_auth_token", use_auth_token)

551
552
553
554
555
556
    # Config is the primordial information item.
    # Instantiate config if needed
    if isinstance(config, str):
        config = AutoConfig.from_pretrained(config, revision=revision, _from_pipeline=task, **model_kwargs)
    elif config is None and isinstance(model, str):
        config = AutoConfig.from_pretrained(model, revision=revision, _from_pipeline=task, **model_kwargs)
557

558
    model_name = model if isinstance(model, str) else None
559

560
561
562
563
564
    # Infer the framework from the model
    # Forced if framework already defined, inferred if it's None
    # Will load the correct model if possible
    model_classes = {"tf": targeted_task["tf"], "pt": targeted_task["pt"]}
    framework, model = infer_framework_load_model(
565
566
567
568
569
570
571
        model,
        model_classes=model_classes,
        config=config,
        framework=framework,
        revision=revision,
        task=task,
        **model_kwargs,
572
    )
573

574
575
    model_config = model.config

576
577
    load_tokenizer = type(model_config) in TOKENIZER_MAPPING or model_config.tokenizer_class is not None
    load_feature_extractor = type(model_config) in FEATURE_EXTRACTOR_MAPPING or feature_extractor is not None
578

579
    if task in NO_TOKENIZER_TASKS:
580
        # These will never require a tokenizer.
581
582
583
584
        # the model on the other hand might have a tokenizer, but
        # the files could be missing from the hub, instead of failing
        # on such repos, we just force to not load it.
        load_tokenizer = False
585
586
    if task in NO_FEATURE_EXTRACTOR_TASKS:
        load_feature_extractor = False
587

588
589
590
591
592
593
594
595
596
597
598
599
600
601
602
603
604
605
606
607
608
609
610
611
612
613
614
615
616
617
618
619
620
621
622
623
624
625
626
627
628
629
630
631
632
633
634
635
636
637
    if load_tokenizer:
        # Try to infer tokenizer from model or config name (if provided as str)
        if tokenizer is None:
            if isinstance(model_name, str):
                tokenizer = model_name
            elif isinstance(config, str):
                tokenizer = config
            else:
                # Impossible to guess what is the right tokenizer here
                raise Exception(
                    "Impossible to guess which tokenizer to use. "
                    "Please provide a PreTrainedTokenizer class or a path/identifier to a pretrained tokenizer."
                )

        # Instantiate tokenizer if needed
        if isinstance(tokenizer, (str, tuple)):
            if isinstance(tokenizer, tuple):
                # For tuple we have (tokenizer name, {kwargs})
                use_fast = tokenizer[1].pop("use_fast", use_fast)
                tokenizer_identifier = tokenizer[0]
                tokenizer_kwargs = tokenizer[1]
            else:
                tokenizer_identifier = tokenizer
                tokenizer_kwargs = model_kwargs

            tokenizer = AutoTokenizer.from_pretrained(
                tokenizer_identifier, revision=revision, use_fast=use_fast, _from_pipeline=task, **tokenizer_kwargs
            )

    if load_feature_extractor:
        # Try to infer feature extractor from model or config name (if provided as str)
        if feature_extractor is None:
            if isinstance(model_name, str):
                feature_extractor = model_name
            elif isinstance(config, str):
                feature_extractor = config
            else:
                # Impossible to guess what is the right feature_extractor here
                raise Exception(
                    "Impossible to guess which feature extractor to use. "
                    "Please provide a PreTrainedFeatureExtractor class or a path/identifier "
                    "to a pretrained feature extractor."
                )

        # Instantiate feature_extractor if needed
        if isinstance(feature_extractor, (str, tuple)):
            feature_extractor = AutoFeatureExtractor.from_pretrained(
                feature_extractor, revision=revision, _from_pipeline=task, **model_kwargs
            )

Nicolas Patry's avatar
Nicolas Patry committed
638
639
640
641
642
643
            if (
                feature_extractor._processor_class
                and feature_extractor._processor_class.endswith("WithLM")
                and isinstance(model_name, str)
            ):
                try:
644
                    import kenlm  # to trigger `ImportError` if not installed
Nicolas Patry's avatar
Nicolas Patry committed
645
646
                    from pyctcdecode import BeamSearchDecoderCTC

647
648
649
650
651
652
653
654
655
                    if os.path.isdir(model_name) or os.path.isfile(model_name):
                        decoder = BeamSearchDecoderCTC.load_from_dir(model_name)
                    else:
                        language_model_glob = os.path.join(
                            BeamSearchDecoderCTC._LANGUAGE_MODEL_SERIALIZED_DIRECTORY, "*"
                        )
                        alphabet_filename = BeamSearchDecoderCTC._ALPHABET_SERIALIZED_FILENAME
                        allow_regex = [language_model_glob, alphabet_filename]
                        decoder = BeamSearchDecoderCTC.load_from_hf_hub(model_name, allow_regex=allow_regex)
Nicolas Patry's avatar
Nicolas Patry committed
656
657

                    kwargs["decoder"] = decoder
658
                except ImportError as e:
Nicolas Patry's avatar
Nicolas Patry committed
659
                    logger.warning(
Sylvain Gugger's avatar
Sylvain Gugger committed
660
661
662
                        f"Could not load the `decoder` for {model_name}. Defaulting to raw CTC. Try to install"
                        " `pyctcdecode` and `kenlm`: (`pip install pyctcdecode`, `pip install"
                        f" https://github.com/kpu/kenlm/archive/master.zip`): Error: {e}"
Nicolas Patry's avatar
Nicolas Patry committed
663
664
                    )

665
666
667
668
669
670
671
672
673
    if task == "translation" and model.config.task_specific_params:
        for key in model.config.task_specific_params:
            if key.startswith("translation"):
                task = key
                warnings.warn(
                    f'"translation" task was used, instead of "translation_XX_to_YY", defaulting to "{task}"',
                    UserWarning,
                )
                break
674

675
676
677
678
679
680
    if tokenizer is not None:
        kwargs["tokenizer"] = tokenizer

    if feature_extractor is not None:
        kwargs["feature_extractor"] = feature_extractor

681
    return pipeline_class(model=model, framework=framework, task=task, **kwargs)