optimum/intel/pipelines/pipeline_base.py

#  Copyright 2024 The HuggingFace Team. All rights reserved.
#
#  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.

from pathlib import Path
from typing import TYPE_CHECKING, Any, Dict, Optional, Union

import torch
from transformers import (
    AudioClassificationPipeline,
    AutoFeatureExtractor,
    AutoImageProcessor,
    AutomaticSpeechRecognitionPipeline,
    AutoTokenizer,
    FeatureExtractionPipeline,
    FillMaskPipeline,
    ImageClassificationPipeline,
    ImageToTextPipeline,
    Pipeline,
    PreTrainedTokenizer,
    QuestionAnsweringPipeline,
    SummarizationPipeline,
    Text2TextGenerationPipeline,
    TextClassificationPipeline,
    TextGenerationPipeline,
    TokenClassificationPipeline,
    TranslationPipeline,
    ZeroShotClassificationPipeline,
)
from transformers import pipeline as transformers_pipeline
from transformers.feature_extraction_utils import PreTrainedFeatureExtractor
from transformers.utils import logging

from optimum.intel.utils.import_utils import (
    IPEX_IMPORT_ERROR,
    OPENVINO_IMPORT_ERROR,
    is_ipex_available,
    is_openvino_available,
)


if is_ipex_available():
    from ..ipex.modeling_base import (
        IPEXModel,
        IPEXModelForAudioClassification,
        IPEXModelForCausalLM,
        IPEXModelForImageClassification,
        IPEXModelForMaskedLM,
        IPEXModelForQuestionAnswering,
        IPEXModelForSeq2SeqLM,
        IPEXModelForSequenceClassification,
        IPEXModelForTokenClassification,
    )

    IPEX_SUPPORTED_TASKS = {
        "text-generation": {
            "impl": TextGenerationPipeline,
            "class": (IPEXModelForCausalLM,),
            "default": "gpt2",
            "type": "text",
        },
        "summarization": {
            "impl": SummarizationPipeline,
            "class": (IPEXModelForSeq2SeqLM,),
            "default": "t5-base",
            "type": "text",
        },
        "translation": {
            "impl": TranslationPipeline,
            "class": (IPEXModelForSeq2SeqLM,),
            "default": "t5-small",
            "type": "text",
        },
        "text2text-generation": {
            "impl": Text2TextGenerationPipeline,
            "class": (IPEXModelForSeq2SeqLM,),
            "default": "t5-small",
            "type": "text",
        },
        "fill-mask": {
            "impl": FillMaskPipeline,
            "class": (IPEXModelForMaskedLM,),
            "default": "bert-base-cased",
            "type": "text",
        },
        "question-answering": {
            "impl": QuestionAnsweringPipeline,
            "class": (IPEXModelForQuestionAnswering,),
            "default": "distilbert-base-cased-distilled-squad",
            "type": "text",
        },
        "image-classification": {
            "impl": ImageClassificationPipeline,
            "class": (IPEXModelForImageClassification,),
            "default": "google/vit-base-patch16-224",
            "type": "image",
        },
        "text-classification": {
            "impl": TextClassificationPipeline,
            "class": (IPEXModelForSequenceClassification,),
            "default": "distilbert-base-uncased-finetuned-sst-2-english",
            "type": "text",
        },
        "token-classification": {
            "impl": TokenClassificationPipeline,
            "class": (IPEXModelForTokenClassification,),
            "default": "dbmdz/bert-large-cased-finetuned-conll03-english",
            "type": "text",
        },
        "audio-classification": {
            "impl": AudioClassificationPipeline,
            "class": (IPEXModelForAudioClassification,),
            "default": "superb/hubert-base-superb-ks",
            "type": "audio",
        },
    }
else:
    IPEX_SUPPORTED_TASKS = {}


