483 lines
16 KiB
Python
483 lines
16 KiB
Python
|
# coding=utf-8
|
||
|
|
# Copyright 2023-present, the HuggingFace Inc. team.
|
||
|
|
#
|
||
|
|
# Licensed under the Apache License, Version 2.0 (the "License");
|
||
|
|
# you may not use this file except in compliance with the License.
|
||
|
|
# You may obtain a copy of the License at
|
||
|
|
#
|
||
|
|
# http://www.apache.org/licenses/LICENSE-2.0
|
||
|
|
#
|
||
|
|
# Unless required by applicable law or agreed to in writing, software
|
||
|
|
# distributed under the License is distributed on an "AS IS" BASIS,
|
||
|
|
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
|
||
|
|
# See the License for the specific language governing permissions and
|
||
|
|
# limitations under the License.
|
||
|
|
"""Contains utilities used by both the sync and async inference clients."""
|
||
|
|
|
||
|
|
import base64
|
||
|
|
import io
|
||
|
|
import json
|
||
|
|
import logging
|
||
|
|
import time
|
||
|
|
from contextlib import contextmanager
|
||
|
|
from dataclasses import dataclass
|
||
|
|
from pathlib import Path
|
||
|
|
from typing import (
|
||
|
|
TYPE_CHECKING,
|
||
|
|
Any,
|
||
|
|
AsyncIterable,
|
||
|
|
BinaryIO,
|
||
|
|
ContextManager,
|
||
|
|
Dict,
|
||
|
|
Generator,
|
||
|
|
Iterable,
|
||
|
|
List,
|
||
|
|
Literal,
|
||
|
|
NoReturn,
|
||
|
|
Optional,
|
||
|
|
Set,
|
||
|
|
Union,
|
||
|
|
overload,
|
||
|
|
)
|
||
|
|
|
||
|
|
from requests import HTTPError
|
||
|
|
|
||
|
|
from huggingface_hub.errors import (
|
||
|
|
GenerationError,
|
||
|
|
IncompleteGenerationError,
|
||
|
|
OverloadedError,
|
||
|
|
TextGenerationError,
|
||
|
|
UnknownError,
|
||
|
|
ValidationError,
|
||
|
|
)
|
||
|
|
|
||
|
|
from ..constants import ENDPOINT
|
||
|
|
from ..utils import (
|
||
|
|
build_hf_headers,
|
||
|
|
get_session,
|
||
|
|
hf_raise_for_status,
|
||
|
|
is_aiohttp_available,
|
||
|
|
is_numpy_available,
|
||
|
|
is_pillow_available,
|
||
|
|
)
|
||
|
|
from ._generated.types import (
|
||
|
|
ChatCompletionStreamOutput,
|
||
|
|
ChatCompletionStreamOutputChoice,
|
||
|
|
ChatCompletionStreamOutputDelta,
|
||
|
|
TextGenerationStreamOutput,
|
||
|
|
)
|
||
|
|
|
||
|
|
|
||
|
|
if TYPE_CHECKING:
|
||
|
|
from aiohttp import ClientResponse, ClientSession
|
||
|
|
from PIL import Image
|
||
|
|
|
||
|
|
# TYPES
|
||
|
|
UrlT = str
|
||
|
|
PathT = Union[str, Path]
|
||
|
|
BinaryT = Union[bytes, BinaryIO]
|
||
|
|
ContentT = Union[BinaryT, PathT, UrlT]
|
||
|
|
|
||
|
|
# Use to set a Accept: image/png header
|
||
|
|
TASKS_EXPECTING_IMAGES = {"text-to-image", "image-to-image"}
|
||
|
|
|
||
|
|
logger = logging.getLogger(__name__)
|
||
|
|
|
||
|
|
|
||
|
|
# Add dataclass for ModelStatus. We use this dataclass in get_model_status function.
|
||
|
|
@dataclass
|
||
|
|
class ModelStatus:
|
||
|
|
"""
|
||
|
|
This Dataclass represents the the model status in the Hugging Face Inference API.
|
||
|
|
|
||
|
|
Args:
|
||
|
|
loaded (`bool`):
|
||
|
|
If the model is currently loaded into Hugging Face's InferenceAPI. Models
|
||
|
|
are loaded on-demand, leading to the user's first request taking longer.
|
||
|
|
If a model is loaded, you can be assured that it is in a healthy state.
|
||
|
|
state (`str`):
|
||
|
|
The current state of the model. This can be 'Loaded', 'Loadable', 'TooBig'.
|
||
|
|
If a model's state is 'Loadable', it's not too big and has a supported
|
||
|
|
backend. Loadable models are automatically loaded when the user first
|
||
|
|
requests inference on the endpoint. This means it is transparent for the
|
||
|
|
user to load a model, except that the first call takes longer to complete.
|
||
|
|
compute_type (`Dict`):
|
||
|
|
Information about the compute resource the model is using or will use, such as 'gpu' type and number of
|
||
|
|
replicas.
|
||
|
|
framework (`str`):
|
||
|
|
The name of the framework that the model was built with, such as 'transformers'
|
||
|
|
or 'text-generation-inference'.
|
||
|
|
"""
|
||
|
|
|
||
|
|
loaded: bool
|
||
|
|
state: str
|
||
|
|
compute_type: Dict
|
||
|
|
framework: str
|
||
|
|
|
||
|
|
|
||
|
|
## IMPORT UTILS
|
||
|
|
|
||
|
|
|
||
|
|
def _import_aiohttp():
|
||
|
|
# Make sure `aiohttp` is installed on the machine.
|
||
|
|
if not is_aiohttp_available():
|
||
|
|
raise ImportError("Please install aiohttp to use `AsyncInferenceClient` (`pip install aiohttp`).")
|
||
|
|
import aiohttp
|
||
|
|
|
||
|
|
return aiohttp
|
||
|
|
|
||
|
|
|
||
|
|
def _import_numpy():
|
||
|
|
"""Make sure `numpy` is installed on the machine."""
|
||
|
|
if not is_numpy_available():
|
||
|
|
raise ImportError("Please install numpy to use deal with embeddings (`pip install numpy`).")
|
||
|
|
import numpy
|
||
|
|
|
||
|
|
return numpy
|
||
|
|
|
||
|
|
|
||
|
|
def _import_pil_image():
|
||
|
|
"""Make sure `PIL` is installed on the machine."""
|
||
|
|
if not is_pillow_available():
|
||
|
|
raise ImportError(
|
||
|
|
"Please install Pillow to use deal with images (`pip install Pillow`). If you don't want the image to be"
|
||
|
|
" post-processed, use `client.post(...)` and get the raw response from the server."
|
||
|
|
)
|
||
|
|
from PIL import Image
|
||
|
|
|
||
|
|
return Image
|
||
|
|
|
||
|
|
|
||
|
|
## RECOMMENDED MODELS
|
||
|
|
|
||
|
|
# Will be globally fetched only once (see '_fetch_recommended_models')
|
||
|
|
_RECOMMENDED_MODELS: Optional[Dict[str, Optional[str]]] = None
|
||
|
|
|
||
|
|
|
||
|
|
def _fetch_recommended_models() -> Dict[str, Optional[str]]:
|
||
|
|
global _RECOMMENDED_MODELS
|
||
|
|
if _RECOMMENDED_MODELS is None:
|
||
|
|
response = get_session().get(f"{ENDPOINT}/api/tasks", headers=build_hf_headers())
|
||
|
|
hf_raise_for_status(response)
|
||
|
|
_RECOMMENDED_MODELS = {
|
||
|
|
task: _first_or_none(details["widgetModels"]) for task, details in response.json().items()
|
||
|
|
}
|
||
|
|
return _RECOMMENDED_MODELS
|
||
|
|
|
||
|
|
|
||
|
|
def _first_or_none(items: List[Any]) -> Optional[Any]:
|
||
|
|
try:
|
||
|
|
return items[0] or None
|
||
|
|
except IndexError:
|
||
|
|
return None
|
||
|
|
|
||
|
|
|
||
|
|
## ENCODING / DECODING UTILS
|
||
|
|
|
||
|
|
|
||
|
|
@overload
|
||
|
|
def _open_as_binary(
|
||
|
|
content: ContentT,
|
||
|
|
) -> ContextManager[BinaryT]: ... # means "if input is not None, output is not None"
|
||
|
|
|
||
|
|
|
||
|
|
@overload
|
||
|
|
def _open_as_binary(
|
||
|
|
content: Literal[None],
|
||
|
|
) -> ContextManager[Literal[None]]: ... # means "if input is None, output is None"
|
||
|
|
|
||
|
|
|
||
|
|
@contextmanager # type: ignore
|
||
|
|
def _open_as_binary(content: Optional[ContentT]) -> Generator[Optional[BinaryT], None, None]:
|
||
|
|
"""Open `content` as a binary file, either from a URL, a local path, or raw bytes.
|
||
|
|
|
||
|
|
Do nothing if `content` is None,
|
||
|
|
|
||
|
|
TODO: handle a PIL.Image as input
|
||
|
|
TODO: handle base64 as input
|
||
|
|
"""
|
||
|
|
# If content is a string => must be either a URL or a path
|
||
|
|
if isinstance(content, str):
|
||
|
|
if content.startswith("https://") or content.startswith("http://"):
|
||
|
|
logger.debug(f"Downloading content from {content}")
|
||
|
|
yield get_session().get(content).content # TODO: retrieve as stream and pipe to post request ?
|
||
|
|
return
|
||
|
|
content = Path(content)
|
||
|
|
if not content.exists():
|
||
|
|
raise FileNotFoundError(
|
||
|
|
f"File not found at {content}. If `data` is a string, it must either be a URL or a path to a local"
|
||
|
|
" file. To pass raw content, please encode it as bytes first."
|
||
|
|
)
|
||
|
|
|
||
|
|
# If content is a Path => open it
|
||
|
|
if isinstance(content, Path):
|
||
|
|
logger.debug(f"Opening content from {content}")
|
||
|
|
with content.open("rb") as f:
|
||
|
|
yield f
|
||
|
|
else:
|
||
|
|
# Otherwise: already a file-like object or None
|
||
|
|
yield content
|
||
|
|
|
||
|
|
|
||
|
|
def _b64_encode(content: ContentT) -> str:
|
||
|
|
"""Encode a raw file (image, audio) into base64. Can be byes, an opened file, a path or a URL."""
|
||
|
|
with _open_as_binary(content) as data:
|
||
|
|
data_as_bytes = data if isinstance(data, bytes) else data.read()
|
||
|
|
return base64.b64encode(data_as_bytes).decode()
|
||
|
|
|
||
|
|
|
||
|
|
def _b64_to_image(encoded_image: str) -> "Image":
|
||
|
|
"""Parse a base64-encoded string into a PIL Image."""
|
||
|
|
Image = _import_pil_image()
|
||
|
|
return Image.open(io.BytesIO(base64.b64decode(encoded_image)))
|
||
|
|
|
||
|
|
|
||
|
|
def _bytes_to_list(content: bytes) -> List:
|
||
|
|
"""Parse bytes from a Response object into a Python list.
|
||
|
|
|
||
|
|
Expects the response body to be JSON-encoded data.
|
||
|
|
|
||
|
|
NOTE: This is exactly the same implementation as `_bytes_to_dict` and will not complain if the returned data is a
|
||
|
|
dictionary. The only advantage of having both is to help the user (and mypy) understand what kind of data to expect.
|
||
|
|
"""
|
||
|
|
return json.loads(content.decode())
|
||
|
|
|
||
|
|
|
||
|
|
def _bytes_to_dict(content: bytes) -> Dict:
|
||
|
|
"""Parse bytes from a Response object into a Python dictionary.
|
||
|
|
|
||
|
|
Expects the response body to be JSON-encoded data.
|
||
|
|
|
||
|
|
NOTE: This is exactly the same implementation as `_bytes_to_list` and will not complain if the returned data is a
|
||
|
|
list. The only advantage of having both is to help the user (and mypy) understand what kind of data to expect.
|
||
|
|
"""
|
||
|
|
return json.loads(content.decode())
|
||
|
|
|
||
|
|
|
||
|
|
def _bytes_to_image(content: bytes) -> "Image":
|
||
|
|
"""Parse bytes from a Response object into a PIL Image.
|
||
|
|
|
||
|
|
Expects the response body to be raw bytes. To deal with b64 encoded images, use `_b64_to_image` instead.
|
||
|
|
"""
|
||
|
|
Image = _import_pil_image()
|
||
|
|
return Image.open(io.BytesIO(content))
|
||
|
|
|
||
|
|
|
||
|
|
## STREAMING UTILS
|
||
|
|
|
||
|
|
|
||
|
|
def _stream_text_generation_response(
|
||
|
|
bytes_output_as_lines: Iterable[bytes], details: bool
|
||
|
|
) -> Union[Iterable[str], Iterable[TextGenerationStreamOutput]]:
|
||
|
|
"""Used in `InferenceClient.text_generation`."""
|
||
|
|
# Parse ServerSentEvents
|
||
|
|
for byte_payload in bytes_output_as_lines:
|
||
|
|
output = _format_text_generation_stream_output(byte_payload, details)
|
||
|
|
if output is not None:
|
||
|
|
yield output
|
||
|
|
|
||
|
|
|
||
|
|
async def _async_stream_text_generation_response(
|
||
|
|
bytes_output_as_lines: AsyncIterable[bytes], details: bool
|
||
|
|
) -> Union[AsyncIterable[str], AsyncIterable[TextGenerationStreamOutput]]:
|
||
|
|
"""Used in `AsyncInferenceClient.text_generation`."""
|
||
|
|
# Parse ServerSentEvents
|
||
|
|
async for byte_payload in bytes_output_as_lines:
|
||
|
|
output = _format_text_generation_stream_output(byte_payload, details)
|
||
|
|
if output is not None:
|
||
|
|
yield output
|
||
|
|
|
||
|
|
|
||
|
|
def _format_text_generation_stream_output(
|
||
|
|
byte_payload: bytes, details: bool
|
||
|
|
) -> Optional[Union[str, TextGenerationStreamOutput]]:
|
||
|
|
if not byte_payload.startswith(b"data:"):
|
||
|
|
return None # empty line
|
||
|
|
|
||
|
|
# Decode payload
|
||
|
|
payload = byte_payload.decode("utf-8")
|
||
|
|
json_payload = json.loads(payload.lstrip("data:").rstrip("/n"))
|
||
|
|
|
||
|
|
# Either an error as being returned
|
||
|
|
if json_payload.get("error") is not None:
|
||
|
|
raise _parse_text_generation_error(json_payload["error"], json_payload.get("error_type"))
|
||
|
|
|
||
|
|
# Or parse token payload
|
||
|
|
output = TextGenerationStreamOutput.parse_obj_as_instance(json_payload)
|
||
|
|
return output.token.text if not details else output
|
||
|
|
|
||
|
|
|
||
|
|
def _stream_chat_completion_response_from_text_generation(
|
||
|
|
text_generation_output: Iterable[TextGenerationStreamOutput],
|
||
|
|
) -> Iterable[ChatCompletionStreamOutput]:
|
||
|
|
"""Used in `InferenceClient.chat_completion`."""
|
||
|
|
created = int(time.time())
|
||
|
|
for item in text_generation_output:
|
||
|
|
yield _format_chat_completion_stream_output_from_text_generation(item, created)
|
||
|
|
|
||
|
|
|
||
|
|
async def _async_stream_chat_completion_response_from_text_generation(
|
||
|
|
text_generation_output: AsyncIterable[TextGenerationStreamOutput],
|
||
|
|
) -> AsyncIterable[ChatCompletionStreamOutput]:
|
||
|
|
"""Used in `AsyncInferenceClient.chat_completion`."""
|
||
|
|
created = int(time.time())
|
||
|
|
async for item in text_generation_output:
|
||
|
|
yield _format_chat_completion_stream_output_from_text_generation(item, created)
|
||
|
|
|
||
|
|
|
||
|
|
def _format_chat_completion_stream_output_from_text_generation(
|
||
|
|
item: TextGenerationStreamOutput, created: int
|
||
|
|
) -> ChatCompletionStreamOutput:
|
||
|
|
if item.details is None:
|
||
|
|
# new token generated => return delta
|
||
|
|
return ChatCompletionStreamOutput(
|
||
|
|
choices=[
|
||
|
|
ChatCompletionStreamOutputChoice(
|
||
|
|
delta=ChatCompletionStreamOutputDelta(
|
||
|
|
role="assistant",
|
||
|
|
content=item.token.text,
|
||
|
|
),
|
||
|
|
finish_reason=None,
|
||
|
|
index=0,
|
||
|
|
)
|
||
|
|
],
|
||
|
|
created=created,
|
||
|
|
)
|
||
|
|
else:
|
||
|
|
# generation is completed => return finish reason
|
||
|
|
return ChatCompletionStreamOutput(
|
||
|
|
choices=[
|
||
|
|
ChatCompletionStreamOutputChoice(
|
||
|
|
delta=ChatCompletionStreamOutputDelta(),
|
||
|
|
finish_reason=item.details.finish_reason,
|
||
|
|
index=0,
|
||
|
|
)
|
||
|
|
],
|
||
|
|
created=created,
|
||
|
|
)
|
||
|
|
|
||
|
|
|
||
|
|
def _stream_chat_completion_response_from_bytes(
|
||
|
|
bytes_lines: Iterable[bytes],
|
||
|
|
) -> Iterable[ChatCompletionStreamOutput]:
|
||
|
|
"""Used in `InferenceClient.chat_completion` if model is served with TGI."""
|
||
|
|
for item in bytes_lines:
|
||
|
|
output = _format_chat_completion_stream_output_from_text_generation_from_bytes(item)
|
||
|
|
if output is not None:
|
||
|
|
yield output
|
||
|
|
|
||
|
|
|
||
|
|
async def _async_stream_chat_completion_response_from_bytes(
|
||
|
|
bytes_lines: AsyncIterable[bytes],
|
||
|
|
) -> AsyncIterable[ChatCompletionStreamOutput]:
|
||
|
|
"""Used in `AsyncInferenceClient.chat_completion`."""
|
||
|
|
async for item in bytes_lines:
|
||
|
|
output = _format_chat_completion_stream_output_from_text_generation_from_bytes(item)
|
||
|
|
if output is not None:
|
||
|
|
yield output
|
||
|
|
|
||
|
|
|
||
|
|
def _format_chat_completion_stream_output_from_text_generation_from_bytes(
|
||
|
|
byte_payload: bytes,
|
||
|
|
) -> Optional[ChatCompletionStreamOutput]:
|
||
|
|
if not byte_payload.startswith(b"data:"):
|
||
|
|
return None # empty line
|
||
|
|
|
||
|
|
# Decode payload
|
||
|
|
payload = byte_payload.decode("utf-8")
|
||
|
|
json_payload = json.loads(payload.lstrip("data:").rstrip("/n"))
|
||
|
|
return ChatCompletionStreamOutput.parse_obj_as_instance(json_payload)
|
||
|
|
|
||
|
|
|
||
|
|
async def _async_yield_from(client: "ClientSession", response: "ClientResponse") -> AsyncIterable[bytes]:
|
||
|
|
async for byte_payload in response.content:
|
||
|
|
yield byte_payload
|
||
|
|
await client.close()
|
||
|
|
|
||
|
|
|
||
|
|
# "TGI servers" are servers running with the `text-generation-inference` backend.
|
||
|
|
# This backend is the go-to solution to run large language models at scale. However,
|
||
|
|
# for some smaller models (e.g. "gpt2") the default `transformers` + `api-inference`
|
||
|
|
# solution is still in use.
|
||
|
|
#
|
||
|
|
# Both approaches have very similar APIs, but not exactly the same. What we do first in
|
||
|
|
# the `text_generation` method is to assume the model is served via TGI. If we realize
|
||
|
|
# it's not the case (i.e. we receive an HTTP 400 Bad Request), we fallback to the
|
||
|
|
# default API with a warning message. We remember for each model if it's a TGI server
|
||
|
|
# or not using `_NON_TGI_SERVERS` global variable.
|
||
|
|
#
|
||
|
|
# In addition, TGI servers have a built-in API route for chat-completion, which is not
|
||
|
|
# available on the default API. We use this route to provide a more consistent behavior
|
||
|
|
# when available.
|
||
|
|
#
|
||
|
|
# For more details, see https://github.com/huggingface/text-generation-inference and
|
||
|
|
# https://huggingface.co/docs/api-inference/detailed_parameters#text-generation-task.
|
||
|
|
|
||
|
|
_NON_TGI_SERVERS: Set[Optional[str]] = set()
|
||
|
|
|
||
|
|
|
||
|
|
def _set_as_non_tgi(model: Optional[str]) -> None:
|
||
|
|
_NON_TGI_SERVERS.add(model)
|
||
|
|
|
||
|
|
|
||
|
|
def _is_tgi_server(model: Optional[str]) -> bool:
|
||
|
|
return model not in _NON_TGI_SERVERS
|
||
|
|
|
||
|
|
|
||
|
|
_NON_CHAT_COMPLETION_SERVER: Set[str] = set()
|
||
|
|
|
||
|
|
|
||
|
|
def _set_as_non_chat_completion_server(model: str) -> None:
|
||
|
|
print("Set as non chat completion", model)
|
||
|
|
_NON_CHAT_COMPLETION_SERVER.add(model)
|
||
|
|
|
||
|
|
|
||
|
|
def _is_chat_completion_server(model: str) -> bool:
|
||
|
|
return model not in _NON_CHAT_COMPLETION_SERVER
|
||
|
|
|
||
|
|
|
||
|
|
# TEXT GENERATION ERRORS
|
||
|
|
# ----------------------
|
||
|
|
# Text-generation errors are parsed separately to handle as much as possible the errors returned by the text generation
|
||
|
|
# inference project (https://github.com/huggingface/text-generation-inference).
|
||
|
|
# ----------------------
|
||
|
|
|
||
|
|
|
||
|
|
def raise_text_generation_error(http_error: HTTPError) -> NoReturn:
|
||
|
|
"""
|
||
|
|
Try to parse text-generation-inference error message and raise HTTPError in any case.
|
||
|
|
|
||
|
|
Args:
|
||
|
|
error (`HTTPError`):
|
||
|
|
The HTTPError that have been raised.
|
||
|
|
"""
|
||
|
|
# Try to parse a Text Generation Inference error
|
||
|
|
|
||
|
|
try:
|
||
|
|
# Hacky way to retrieve payload in case of aiohttp error
|
||
|
|
payload = getattr(http_error, "response_error_payload", None) or http_error.response.json()
|
||
|
|
error = payload.get("error")
|
||
|
|
error_type = payload.get("error_type")
|
||
|
|
except Exception: # no payload
|
||
|
|
raise http_error
|
||
|
|
|
||
|
|
# If error_type => more information than `hf_raise_for_status`
|
||
|
|
if error_type is not None:
|
||
|
|
exception = _parse_text_generation_error(error, error_type)
|
||
|
|
raise exception from http_error
|
||
|
|
|
||
|
|
# Otherwise, fallback to default error
|
||
|
|
raise http_error
|
||
|
|
|
||
|
|
|
||
|
|
def _parse_text_generation_error(error: Optional[str], error_type: Optional[str]) -> TextGenerationError:
|
||
|
|
if error_type == "generation":
|
||
|
|
return GenerationError(error) # type: ignore
|
||
|
|
if error_type == "incomplete_generation":
|
||
|
|
return IncompleteGenerationError(error) # type: ignore
|
||
|
|
if error_type == "overloaded":
|
||
|
|
return OverloadedError(error) # type: ignore
|
||
|
|
if error_type == "validation":
|
||
|
|
return ValidationError(error) # type: ignore
|
||
|
|
return UnknownError(error) # type: ignore
|