optimum/graphcore/pipelines/__init__.py

# Copyright 2018 The HuggingFace Inc. team.
# Copyright (c) 2022 Graphcore Ltd. 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 typing import Any, Dict, List, Optional, Union

import poptorch
import torch
import transformers.pipelines
from peft import PeftModel
from transformers import (
    AudioClassificationPipeline,
    AutoFeatureExtractor,
    AutoModelForAudioClassification,
    AutoModelForCausalLM,
    AutoModelForCTC,
    AutoModelForImageClassification,
    AutoModelForMaskedLM,
    AutoModelForQuestionAnswering,
    AutoModelForSeq2SeqLM,
    AutoModelForSequenceClassification,
    AutoModelForSpeechSeq2Seq,
    AutoModelForTokenClassification,
    AutoTokenizer,
    ImageClassificationPipeline,
    Pipeline,
    PreTrainedTokenizer,
    QuestionAnsweringPipeline,
    TextClassificationPipeline,
    TextGenerationPipeline,
    WhisperForConditionalGeneration,
)
from transformers.feature_extraction_utils import PreTrainedFeatureExtractor
from transformers.modeling_utils import PreTrainedModel
from transformers.pipelines import get_task
from transformers.utils import HUGGINGFACE_CO_RESOLVE_ENDPOINT, logging

from optimum.graphcore.generation.utils import MODELS_SUPPORTING_KV_CACHE, IPUGenerationMixin
from optimum.graphcore.ipu_configuration import IncompatibleIPUConfigError, IPUConfig
from optimum.graphcore.modeling_utils import to_pipelined

from .automatic_speech_recognition import IPUAutomaticSpeechRecognitionPipeline
from .fill_mask import IPUFillMaskPipeline
from .text2text_generation import IPUSummarizationPipeline, IPUText2TextGenerationPipeline, IPUTranslationPipeline
from .token_classification import IPUTokenClassificationPipeline
from .zero_shot_classification import IPUZeroShotClassificationPipeline


logger = logging.get_logger(__name__)

TASK_ALIASES = {
    "sentiment-analysis": "text-classification",
    "ner": "token-classification",
}
SUPPORTED_TASKS = {
    "audio-classification": {
        "impl": AudioClassificationPipeline,
        "class": (AutoModelForAudioClassification,),
        "default": {
            "model": ("superb/hubert-base-superb-ks", "d7e0efe"),
            "ipu_config": "Graphcore/hubert-base-ipu",
        },
        "type": "audio",
    },
    "automatic-speech-recognition": {
        "impl": IPUAutomaticSpeechRecognitionPipeline,
        "class": (AutoModelForCTC, AutoModelForSpeechSeq2Seq),
        "default": {
            "model": ("facebook/wav2vec2-base-960h", "55bb623"),
            "ipu_config": "Graphcore/wav2vec2-ctc-base-ipu",
        },
        "type": "multimodal",
    },
    "fill-mask": {
        "impl": IPUFillMaskPipeline,
        "class": (AutoModelForMaskedLM,),
        "default": {
            "model": ("distilroberta-base", "ec58a5b"),
            "ipu_config": "Graphcore/distilroberta-base-ipu",
            "max_length": 128,
        },
        "type": "text",
    },
    "image-classification": {
        "impl": ImageClassificationPipeline,
        "class": (AutoModelForImageClassification,),
        "default": {
            "model": ("google/vit-base-patch16-224", "5dca96d"),
            "ipu_config": "Graphcore/vit-base-ipu",
        },
        "type": "image",
    },
    "question-answering": {
        "impl": QuestionAnsweringPipeline,
        "class": (AutoModelForQuestionAnswering,),
        "default": {
            "model": ("distilbert-base-cased-distilled-squad", "626af31"),
            "ipu_config": "Graphcore/distilbert-base-ipu",
        },
        "type": "text",
    },
    "text-classification": {
        "impl": TextClassificationPipeline,
        "class": (AutoModelForSequenceClassification,),
        "default": {
            "model": ("distilbert-base-uncased-finetuned-sst-2-english", "af0f99b"),
            "ipu_config": "Graphcore/distilbert-base-ipu",
            "max_length": 128,
        },
        "type": "text",
    },
    "token-classification": {
        "impl": IPUTokenClassificationPipeline,
        "class": (AutoModelForTokenClassification,),
        "default": {
            "model": ("dbmdz/bert-large-cased-finetuned-conll03-english", "f2482bf"),
            "ipu_config": "Graphcore/bert-large-ipu",
            "max_length": 128,
        },
        "type": "text",
    },
    "text-generation": {
        "impl": TextGenerationPipeline,
        "class": (AutoModelForCausalLM,),
        "default": {
            "model": ("gpt2", "e7da7f2"),
            "ipu_config": IPUConfig(),
            "max_length": 50,
        },
        "type": "text",
    },
    "summarization": {
        "impl": IPUSummarizationPipeline,
        "class": (AutoModelForSeq2SeqLM,),
        "default": {
            "model": ("ainize/bart-base-cnn", "b90bc9a"),
            "ipu_config": IPUConfig(ipus_per_replica=2),
            "max_input_length": 50,
            "max_length": 20,
            "truncation": "only_first",
        },
        "type": "text",
    },
    # This task is a special case as it's parametrized by SRC, TGT languages.
    "translation": {
        "impl": IPUTranslationPipeline,
        "class": (AutoModelForSeq2SeqLM,),
        "default": {
            "model": ("t5-small", "9507060"),
            "ipu_config": IPUConfig(ipus_per_replica=2),
            "max_length": 50,
            "max_input_length": 45,
            "truncation": "only_first",
        },
        "type": "text",
    },
    "text2text-generation": {
        "impl": IPUText2TextGenerationPipeline,
        "class": (AutoModelForSeq2SeqLM,),
        "default": {
            "model": ("t5-small", "9507060"),
            "ipu_config": IPUConfig(ipus_per_replica=2),
            "max_length": 50,
            "max_input_length": 50,
            "truncation": "only_first",
        },
        "type": "text",
    },
    "zero-shot-classification": {
        "impl": IPUZeroShotClassificationPipeline,
        "class": (AutoModelForSequenceClassification,),
        "default": {
            "model": ("roberta-large-mnli", "130fb28"),
            "ipu_config": "Graphcore/roberta-large-ipu",
            "max_length": 128,
        },
        "type": "text",
    },
}
SUPPORTED_GENERATION_TASKS = {
    "summarization",
    "text-generation",
    "text2text-generation",
    "translation",
}
SUPPORTED_SEQ2SEQ_GENERATION_TASKS = {"summarization", "text2text-generation", "translation"}

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']}")