if is_openvino_available():
    from ..openvino import (
        OVModelForAudioClassification,
        OVModelForCausalLM,
        OVModelForFeatureExtraction,
        OVModelForImageClassification,
        OVModelForMaskedLM,
        OVModelForQuestionAnswering,
        OVModelForSeq2SeqLM,
        OVModelForSequenceClassification,
        OVModelForSpeechSeq2Seq,
        OVModelForTokenClassification,
        OVModelForVision2Seq,
    )
    from ..openvino.modeling_base import OVBaseModel

    OPENVINO_SUPPORTED_TASKS = {
        "feature-extraction": {
            "impl": FeatureExtractionPipeline,
            "class": (OVModelForFeatureExtraction,),
            "default": "distilbert-base-cased",
            "type": "text",  # feature extraction is only supported for text at the moment
        },
        "fill-mask": {
            "impl": FillMaskPipeline,
            "class": (OVModelForMaskedLM,),
            "default": "bert-base-cased",
            "type": "text",
        },
        "image-classification": {
            "impl": ImageClassificationPipeline,
            "class": (OVModelForImageClassification,),
            "default": "google/vit-base-patch16-224",
            "type": "image",
        },
        "question-answering": {
            "impl": QuestionAnsweringPipeline,
            "class": (OVModelForQuestionAnswering,),
            "default": "distilbert-base-cased-distilled-squad",
            "type": "text",
        },
        "text-classification": {
            "impl": TextClassificationPipeline,
            "class": (OVModelForSequenceClassification,),
            "default": "distilbert-base-uncased-finetuned-sst-2-english",
            "type": "text",
        },
        "text-generation": {
            "impl": TextGenerationPipeline,
            "class": (OVModelForCausalLM,),
            "default": "distilgpt2",
            "type": "text",
        },
        "token-classification": {
            "impl": TokenClassificationPipeline,
            "class": (OVModelForTokenClassification,),
            "default": "dbmdz/bert-large-cased-finetuned-conll03-english",
            "type": "text",
        },
        "zero-shot-classification": {
            "impl": ZeroShotClassificationPipeline,
            "class": (OVModelForSequenceClassification,),
            "default": "facebook/bart-large-mnli",
            "type": "text",
        },
        "summarization": {
            "impl": SummarizationPipeline,
            "class": (OVModelForSeq2SeqLM,),
            "default": "t5-base",
            "type": "text",
        },
        "translation": {
            "impl": TranslationPipeline,
            "class": (OVModelForSeq2SeqLM,),
            "default": "t5-small",
            "type": "text",
        },
        "text2text-generation": {
            "impl": Text2TextGenerationPipeline,
            "class": (OVModelForSeq2SeqLM,),
            "default": "t5-small",
            "type": "text",
        },
        "automatic-speech-recognition": {
            "impl": AutomaticSpeechRecognitionPipeline,
            "class": (OVModelForSpeechSeq2Seq,),
            "default": "openai/whisper-tiny.en",
            "type": "multimodal",
        },
        "image-to-text": {
            "impl": ImageToTextPipeline,
            "class": (OVModelForVision2Seq,),
            "default": "nlpconnect/vit-gpt2-image-captioning",
            "type": "multimodal",
        },
        "audio-classification": {
            "impl": AudioClassificationPipeline,
            "class": (OVModelForAudioClassification,),
            "default": "superb/hubert-base-superb-ks",
            "type": "audio",
        },
    }
else:
    OPENVINO_SUPPORTED_TASKS = {}


def load_openvino_model(
    model,
    targeted_task,
    SUPPORTED_TASKS,
    hub_kwargs: Optional[Dict[str, Any]] = None,
    model_kwargs: Optional[Dict[str, Any]] = None,
):
    hub_kwargs = hub_kwargs or {}
    model_kwargs = model_kwargs or {}
    ov_model_class = SUPPORTED_TASKS[targeted_task]["class"][0]

    if isinstance(model, str) or model is None:
        model_id = model or SUPPORTED_TASKS[targeted_task]["default"]
        model = ov_model_class.from_pretrained(model_id, **hub_kwargs, **model_kwargs)
    elif isinstance(model, OVBaseModel):
        model_id = model.model_save_dir
    else:
        raise ValueError(
            f"""Model {model} is not supported. Please provide a valid model either as string or ORTModel.
            You can also provide non model then a default one will be used"""
        )
    return model, model_id


