in src/huggingface_inference_toolkit/handler.py [0:0]
def __call__(self, data: Dict[str, Any]) -> Dict[str, Any]:
"""
Handles an inference request with input data and makes a prediction.
Args:
:data: (obj): the raw request body data.
:return: prediction output
"""
inputs = data.pop("inputs", data)
parameters = data.pop("parameters", {})
# diffusers and sentence transformers pipelines do not have the `task` arg
if not hasattr(self.pipeline, "task"):
# sentence transformers parameters not supported yet
if any(isinstance(self.pipeline, v) for v in SENTENCE_TRANSFORMERS_TASKS.values()):
return ( # type: ignore
self.pipeline(**inputs) if isinstance(inputs, dict) else self.pipeline(inputs)
)
# diffusers does support kwargs
return ( # type: ignore
self.pipeline(**inputs, **parameters)
if isinstance(inputs, dict)
else self.pipeline(inputs, **parameters)
)
if self.pipeline.task == "question-answering":
if not isinstance(inputs, dict):
raise ValueError(f"inputs must be a dict, but a `{type(inputs)}` was provided instead.")
if not all(k in inputs for k in {"question", "context"}):
raise ValueError(
f"{self.pipeline.task} expects `inputs` to be a dict containing both `question` and "
"`context` as the keys, both of them being either a `str` or a `List[str]`."
)
if self.pipeline.task == "table-question-answering":
if not isinstance(inputs, dict):
raise ValueError(f"inputs must be a dict, but a `{type(inputs)}` was provided instead.")
if "question" in inputs:
inputs["query"] = inputs.pop("question")
if not all(k in inputs for k in {"table", "query"}):
raise ValueError(
f"{self.pipeline.task} expects `inputs` to be a dict containing the keys `table` and "
"either `question` or `query`."
)
if self.pipeline.task.__contains__("translation") or self.pipeline.task in {
"text-generation",
"image-to-text",
"automatic-speech-recognition",
"text-to-audio",
"text-to-speech",
}:
# `generate_kwargs` needs to be a dict, `generation_parameters` is here for forward compatibility
if "generation_parameters" in parameters:
parameters["generate_kwargs"] = parameters.pop("generation_parameters")
if self.pipeline.task.__contains__("translation") or self.pipeline.task in {"text-generation"}:
# flatten the values of `generate_kwargs` as it's not supported as is, but via top-level parameters
generate_kwargs = parameters.pop("generate_kwargs", {})
for key, value in generate_kwargs.items():
parameters[key] = value
if self.pipeline.task.__contains__("zero-shot-classification"):
if "candidateLabels" in parameters:
parameters["candidate_labels"] = parameters.pop("candidateLabels")
if not isinstance(inputs, dict):
inputs = {"sequences": inputs}
if "text" in inputs:
inputs["sequences"] = inputs.pop("text")
if not all(k in inputs for k in {"sequences"}) or not all(k in parameters for k in {"candidate_labels"}):
raise ValueError(
f"{self.pipeline.task} expects `inputs` to be either a string or a dict containing the "
"key `text` or `sequences`, and `parameters` to be a dict containing either `candidate_labels` "
"or `candidateLabels`."
)
return (
self.pipeline(**inputs, **parameters) if isinstance(inputs, dict) else self.pipeline(inputs, **parameters) # type: ignore
)