def list_tasks() -> List[str]:
    """Lists the supported tasks and their aliases"""
    return sorted([*{*SUPPORTED_TASKS, *TASK_ALIASES}])


def get_poplar_executor(
    task: str,
    model: PreTrainedModel,
    ipu_config: Union[IPUConfig, str, dict] = None,
    fp16: bool = True,
    for_generation: bool = False,
) -> PreTrainedModel:
    ipu_config_arg = ipu_config

    if isinstance(ipu_config, str):
        ipu_config = IPUConfig.from_pretrained(ipu_config)
    elif isinstance(ipu_config, dict):
        ipu_config = IPUConfig.from_dict(ipu_config)
    elif not isinstance(ipu_config, IPUConfig):
        raise ValueError("ipu_config must be an IPUConfig, string, or a dictionary.")

    # So that IPUConfig returns inference versions of any parameters
    # that are different in training and inference
    ipu_config.eval()

    parallelize_kwargs = ipu_config.inference_parallelize_kwargs

    ipu_config.inference_device_iterations = 1
    if not parallelize_kwargs.get("use_cond_encoder", False):
        ipu_config.inference_replication_factor = 1
    if not fp16:
        ipu_config.enable_half_partials = False
    try:
        model = to_pipelined(model, ipu_config, force=False)
        if model.config.is_encoder_decoder and isinstance(model, IPUGenerationMixin):
            if "use_cache" not in parallelize_kwargs and model.__class__ in MODELS_SUPPORTING_KV_CACHE:
                parallelize_kwargs["use_cache"] = True
            model.parallelize(for_generation=for_generation, **parallelize_kwargs)
        else:
            model.parallelize(**parallelize_kwargs)
    except Exception as error:
        new_message = (
            "The model and ipu_config seem to be incompatible,"
            " please try a different IPU config or customize it for the model."
            f" The config provided is '{ipu_config_arg}'\n"
            f"{error}"
        )
        raise IncompatibleIPUConfigError(new_message) from error
    if fp16:
        model.half()
    opts = ipu_config.to_options(for_inference=True)
    opts.setExecutionStrategy(poptorch.ShardedExecution())

    # Text generation models have an internal Poplar executor so don't wrap model in that case
    if not for_generation:
        model = poptorch.inferenceModel(model.eval(), opts)
    return model