def load_ipex_model(
    model,
    targeted_task,
    SUPPORTED_TASKS,
    hub_kwargs: Optional[Dict[str, Any]] = None,
    model_kwargs: Optional[Dict[str, Any]] = None,
    device_map: Optional[torch.device] = None,
):
    hub_kwargs = hub_kwargs or {}
    model_kwargs = model_kwargs or {}
    ipex_model_class = SUPPORTED_TASKS[targeted_task]["class"][0]

    if model is None:
        model_id = SUPPORTED_TASKS[targeted_task]["default"]
        model = ipex_model_class.from_pretrained(model_id, **hub_kwargs, **model_kwargs, device_map=device_map)
    elif isinstance(model, str):
        model_id = model
        model = ipex_model_class.from_pretrained(model, **hub_kwargs, **model_kwargs, device_map=device_map)
    elif isinstance(model, IPEXModel):
        model_id = getattr(model.config, "name_or_path", None)
    else:
        raise ValueError(
            f"""Model {model} is not supported. Please provide a valid model name or path or a IPEXModel.
            You can also provide non model then a default one will be used"""
        )

    return model, model_id


MAPPING_LOADING_FUNC = {
    "ipex": load_ipex_model,
    "openvino": load_openvino_model,
}


if TYPE_CHECKING:
    from transformers.modeling_utils import PreTrainedModel
    from transformers.tokenization_utils_fast import PreTrainedTokenizerFast


logger = logging.get_logger(__name__)


