BentoML
BentoML copied to clipboard
bug: Can't save bento with custom transformers pipeline
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
Did you register your pipeline class with huggingface?
Can you send a small reproducer here?