BentoML icon indicating copy to clipboard operation
BentoML copied to clipboard

bug: Can't save bento with custom transformers pipeline

Open beatanyari opened this issue 11 months ago • 1 comments

Describe the bug

The support for custom transformers pipeline seems to be broken, I believe introduced by this commit When trying to save a custom pipeline, I'm getting this error:

BentoMLException: Argument 'pipeline' is not an instance of <class '__main__.MyPipeline'>. It is an instance of <class '__main__.MyPipeline'>.

To reproduce

from transformers import pipeline
from transformers import Pipeline
from transformers import AutoTokenizer
from transformers import AutoModelForSequenceClassification
from transformers.pipelines import SUPPORTED_TASKS

import typing as t

import bentoml


class MyPipeline(Pipeline):
    def _sanitize_parameters(self, **kwargs):
        preprocess_kwargs = {}
        if "maybe_arg" in kwargs:
            preprocess_kwargs["maybe_arg"] = kwargs["maybe_arg"]
        return preprocess_kwargs, {}, {}

    def preprocess(self, text, maybe_arg=2):
        input_ids = self.tokenizer(text, return_tensors="pt")
        return input_ids

    def _forward(self, model_inputs):
        outputs = self.model(**model_inputs)
        return outputs

    def postprocess(self, model_outputs):
        return model_outputs["logits"].softmax(-1).numpy()


TASK_NAME: str = "my-classification-task"
TASK_DEFINITION: t.Dict[str, t.Any] = {
    "impl": MyPipeline,
    "tf": (),
    "pt": (AutoModelForSequenceClassification,),
    "default": {},
    "type": "text",
}
SUPPORTED_TASKS[TASK_NAME] = TASK_DEFINITION

pipe = pipeline(
    task=TASK_NAME,
    model=AutoModelForSequenceClassification.from_pretrained(
        "distilbert-base-uncased-finetuned-sst-2-english"
    ),
    tokenizer=AutoTokenizer.from_pretrained(
        "distilbert-base-uncased-finetuned-sst-2-english"
    ),
)

saved_pipe = bentoml.transformers.save_model(
    "my_classification_model",
    pipeline=pipe,
    task_name=TASK_NAME,
    task_definition=TASK_DEFINITION,
)

Expected behavior

No response

Environment

bentoml: 1.3.16

beatanyari avatar Dec 16 '24 13:12 beatanyari

Did you register your pipeline class with huggingface?

Can you send a small reproducer here?

aarnphm avatar Dec 17 '24 08:12 aarnphm