def pipeline(
    task: str = None,
    model: Optional[Union[str, "PreTrainedModel"]] = None,
    tokenizer: Optional[Union[str, PreTrainedTokenizer, "PreTrainedTokenizerFast"]] = None,
    feature_extractor: Optional[Union[str, PreTrainedFeatureExtractor]] = None,
    use_fast: bool = True,
    token: Optional[Union[str, bool]] = None,
    accelerator: Optional[str] = None,
    revision: Optional[str] = None,
    trust_remote_code: Optional[bool] = None,
    torch_dtype: Optional[Union[str, torch.dtype]] = None,
    model_kwargs: Dict[str, Any] = None,
    **kwargs,
) -> Pipeline:
    """
    Utility factory method to build a [`Pipeline`].

    Pipelines are made of:

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

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

            - `"text-generation"`: will return a [`TextGenerationPipeline`]:.
            - `"fill-mask"`: will return a [`FillMaskPipeline`].
            - `"question-answering"`: will return a [`QuestionAnsweringPipeline`].
            - `"image-classificatio"`: will return a [`ImageClassificationPipeline`].
            - `"text-classification"`: will return a [`TextClassificationPipeline`].
            - `"token-classification"`: will return a [`TokenClassificationPipeline`].
            - `"audio-classification"`: will return a [`AudioClassificationPipeline`].

        model (`str` or [`PreTrainedModel`], *optional*):
            The model that will be used by the pipeline to make predictions. This can be a model identifier or an
            actual instance of a pretrained model inheriting from [`PreTrainedModel`] (for PyTorch).

            If not provided, the default for the `task` will be loaded.
        tokenizer (`str` or [`PreTrainedTokenizer`], *optional*):
            The tokenizer that will be used by the pipeline to encode data for the model. This can be a model
            identifier or an actual pretrained tokenizer inheriting from [`PreTrainedTokenizer`].

            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.
        accelerator (`str`, *optional*):
            The optimization backends, choose from ["ipex", "inc", "openvino"].
        use_fast (`bool`, *optional*, defaults to `True`):
            Whether or not to use a Fast tokenizer if possible (a [`PreTrainedTokenizerFast`]).
        torch_dtype (`str` or `torch.dtype`, *optional*):
            Sent directly as `model_kwargs` (just a simpler shortcut) to use the available precision for this model
            (`torch.float16`, `torch.bfloat16`, ... or `"auto"`).
        model_kwargs (`Dict[str, Any]`, *optional*):
            Additional dictionary of keyword arguments passed along to the model's `from_pretrained(...,
            **model_kwargs)` function.

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

    Examples:

    ```python
    >>> import torch
    >>> from optimum.intel.pipelines import pipeline

    >>> pipe = pipeline('text-generation', 'gpt2', torch_dtype=torch.bfloat16)
    >>> pipe("Describe a real-world application of AI in sustainable energy.")
    ```"""

    if model_kwargs is None:
        model_kwargs = {}

    commit_hash = kwargs.pop("_commit_hash", None)

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

    if model is None and tokenizer is not None:
        raise RuntimeError(
            "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."
        )

    if accelerator not in MAPPING_LOADING_FUNC:
        if accelerator is None:
            msg = "Impossible to instantiate a pipeline without specifying an `accelerator`."
        else:
            msg = f"`accelerator` {accelerator} is not supported."

        raise ValueError(msg + f" Supported list of `accelerator` is : {', '.join(MAPPING_LOADING_FUNC)}.")

    if accelerator == "ipex":
        if not is_ipex_available():
            raise RuntimeError(IPEX_IMPORT_ERROR.format("`accelerator=ipex`"))
        supported_tasks = IPEX_SUPPORTED_TASKS

    if accelerator == "openvino":
        if not is_openvino_available():
            raise RuntimeError(OPENVINO_IMPORT_ERROR.format("`accelerator=openvino`"))
        supported_tasks = OPENVINO_SUPPORTED_TASKS

    if task not in supported_tasks:
        raise ValueError(
            f"Task {task} is not supported for the {accelerator} pipelines. Supported tasks are {', '.join(supported_tasks)}"
        )

    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 {"image", "video"}:
            no_tokenizer_tasks.add(_task)
        elif values["type"] in {"audio"}:
            no_tokenizer_tasks.add(_task)
        elif values["type"] not in ["multimodal", "audio", "video"]:
            raise ValueError(f"SUPPORTED_TASK {_task} contains invalid type {values['type']}")

    load_tokenizer = task not in no_tokenizer_tasks
    load_feature_extractor = task not in no_feature_extractor_tasks

    hub_kwargs = {
        "revision": revision,
        "token": token,
        "trust_remote_code": trust_remote_code,
        "_commit_hash": commit_hash,
    }

    if isinstance(model, Path):
        model = str(model)

    tokenizer_kwargs = model_kwargs.copy()
    if torch_dtype is not None:
        if "torch_dtype" in model_kwargs:
            raise ValueError(
                'You cannot use both `pipeline(... torch_dtype=..., model_kwargs={"torch_dtype":...})` as those'
                " arguments might conflict, use only one.)"
            )
        model_kwargs["torch_dtype"] = torch_dtype

    # Load the correct model and convert it to the expected format if needed
    model, model_id = MAPPING_LOADING_FUNC[accelerator](
        model,
        task,
        SUPPORTED_TASKS=supported_tasks,
        hub_kwargs=hub_kwargs,
        model_kwargs=model_kwargs,
        **kwargs,
    )

    if load_tokenizer and tokenizer is None:
        tokenizer = AutoTokenizer.from_pretrained(model_id, use_fast=use_fast, **hub_kwargs, **tokenizer_kwargs)
    if load_feature_extractor and feature_extractor is None:
        try:
            feature_extractor = AutoFeatureExtractor.from_pretrained(model_id, **hub_kwargs, **tokenizer_kwargs)
        except Exception:
            feature_extractor = AutoImageProcessor.from_pretrained(model_id, **hub_kwargs, **tokenizer_kwargs)

    return transformers_pipeline(
        task,
        model=model,
        tokenizer=tokenizer,
        feature_extractor=feature_extractor,
        torch_dtype=torch_dtype,
    )