Skip to content

Commit

Permalink
[Inference Client] Factorize inference payload build (#2601)
Browse files Browse the repository at this point in the history 
* Factorize inference payload build and add test

* Add comments

* Add method description

* fix style

* fix style again

* fix prepare payload helper

* experiment: try old version of workflow

* revert experiment: try old version of workflow

* Add docstring

* update docstring

* simplify json payload construction when inputs is a dict

* ignore mypy str bytes warning

* fix encoding condition

* remove unnecessary checks for parameters
  • Loading branch information
@hanouticelina
hanouticelina authored Oct 15, 2024
1 parent c9d7865 commit a4bc2e5
Show file tree
Hide file tree
Showing 4 changed files with 306 additions and 277 deletions.
196 changes: 60 additions & 136 deletions src/huggingface_hub/inference/_client.py
View file
Original file line number Diff line number Diff line change
Expand Up @@ -57,6 +57,7 @@
_get_unsupported_text_generation_kwargs,
_import_numpy,
_open_as_binary,
_prepare_payload,
_set_unsupported_text_generation_kwargs,
_stream_chat_completion_response,
_stream_text_generation_response,
Expand Down Expand Up @@ -364,18 +365,8 @@ def audio_classification(
```
"""
parameters = {"function_to_apply": function_to_apply, "top_k": top_k}
if all(parameter is None for parameter in parameters.values()):
# if no parameters are provided, send audio as raw data
data = audio
payload: Optional[Dict[str, Any]] = None
else:
# Or some parameters are provided -> send audio as base64 encoded string
data = None
payload = {"inputs": _b64_encode(audio)}
for key, value in parameters.items():
if value is not None:
payload.setdefault("parameters", {})[key] = value
response = self.post(json=payload, data=data, model=model, task="audio-classification")
payload = _prepare_payload(audio, parameters=parameters, expect_binary=True)
response = self.post(**payload, model=model, task="audio-classification")
return AudioClassificationOutputElement.parse_obj_as_list(response)

def audio_to_audio(
Expand Down Expand Up @@ -988,7 +979,7 @@ def document_question_answering(
[DocumentQuestionAnsweringOutputElement(answer='us-001', end=16, score=0.9999666213989258, start=16, words=None)]
```
"""
payload: Dict[str, Any] = {"question": question, "image": _b64_encode(image)}
inputs: Dict[str, Any] = {"question": question, "image": _b64_encode(image)}
parameters = {
"doc_stride": doc_stride,
"handle_impossible_answer": handle_impossible_answer,
Expand All @@ -999,10 +990,8 @@ def document_question_answering(
"top_k": top_k,
"word_boxes": word_boxes,
}
for key, value in parameters.items():
if value is not None:
payload.setdefault("parameters", {})[key] = value
response = self.post(json=payload, model=model, task="document-question-answering")
payload = _prepare_payload(inputs, parameters=parameters)
response = self.post(**payload, model=model, task="document-question-answering")
return DocumentQuestionAnsweringOutputElement.parse_obj_as_list(response)

def feature_extraction(
Expand Down Expand Up @@ -1060,17 +1049,14 @@ def feature_extraction(
[ 0.28552425, -0.928395 , -1.2077185 , ..., 0.76810825, -2.1069427 , 0.6236161 ]], dtype=float32)
```
"""
payload: Dict = {"inputs": text}
parameters = {
"normalize": normalize,
"prompt_name": prompt_name,
"truncate": truncate,
"truncation_direction": truncation_direction,
}
for key, value in parameters.items():
if value is not None:
payload.setdefault("parameters", {})[key] = value
response = self.post(json=payload, model=model, task="feature-extraction")
payload = _prepare_payload(text, parameters=parameters)
response = self.post(**payload, model=model, task="feature-extraction")
np = _import_numpy()
return np.array(_bytes_to_dict(response), dtype="float32")

Expand Down Expand Up @@ -1119,12 +1105,9 @@ def fill_mask(
]
```
"""
payload: Dict = {"inputs": text}
parameters = {"targets": targets, "top_k": top_k}
for key, value in parameters.items():
if value is not None:
payload.setdefault("parameters", {})[key] = value
response = self.post(json=payload, model=model, task="fill-mask")
payload = _prepare_payload(text, parameters=parameters)
response = self.post(**payload, model=model, task="fill-mask")
return FillMaskOutputElement.parse_obj_as_list(response)

def image_classification(
Expand Down Expand Up @@ -1166,19 +1149,8 @@ def image_classification(
```
"""
parameters = {"function_to_apply": function_to_apply, "top_k": top_k}

if all(parameter is None for parameter in parameters.values()):
data = image
payload: Optional[Dict[str, Any]] = None

else:
data = None
payload = {"inputs": _b64_encode(image)}
for key, value in parameters.items():
if value is not None:
payload.setdefault("parameters", {})[key] = value

response = self.post(json=payload, data=data, model=model, task="image-classification")
payload = _prepare_payload(image, parameters=parameters, expect_binary=True)
response = self.post(**payload, model=model, task="image-classification")
return ImageClassificationOutputElement.parse_obj_as_list(response)

def image_segmentation(
Expand Down Expand Up @@ -1237,18 +1209,8 @@ def image_segmentation(
"subtask": subtask,
"threshold": threshold,
}
if all(parameter is None for parameter in parameters.values()):
# if no parameters are provided, the image can be raw bytes, an image file, or URL to an online image
data = image
payload: Optional[Dict[str, Any]] = None
else:
# if parameters are provided, the image needs to be a base64-encoded string
data = None
payload = {"inputs": _b64_encode(image)}
for key, value in parameters.items():
if value is not None:
payload.setdefault("parameters", {})[key] = value
response = self.post(json=payload, data=data, model=model, task="image-segmentation")
payload = _prepare_payload(image, parameters=parameters, expect_binary=True)
response = self.post(**payload, model=model, task="image-segmentation")
output = ImageSegmentationOutputElement.parse_obj_as_list(response)
for item in output:
item.mask = _b64_to_image(item.mask) # type: ignore [assignment]
Expand Down Expand Up @@ -1323,19 +1285,8 @@ def image_to_image(
"guidance_scale": guidance_scale,
**kwargs,
}
if all(parameter is None for parameter in parameters.values()):
# Either only an image to send => send as raw bytes
data = image
payload: Optional[Dict[str, Any]] = None
else:
# if parameters are provided, the image needs to be a base64-encoded string
data = None
payload = {"inputs": _b64_encode(image)}
for key, value in parameters.items():
if value is not None:
payload.setdefault("parameters", {})[key] = value

response = self.post(json=payload, data=data, model=model, task="image-to-image")
payload = _prepare_payload(image, parameters=parameters, expect_binary=True)
response = self.post(**payload, model=model, task="image-to-image")
return _bytes_to_image(response)

def image_to_text(self, image: ContentT, *, model: Optional[str] = None) -> ImageToTextOutput:
Expand Down Expand Up @@ -1493,25 +1444,15 @@ def object_detection(
```py
>>> from huggingface_hub import InferenceClient
>>> client = InferenceClient()
>>> client.object_detection("people.jpg"):
>>> client.object_detection("people.jpg")
[ObjectDetectionOutputElement(score=0.9486683011054993, label='person', box=ObjectDetectionBoundingBox(xmin=59, ymin=39, xmax=420, ymax=510)), ...]
```
"""
parameters = {
"threshold": threshold,
}
if all(parameter is None for parameter in parameters.values()):
# if no parameters are provided, the image can be raw bytes, an image file, or URL to an online image
data = image
payload: Optional[Dict[str, Any]] = None
else:
# if parameters are provided, the image needs to be a base64-encoded string
data = None
payload = {"inputs": _b64_encode(image)}
for key, value in parameters.items():
if value is not None:
payload.setdefault("parameters", {})[key] = value
response = self.post(json=payload, data=data, model=model, task="object-detection")
payload = _prepare_payload(image, parameters=parameters, expect_binary=True)
response = self.post(**payload, model=model, task="object-detection")
return ObjectDetectionOutputElement.parse_obj_as_list(response)

def question_answering(
Expand Down Expand Up @@ -1587,12 +1528,10 @@ def question_answering(
"max_seq_len": max_seq_len,
"top_k": top_k,
}
payload: Dict[str, Any] = {"question": question, "context": context}
for key, value in parameters.items():
if value is not None:
payload.setdefault("parameters", {})[key] = value
inputs: Dict[str, Any] = {"question": question, "context": context}
payload = _prepare_payload(inputs, parameters=parameters)
response = self.post(
json=payload,
**payload,
model=model,
task="question-answering",
)
Expand Down Expand Up @@ -1700,19 +1639,14 @@ def summarization(
SummarizationOutput(generated_text="The Eiffel tower is one of the most famous landmarks in the world....")
```
"""
payload: Dict[str, Any] = {"inputs": text}
if parameters is not None:
payload["parameters"] = parameters
else:
if parameters is None:
parameters = {
"clean_up_tokenization_spaces": clean_up_tokenization_spaces,
"generate_parameters": generate_parameters,
"truncation": truncation,
}
for key, value in parameters.items():
if value is not None:
payload.setdefault("parameters", {})[key] = value
response = self.post(json=payload, model=model, task="summarization")
payload = _prepare_payload(text, parameters=parameters)
response = self.post(**payload, model=model, task="summarization")
return SummarizationOutput.parse_obj_as_list(response)[0]

def table_question_answering(
Expand Down Expand Up @@ -1757,15 +1691,13 @@ def table_question_answering(
TableQuestionAnsweringOutputElement(answer='36542', coordinates=[[0, 1]], cells=['36542'], aggregator='AVERAGE')
```
"""
payload: Dict[str, Any] = {
inputs = {
"query": query,
"table": table,
}

if parameters is not None:
payload["parameters"] = parameters
payload = _prepare_payload(inputs, parameters=parameters)
response = self.post(
json=payload,
**payload,
model=model,
task="table-question-answering",
)
Expand Down Expand Up @@ -1813,7 +1745,11 @@ def tabular_classification(self, table: Dict[str, Any], *, model: Optional[str]
["5", "5", "5"]
```
"""
response = self.post(json={"table": table}, model=model, task="tabular-classification")
response = self.post(
json={"table": table},
model=model,
task="tabular-classification",
)
return _bytes_to_list(response)

def tabular_regression(self, table: Dict[str, Any], *, model: Optional[str] = None) -> List[float]:
Expand Down Expand Up @@ -1899,15 +1835,16 @@ def text_classification(
]
```
"""
payload: Dict[str, Any] = {"inputs": text}
parameters = {
"function_to_apply": function_to_apply,
"top_k": top_k,
}
for key, value in parameters.items():
if value is not None:
payload.setdefault("parameters", {})[key] = value
response = self.post(json=payload, model=model, task="text-classification")
payload = _prepare_payload(text, parameters=parameters)
response = self.post(
**payload,
model=model,
task="text-classification",
)
return TextClassificationOutputElement.parse_obj_as_list(response)[0] # type: ignore [return-value]

@overload
Expand Down Expand Up @@ -2481,7 +2418,7 @@ def text_to_image(
>>> image.save("better_astronaut.png")
```
"""
payload = {"inputs": prompt}

parameters = {
"negative_prompt": negative_prompt,
"height": height,
Expand All @@ -2493,10 +2430,8 @@ def text_to_image(
"seed": seed,
**kwargs,
}
for key, value in parameters.items():
if value is not None:
payload.setdefault("parameters", {})[key] = value # type: ignore
response = self.post(json=payload, model=model, task="text-to-image")
payload = _prepare_payload(prompt, parameters=parameters)
response = self.post(**payload, model=model, task="text-to-image")
return _bytes_to_image(response)

def text_to_speech(
Expand Down Expand Up @@ -2599,7 +2534,6 @@ def text_to_speech(
>>> Path("hello_world.flac").write_bytes(audio)
```
"""
payload: Dict[str, Any] = {"inputs": text}
parameters = {
"do_sample": do_sample,
"early_stopping": early_stopping,
Expand All @@ -2618,10 +2552,8 @@ def text_to_speech(
"typical_p": typical_p,
"use_cache": use_cache,
}
for key, value in parameters.items():
if value is not None:
payload.setdefault("parameters", {})[key] = value
response = self.post(json=payload, model=model, task="text-to-speech")
payload = _prepare_payload(text, parameters=parameters)
response = self.post(**payload, model=model, task="text-to-speech")
return response

def token_classification(
Expand Down Expand Up @@ -2683,17 +2615,15 @@ def token_classification(
]
```
"""
payload: Dict[str, Any] = {"inputs": text}

parameters = {
"aggregation_strategy": aggregation_strategy,
"ignore_labels": ignore_labels,
"stride": stride,
}
for key, value in parameters.items():
if value is not None:
payload.setdefault("parameters", {})[key] = value
payload = _prepare_payload(text, parameters=parameters)
response = self.post(
json=payload,
**payload,
model=model,
task="token-classification",
)
Expand Down Expand Up @@ -2769,18 +2699,15 @@ def translation(

if src_lang is None and tgt_lang is not None:
raise ValueError("You cannot specify `tgt_lang` without specifying `src_lang`.")
payload: Dict[str, Any] = {"inputs": text}
parameters = {
"src_lang": src_lang,
"tgt_lang": tgt_lang,
"clean_up_tokenization_spaces": clean_up_tokenization_spaces,
"truncation": truncation,
"generate_parameters": generate_parameters,
}
for key, value in parameters.items():
if value is not None:
payload.setdefault("parameters", {})[key] = value
response = self.post(json=payload, model=model, task="translation")
payload = _prepare_payload(text, parameters=parameters)
response = self.post(**payload, model=model, task="translation")
return TranslationOutput.parse_obj_as_list(response)[0]

def visual_question_answering(
Expand Down Expand Up @@ -2921,15 +2848,14 @@ def zero_shot_classification(
```
"""

parameters = {"candidate_labels": labels, "multi_label": multi_label}
if hypothesis_template is not None:
parameters["hypothesis_template"] = hypothesis_template

parameters = {
"candidate_labels": labels,
"multi_label": multi_label,
"hypothesis_template": hypothesis_template,
}
payload = _prepare_payload(text, parameters=parameters)
response = self.post(
json={
"inputs": text,
"parameters": parameters,
},
**payload,
task="zero-shot-classification",
model=model,
)
Expand Down Expand Up @@ -2986,13 +2912,11 @@ def zero_shot_image_classification(
if len(labels) < 2:
raise ValueError("You must specify at least 2 classes to compare.")

payload = {
"inputs": {"image": _b64_encode(image), "candidateLabels": ",".join(labels)},
}
if hypothesis_template is not None:
payload.setdefault("parameters", {})["hypothesis_template"] = hypothesis_template
inputs = {"image": _b64_encode(image), "candidateLabels": ",".join(labels)}
parameters = {"hypothesis_template": hypothesis_template}
payload = _prepare_payload(inputs, parameters=parameters)
response = self.post(
json=payload,
**payload,
model=model,
task="zero-shot-image-classification",
)
Expand Down
Loading

0 comments on commit a4bc2e5

Please sign in to comment.