Skip to content

Commit 52c3de4

Browse files
samrudsjiapinw
authored andcommitted
doc: Add supported task types to schema builder omission (aws#4584)
* Add task types supported in HF_TASK metadata override documentation * Formatting changes * Minor update
1 parent 56ea0eb commit 52c3de4

File tree

1 file changed

+10
-5
lines changed

1 file changed

+10
-5
lines changed

src/sagemaker/serve/builder/model_builder.py

Lines changed: 10 additions & 5 deletions
Original file line numberDiff line numberDiff line change
@@ -143,6 +143,9 @@ class ModelBuilder(Triton, DJL, JumpStart, TGI, Transformers, FastAPIServe):
143143
The schema builder translates the input into bytes and converts the response
144144
into a stream. All translations between the server and the client are handled
145145
automatically with the specified input and output.
146+
The schema builder can be omitted for HuggingFace models with task types TextGeneration,
147+
TextClassification, and QuestionAnswering. Omitting SchemaBuilder is in
148+
beta for FillMask, and AutomaticSpeechRecognition use-cases.
146149
model (Optional[Union[object, str]): Model object (with ``predict`` method to perform
147150
inference) or a HuggingFace/JumpStart Model ID. Either ``model`` or ``inference_spec``
148151
is required for the model builder to build the artifact.
@@ -166,10 +169,11 @@ class ModelBuilder(Triton, DJL, JumpStart, TGI, Transformers, FastAPIServe):
166169
``TORCHSERVE``, ``MMS``, ``TENSORFLOW_SERVING``, ``DJL_SERVING``,
167170
``TRITON``, and``TGI``.
168171
model_metadata (Optional[Dict[str, Any]): Dictionary used to override model metadata.
169-
Currently, ``HF_TASK`` is overridable for HuggingFace model. ``MLFLOW_MODEL_PATH``
170-
is available for providing local path or s3 path to MLflow artifacts. However,
171-
``MLFLOW_MODEL_PATH`` is experimental and is not intended for production use
172-
at this moment.
172+
Currently, ``HF_TASK`` is overridable for HuggingFace model. HF_TASK should be set for
173+
new models without task metadata in the Hub, adding unsupported task types will throw
174+
an exception. ``MLFLOW_MODEL_PATH`` is available for providing local path or s3 path
175+
to MLflow artifacts. However, ``MLFLOW_MODEL_PATH`` is experimental and is not
176+
intended for production use at this moment.
173177
"""
174178

175179
model_path: Optional[str] = field(
@@ -270,7 +274,8 @@ class ModelBuilder(Triton, DJL, JumpStart, TGI, Transformers, FastAPIServe):
270274
default=None,
271275
metadata={
272276
"help": "Define the model metadata to override, currently supports `HF_TASK`, "
273-
"`MLFLOW_MODEL_PATH`"
277+
"`MLFLOW_MODEL_PATH`. HF_TASK should be set for new models without task metadata in "
278+
"the Hub, Adding unsupported task types will throw an exception"
274279
},
275280
)
276281

0 commit comments

Comments
 (0)