fix: add task aliases of supported tasks

bentoml · Dec 9, 2022 · 6aa0442 · 6aa0442
1 parent d7254c0
commit 6aa0442
Showing 1 changed file with 2 additions and 1 deletion.
diff --git a/src/bentoml/_internal/frameworks/transformers.py b/src/bentoml/_internal/frameworks/transformers.py
@@ -191,6 +191,7 @@ def load_model(
             f"Model {bento_model.tag} was saved with module {bento_model.info.module}, not loading with {MODULE_NAME}."
         )
 
+    from transformers.pipelines import TASK_ALIASES
     from transformers.pipelines import SUPPORTED_TASKS
 
     if TYPE_CHECKING:
@@ -199,7 +200,7 @@ def load_model(
         options = bento_model.info.options
 
     task: str = bento_model.info.options.task  # type: ignore
-    if task not in SUPPORTED_TASKS:
+    if task not in SUPPORTED_TASKS and task not in TASK_ALIASES:
         try:
             import cloudpickle  # type: ignore
         except ImportError:  # pragma: no cover