def check_model_type(self, supported_models: Union[List[str], dict]):
    """
    Check if the model class is supported by the pipeline.

    Args:
        supported_models (`List[str]` or `dict`):
            The list of models supported by the pipeline, or a dictionary with model class values.
    """
    if not isinstance(supported_models, list):  # Create from a model mapping
        supported_models_names = []
        for config, model in supported_models.items():
            # Mapping can now contain tuples of models for the same configuration.
            if isinstance(model, tuple):
                supported_models_names.extend([_model.__name__ for _model in model])
            else:
                supported_models_names.append(model.__name__)
        supported_models = supported_models_names

    if isinstance(self.model, poptorch.PoplarExecutor):
        model_class_name = self.model._user_model.__class__.__bases__[0].__name__
    elif isinstance(self.model, IPUGenerationMixin):
        model_class_name = self.model.__class__.__bases__[0].__name__
    else:
        model_class_name = self.model.__class__.__name__

    if model_class_name not in supported_models:
        logger.error(
            f"The model '{model_class_name}' is not supported for {self.task}. Supported models are"
            f" {supported_models}."
        )


def pipeline(
    task: str = None,
    model: Optional[Any] = None,
    ipu_config: Union[IPUConfig, str, dict] = None,
    tokenizer: Optional[Union[str, PreTrainedTokenizer]] = None,
    feature_extractor: Optional[Union[str, PreTrainedFeatureExtractor]] = None,
    revision: Optional[str] = None,
    use_auth_token: Optional[Union[str, bool]] = None,
    pipeline_class: Optional[Any] = None,
    fp16: bool = True,
    parallelize_kwargs: Optional[Dict[str, Any]] = None,
    **kwargs,
) -> Pipeline:
    """Utility factory method to build a [ Pipeline ] for IPU models.

    Arguments:
        task : The task, see docs for ``transformers.pipeline`` for supported options.
        model : A pre-trained model, see docs for ``transformers.pipeline`` for supported options.
        ipu_config : An IPU config, can either be the path to a model from the HuggingFace Hub
            which defines a ``ipu_config.json`` or a dictionary with the same options.
        tokenizer : The tokenizer, see docs for ``transformers.pipeline`` for supported options.
        feature_extractor : The feature extractor, see docs for ``transformers.pipeline`` for supported options.
        revision : Revision of the model.
        use_auth_token : An authorization token to use for calls to the Hub.
        pipeline_class : Override the `Pipeline` class defined by the task.
        fp16 : If `True`, uses float16.

        **kwargs: Additional keyword arguments that are passed to the ``transformers.pipeline`` function

    Returns:
        The pipeline object for the specified task.
    """

    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 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)

    if task in TASK_ALIASES:
        task = TASK_ALIASES[task]

    targeted_task = "translation" if task.startswith("translation") else task

    if targeted_task not in SUPPORTED_TASKS:
        raise ValueError(f"Task {targeted_task} is not supported. Supported tasks are {list(SUPPORTED_TASKS.keys())}")

    # These will never require a tokenizer.
    # 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 = targeted_task not in NO_TOKENIZER_TASKS
    load_feature_extractor = targeted_task not in NO_FEATURE_EXTRACTOR_TASKS

    if model is None:
        model_id, revision = SUPPORTED_TASKS[targeted_task]["default"]["model"]
        logger.warning(
            f"No model was supplied, defaulted to {model_id} and revision"
            f" {revision} ({HUGGINGFACE_CO_RESOLVE_ENDPOINT}/{model_id}).\n"
            "Using a pipeline without specifying a model name and revision in production is not recommended."
        )
        model = SUPPORTED_TASKS[targeted_task]["class"][0].from_pretrained(model_id, revision=revision)
    elif isinstance(model, str):
        model_id = model
        for cl in SUPPORTED_TASKS[targeted_task]["class"]:
            try:
                model = cl.from_pretrained(model_id, revision=revision)
                break
            except ValueError:
                continue
    elif isinstance(model, PeftModel):
        raise TypeError(
            "Instead of providing `model` as an instance of `PeftModel`, please call `merge_and_unload()` if LoRA "
            "or equivalent to obtain the original `PreTrainedModel` back with adapter weights merged in."
        )
    elif isinstance(model, PreTrainedModel):
        if tokenizer is None and load_tokenizer:
            raise ValueError("If you pass a model as a PreTrainedModel, you must pass a tokenizer as well")
        if feature_extractor is None and load_feature_extractor:
            raise ValueError("If you pass a model as a PreTrainedModel, you must pass a feature extractor as well")

    for_generation = targeted_task in SUPPORTED_GENERATION_TASKS
    if isinstance(model, PreTrainedModel):
        if ipu_config is None:
            ipu_config = SUPPORTED_TASKS[targeted_task]["default"]["ipu_config"]

        parallelize_kwargs = parallelize_kwargs or {}
        # Task of automatic speech recognition is a bit of an edge case where it separates into CTC (not generation) and seq2seq (generation).
        # This check will do for now.
        for_generation |= isinstance(model, WhisperForConditionalGeneration)
        model = get_poplar_executor(
            targeted_task, model, ipu_config=ipu_config, fp16=fp16, for_generation=for_generation, **parallelize_kwargs
        )
    elif isinstance(model, poptorch._poplar_executor.PoplarExecutor):
        if tokenizer is None and load_tokenizer:
            raise ValueError(
                "If you pass a model as a poptorch._poplar_executor.PoplarExecutor, you must pass a tokenizer as well"
            )
        if feature_extractor is None and load_feature_extractor:
            raise ValueError(
                "If you pass a model as a poptorch._poplar_executor.PoplarExecutor, you must pass a feature extractor as well"
            )
    else:
        raise ValueError(
            f"""Model {model} is not supported. Please provide a valid model either as string, PreTrainedModel or
            poptorch._poplar_executor.PoplarExecutor. If you don't provide a model, a default model will be used."""
        )

    # Upstream pipeline creation does not easily support loading these when an actual model
    # is provided, so we load them here.
    if tokenizer is None and load_tokenizer:
        tokenizer = AutoTokenizer.from_pretrained(model_id)

    if feature_extractor is None and load_feature_extractor:
        feature_extractor = AutoFeatureExtractor.from_pretrained(model_id)

    # Override Pipeline methods
    Pipeline.check_model_type = check_model_type

    if pipeline_class is None:
        pipeline_class = SUPPORTED_TASKS[targeted_task]["impl"]

    # Override pipelines' _forward
    old_forward = pipeline_class._forward

    def new_forward(self, model_inputs, *args, **kwargs):
        if isinstance(self.model, poptorch.PoplarExecutor) and not for_generation:
            # For non-text generation models, support batch size changes.
            poplar_executor = self.model
            if poplar_executor._executable_inputs:
                for arg in poplar_executor._executable_inputs.args:
                    if isinstance(arg, torch.Tensor):
                        compiled_bs = arg.shape[0]
                        break
                for input in model_inputs.values():
                    if isinstance(input, torch.Tensor):
                        input_bs = input.shape[0]
                        break
                if compiled_bs != input_bs:
                    poplar_executor.destroy()

        if isinstance(self.model, poptorch.PoplarExecutor) or isinstance(self.model, IPUGenerationMixin):
            if fp16:
                # Support fp16
                for key, input in model_inputs.items():
                    if isinstance(input, torch.Tensor) and input.dtype == torch.float32:
                        model_inputs[key] = input.half()
        return old_forward(self, model_inputs, *args, **kwargs)

    pipeline_class._forward = new_forward

    # Implement pipelines __del__ to clean up poplar exector
    def _del(self):
        # For text generation models, deallocate the internal poplar executors
        if hasattr(self.model, "poptorch_decoder"):
            self.model.poptorch_decoder.destroy()
        if hasattr(self.model, "poptorch_encoder"):
            self.model.poptorch_encoder.destroy()

    pipeline_class.__del__ = _del

    # Auto padding for some tasks
    if "max_length" in SUPPORTED_TASKS[targeted_task]["default"]:
        default_max_length = SUPPORTED_TASKS[targeted_task]["default"]["max_length"]
        if not for_generation:
            kwargs["padding"] = kwargs.get("padding", "max_length")
            if kwargs.get("max_length") is None:
                logger.warning(
                    f"No padding arguments specified, so padding to {default_max_length} by default. "
                    f"Inputs longer than {default_max_length} will be truncated."
                    " To change this behaviour, pass the `padding='max_length'` and"
                    "`max_length=<your desired input length>` arguments to the pipeline function."
                )
        kwargs["max_length"] = kwargs.get("max_length", default_max_length)

    if targeted_task in SUPPORTED_SEQ2SEQ_GENERATION_TASKS:
        default_max_input_length = SUPPORTED_TASKS[targeted_task]["default"]["max_input_length"]
        kwargs["max_input_length"] = kwargs.get("max_input_length", default_max_input_length)
        default_truncation = SUPPORTED_TASKS[targeted_task]["default"]["truncation"]
        kwargs["truncation"] = kwargs.get("truncation", default_truncation)

    # question-answering already has its own default padding length `max_seq_len` defined, so we just enable padding to max length.
    if targeted_task in {"question-answering"}:
        kwargs["padding"] = kwargs.get("padding", "max_length")
        logger.warning(
            "No padding arguments specified, so padding to 384 by default. Inputs longer than 384 will be truncated."
        )

    # Set pad_token for models that do not have pad_token
    if model.config.model_type in {"gpt2"}:
        tokenizer.pad_token = tokenizer.eos_token
        model.config.pad_token_id = model.config.eos_token_id

    return transformers.pipelines.pipeline(
        task,
        model=model,
        tokenizer=tokenizer,
        feature_extractor=feature_extractor,
        use_auth_token=use_auth_token,
        pipeline_class=pipeline_class,
        **kwargs,
    )