kolandr/gridbot
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

Initial import: grid-bot — grid trading bot for BTC-USDT on Cifra Markets

Browse Source
This commit is contained in:
Kolp
2026-09-24 13:22:23 +07:00
commit 642cc11a9f
18968 changed files with 5683248 additions and 0 deletions
Download Patch File Download Diff File Expand all files Collapse all files
kronos-venv/lib/python3.12/site-packages/huggingface_hub/__init__.py
+1484
View File
File diff suppressed because it is too large Load Diff
kronos-venv/lib/python3.12/site-packages/huggingface_hub/_commit_api.py
+915
View File
@@ -0,0 +1,915 @@
"""
Type definitions and utilities for the `create_commit` API
"""
import base64
import io
import math
import os
import warnings
from collections import defaultdict
from contextlib import contextmanager
from dataclasses import dataclass, field
from itertools import groupby
from pathlib import Path, PurePosixPath
from typing import TYPE_CHECKING, Any, BinaryIO, Dict, Iterable, Iterator, List, Literal, Optional, Tuple, Union
from tqdm.contrib.concurrent import thread_map
from . import constants
from .errors import EntryNotFoundError, HfHubHTTPError, XetAuthorizationError, XetRefreshTokenError
from .file_download import hf_hub_url
from .lfs import UploadInfo, lfs_upload, post_lfs_batch_info
from .utils import (
FORBIDDEN_FOLDERS,
XetTokenType,
chunk_iterable,
fetch_xet_connection_info_from_repo_info,
get_session,
hf_raise_for_status,
logging,
sha,
tqdm_stream_file,
validate_hf_hub_args,
)
from .utils import tqdm as hf_tqdm
from .utils.tqdm import _get_progress_bar_context
if TYPE_CHECKING:
from .hf_api import RepoFile
logger = logging.get_logger(__name__)
UploadMode = Literal["lfs", "regular"]
# Max is 1,000 per request on the Hub for HfApi.get_paths_info
# Otherwise we get:
# HfHubHTTPError: 413 Client Error: Payload Too Large for url: https://huggingface.co/api/datasets/xxx (Request ID: xxx)\n\ntoo many parameters
# See https://github.com/huggingface/huggingface_hub/issues/1503
FETCH_LFS_BATCH_SIZE = 500
UPLOAD_BATCH_MAX_NUM_FILES = 256
@dataclass
class CommitOperationDelete:
"""
Data structure holding necessary info to delete a file or a folder from a repository
on the Hub.
Args:
path_in_repo (`str`):
Relative filepath in the repo, for example: `"checkpoints/1fec34a/weights.bin"`
for a file or `"checkpoints/1fec34a/"` for a folder.
is_folder (`bool` or `Literal["auto"]`, *optional*)
Whether the Delete Operation applies to a folder or not. If "auto", the path
type (file or folder) is guessed automatically by looking if path ends with
a "/" (folder) or not (file). To explicitly set the path type, you can set
`is_folder=True` or `is_folder=False`.
"""
path_in_repo: str
is_folder: Union[bool, Literal["auto"]] = "auto"
def __post_init__(self):
self.path_in_repo = _validate_path_in_repo(self.path_in_repo)
if self.is_folder == "auto":
self.is_folder = self.path_in_repo.endswith("/")
if not isinstance(self.is_folder, bool):
raise ValueError(
f"Wrong value for `is_folder`. Must be one of [`True`, `False`, `'auto'`]. Got '{self.is_folder}'."
)
@dataclass
class CommitOperationCopy:
"""
Data structure holding necessary info to copy a file in a repository on the Hub.
Limitations:
- Only LFS files can be copied. To copy a regular file, you need to download it locally and re-upload it
- Cross-repository copies are not supported.
Note: you can combine a [`CommitOperationCopy`] and a [`CommitOperationDelete`] to rename an LFS file on the Hub.
Args:
src_path_in_repo (`str`):
Relative filepath in the repo of the file to be copied, e.g. `"checkpoints/1fec34a/weights.bin"`.
path_in_repo (`str`):
Relative filepath in the repo where to copy the file, e.g. `"checkpoints/1fec34a/weights_copy.bin"`.
src_revision (`str`, *optional*):
The git revision of the file to be copied. Can be any valid git revision.
Default to the target commit revision.
"""
src_path_in_repo: str
path_in_repo: str
src_revision: Optional[str] = None
# set to the OID of the file to be copied if it has already been uploaded
# useful to determine if a commit will be empty or not.
_src_oid: Optional[str] = None
# set to the OID of the file to copy to if it has already been uploaded
# useful to determine if a commit will be empty or not.
_dest_oid: Optional[str] = None
def __post_init__(self):
self.src_path_in_repo = _validate_path_in_repo(self.src_path_in_repo)
self.path_in_repo = _validate_path_in_repo(self.path_in_repo)
@dataclass
class CommitOperationAdd:
"""
Data structure holding necessary info to upload a file to a repository on the Hub.
Args:
path_in_repo (`str`):
Relative filepath in the repo, for example: `"checkpoints/1fec34a/weights.bin"`
path_or_fileobj (`str`, `Path`, `bytes`, or `BinaryIO`):
Either:
- a path to a local file (as `str` or `pathlib.Path`) to upload
- a buffer of bytes (`bytes`) holding the content of the file to upload
- a "file object" (subclass of `io.BufferedIOBase`), typically obtained
with `open(path, "rb")`. It must support `seek()` and `tell()` methods.
Raises:
[`ValueError`](https://docs.python.org/3/library/exceptions.html#ValueError)
If `path_or_fileobj` is not one of `str`, `Path`, `bytes` or `io.BufferedIOBase`.
[`ValueError`](https://docs.python.org/3/library/exceptions.html#ValueError)
If `path_or_fileobj` is a `str` or `Path` but not a path to an existing file.
[`ValueError`](https://docs.python.org/3/library/exceptions.html#ValueError)
If `path_or_fileobj` is a `io.BufferedIOBase` but it doesn't support both
`seek()` and `tell()`.
"""
path_in_repo: str
path_or_fileobj: Union[str, Path, bytes, BinaryIO]
upload_info: UploadInfo = field(init=False, repr=False)
# Internal attributes
# set to "lfs" or "regular" once known
_upload_mode: Optional[UploadMode] = field(init=False, repr=False, default=None)
# set to True if .gitignore rules prevent the file from being uploaded as LFS
# (server-side check)
_should_ignore: Optional[bool] = field(init=False, repr=False, default=None)
# set to the remote OID of the file if it has already been uploaded
# useful to determine if a commit will be empty or not
_remote_oid: Optional[str] = field(init=False, repr=False, default=None)
# set to True once the file has been uploaded as LFS
_is_uploaded: bool = field(init=False, repr=False, default=False)
# set to True once the file has been committed
_is_committed: bool = field(init=False, repr=False, default=False)
def __post_init__(self) -> None:
"""Validates `path_or_fileobj` and compute `upload_info`."""
self.path_in_repo = _validate_path_in_repo(self.path_in_repo)
# Validate `path_or_fileobj` value
if isinstance(self.path_or_fileobj, Path):
self.path_or_fileobj = str(self.path_or_fileobj)
if isinstance(self.path_or_fileobj, str):
path_or_fileobj = os.path.normpath(os.path.expanduser(self.path_or_fileobj))
if not os.path.isfile(path_or_fileobj):
raise ValueError(f"Provided path: '{path_or_fileobj}' is not a file on the local file system")
elif not isinstance(self.path_or_fileobj, (io.BufferedIOBase, bytes)):
# ^^ Inspired from: https://stackoverflow.com/questions/44584829/how-to-determine-if-file-is-opened-in-binary-or-text-mode
raise ValueError(
"path_or_fileobj must be either an instance of str, bytes or"
" io.BufferedIOBase. If you passed a file-like object, make sure it is"
" in binary mode."
)
if isinstance(self.path_or_fileobj, io.BufferedIOBase):
try:
self.path_or_fileobj.tell()
self.path_or_fileobj.seek(0, os.SEEK_CUR)
except (OSError, AttributeError) as exc:
raise ValueError(
"path_or_fileobj is a file-like object but does not implement seek() and tell()"
) from exc
# Compute "upload_info" attribute
if isinstance(self.path_or_fileobj, str):
self.upload_info = UploadInfo.from_path(self.path_or_fileobj)
elif isinstance(self.path_or_fileobj, bytes):
self.upload_info = UploadInfo.from_bytes(self.path_or_fileobj)
else:
self.upload_info = UploadInfo.from_fileobj(self.path_or_fileobj)
@contextmanager
def as_file(self, with_tqdm: bool = False) -> Iterator[BinaryIO]:
"""
A context manager that yields a file-like object allowing to read the underlying
data behind `path_or_fileobj`.
Args:
with_tqdm (`bool`, *optional*, defaults to `False`):
If True, iterating over the file object will display a progress bar. Only
works if the file-like object is a path to a file. Pure bytes and buffers
are not supported.
Example:
```python
>>> operation = CommitOperationAdd(
... path_in_repo="remote/dir/weights.h5",
... path_or_fileobj="./local/weights.h5",
... )
CommitOperationAdd(path_in_repo='remote/dir/weights.h5', path_or_fileobj='./local/weights.h5')
>>> with operation.as_file() as file:
... content = file.read()
>>> with operation.as_file(with_tqdm=True) as file:
... while True:
... data = file.read(1024)
... if not data:
... break
config.json: 100%|█████████████████████████| 8.19k/8.19k [00:02<00:00, 3.72kB/s]
>>> with operation.as_file(with_tqdm=True) as file:
... requests.put(..., data=file)
config.json: 100%|█████████████████████████| 8.19k/8.19k [00:02<00:00, 3.72kB/s]
```
"""
if isinstance(self.path_or_fileobj, str) or isinstance(self.path_or_fileobj, Path):
if with_tqdm:
with tqdm_stream_file(self.path_or_fileobj) as file:
yield file
else:
with open(self.path_or_fileobj, "rb") as file:
yield file
elif isinstance(self.path_or_fileobj, bytes):
yield io.BytesIO(self.path_or_fileobj)
elif isinstance(self.path_or_fileobj, io.BufferedIOBase):
prev_pos = self.path_or_fileobj.tell()
yield self.path_or_fileobj
self.path_or_fileobj.seek(prev_pos, io.SEEK_SET)
def b64content(self) -> bytes:
"""
The base64-encoded content of `path_or_fileobj`
Returns: `bytes`
"""
with self.as_file() as file:
return base64.b64encode(file.read())
@property
def _local_oid(self) -> Optional[str]:
"""Return the OID of the local file.
This OID is then compared to `self._remote_oid` to check if the file has changed compared to the remote one.
If the file did not change, we won't upload it again to prevent empty commits.
For LFS files, the OID corresponds to the SHA256 of the file content (used a LFS ref).
For regular files, the OID corresponds to the SHA1 of the file content.
Note: this is slightly different to git OID computation since the oid of an LFS file is usually the git-SHA1 of the
pointer file content (not the actual file content). However, using the SHA256 is enough to detect changes
and more convenient client-side.
"""
if self._upload_mode is None:
return None
elif self._upload_mode == "lfs":
return self.upload_info.sha256.hex()
else:
# Regular file => compute sha1
# => no need to read by chunk since the file is guaranteed to be <=5MB.
with self.as_file() as file:
return sha.git_hash(file.read())
def _validate_path_in_repo(path_in_repo: str) -> str:
# Validate `path_in_repo` value to prevent a server-side issue
if path_in_repo.startswith("/"):
path_in_repo = path_in_repo[1:]
if path_in_repo == "." or path_in_repo == ".." or path_in_repo.startswith("../"):
raise ValueError(f"Invalid `path_in_repo` in CommitOperation: '{path_in_repo}'")
if path_in_repo.startswith("./"):
path_in_repo = path_in_repo[2:]
for forbidden in FORBIDDEN_FOLDERS:
if any(part == forbidden for part in path_in_repo.split("/")):
raise ValueError(
f"Invalid `path_in_repo` in CommitOperation: cannot update files under a '{forbidden}/' folder (path:"
f" '{path_in_repo}')."
)
return path_in_repo
CommitOperation = Union[CommitOperationAdd, CommitOperationCopy, CommitOperationDelete]
def _warn_on_overwriting_operations(operations: List[CommitOperation]) -> None:
"""
Warn user when a list of operations is expected to overwrite itself in a single
commit.
Rules:
- If a filepath is updated by multiple `CommitOperationAdd` operations, a warning
message is triggered.
- If a filepath is updated at least once by a `CommitOperationAdd` and then deleted
by a `CommitOperationDelete`, a warning is triggered.
- If a `CommitOperationDelete` deletes a filepath that is then updated by a
`CommitOperationAdd`, no warning is triggered. This is usually useless (no need to
delete before upload) but can happen if a user deletes an entire folder and then
add new files to it.
"""
nb_additions_per_path: Dict[str, int] = defaultdict(int)
for operation in operations:
path_in_repo = operation.path_in_repo
if isinstance(operation, CommitOperationAdd):
if nb_additions_per_path[path_in_repo] > 0:
warnings.warn(
"About to update multiple times the same file in the same commit:"
f" '{path_in_repo}'. This can cause undesired inconsistencies in"
" your repo."
)
nb_additions_per_path[path_in_repo] += 1
for parent in PurePosixPath(path_in_repo).parents:
# Also keep track of number of updated files per folder
# => warns if deleting a folder overwrite some contained files
nb_additions_per_path[str(parent)] += 1
if isinstance(operation, CommitOperationDelete):
if nb_additions_per_path[str(PurePosixPath(path_in_repo))] > 0:
if operation.is_folder:
warnings.warn(
"About to delete a folder containing files that have just been"
f" updated within the same commit: '{path_in_repo}'. This can"
" cause undesired inconsistencies in your repo."
)
else:
warnings.warn(
"About to delete a file that have just been updated within the"
f" same commit: '{path_in_repo}'. This can cause undesired"
" inconsistencies in your repo."
)
@validate_hf_hub_args
def _upload_lfs_files(
*,
additions: List[CommitOperationAdd],
repo_type: str,
repo_id: str,
headers: Dict[str, str],
endpoint: Optional[str] = None,
num_threads: int = 5,
revision: Optional[str] = None,
):
"""
Uploads the content of `additions` to the Hub using the large file storage protocol.
Relevant external documentation:
- LFS Batch API: https://github.com/git-lfs/git-lfs/blob/main/docs/api/batch.md
Args:
additions (`List` of `CommitOperationAdd`):
The files to be uploaded
repo_type (`str`):
Type of the repo to upload to: `"model"`, `"dataset"` or `"space"`.
repo_id (`str`):
A namespace (user or an organization) and a repo name separated
by a `/`.
headers (`Dict[str, str]`):
Headers to use for the request, including authorization headers and user agent.
num_threads (`int`, *optional*):
The number of concurrent threads to use when uploading. Defaults to 5.
revision (`str`, *optional*):
The git revision to upload to.
Raises:
[`EnvironmentError`](https://docs.python.org/3/library/exceptions.html#EnvironmentError)
If an upload failed for any reason
[`ValueError`](https://docs.python.org/3/library/exceptions.html#ValueError)
If the server returns malformed responses
[`HTTPError`](https://requests.readthedocs.io/en/latest/api/#requests.HTTPError)
If the LFS batch endpoint returned an HTTP error.
"""
# Step 1: retrieve upload instructions from the LFS batch endpoint.
# Upload instructions are retrieved by chunk of 256 files to avoid reaching
# the payload limit.
batch_actions: List[Dict] = []
for chunk in chunk_iterable(additions, chunk_size=UPLOAD_BATCH_MAX_NUM_FILES):
batch_actions_chunk, batch_errors_chunk = post_lfs_batch_info(
upload_infos=[op.upload_info for op in chunk],
repo_id=repo_id,
repo_type=repo_type,
revision=revision,
endpoint=endpoint,
headers=headers,
token=None, # already passed in 'headers'
)
# If at least 1 error, we do not retrieve information for other chunks
if batch_errors_chunk:
message = "\n".join(
[
f"Encountered error for file with OID {err.get('oid')}: `{err.get('error', {}).get('message')}"
for err in batch_errors_chunk
]
)
raise ValueError(f"LFS batch endpoint returned errors:\n{message}")
batch_actions += batch_actions_chunk
oid2addop = {add_op.upload_info.sha256.hex(): add_op for add_op in additions}
# Step 2: ignore files that have already been uploaded
filtered_actions = []
for action in batch_actions:
if action.get("actions") is None:
logger.debug(
f"Content of file {oid2addop[action['oid']].path_in_repo} is already"
" present upstream - skipping upload."
)
else:
filtered_actions.append(action)
if len(filtered_actions) == 0:
logger.debug("No LFS files to upload.")
return
# Step 3: upload files concurrently according to these instructions
def _wrapped_lfs_upload(batch_action) -> None:
try:
operation = oid2addop[batch_action["oid"]]
lfs_upload(operation=operation, lfs_batch_action=batch_action, headers=headers, endpoint=endpoint)
except Exception as exc:
raise RuntimeError(f"Error while uploading '{operation.path_in_repo}' to the Hub.") from exc
if constants.HF_HUB_ENABLE_HF_TRANSFER:
logger.debug(f"Uploading {len(filtered_actions)} LFS files to the Hub using `hf_transfer`.")
for action in hf_tqdm(filtered_actions, name="huggingface_hub.lfs_upload"):
_wrapped_lfs_upload(action)
elif len(filtered_actions) == 1:
logger.debug("Uploading 1 LFS file to the Hub")
_wrapped_lfs_upload(filtered_actions[0])
else:
logger.debug(
f"Uploading {len(filtered_actions)} LFS files to the Hub using up to {num_threads} threads concurrently"
)
thread_map(
_wrapped_lfs_upload,
filtered_actions,
desc=f"Upload {len(filtered_actions)} LFS files",
max_workers=num_threads,
tqdm_class=hf_tqdm,
)
@validate_hf_hub_args
def _upload_xet_files(
*,
additions: List[CommitOperationAdd],
repo_type: str,
repo_id: str,
headers: Dict[str, str],
endpoint: Optional[str] = None,
revision: Optional[str] = None,
create_pr: Optional[bool] = None,
):
"""
Uploads the content of `additions` to the Hub using the xet storage protocol.
This chunks the files and deduplicates the chunks before uploading them to xetcas storage.
Args:
additions (`List` of `CommitOperationAdd`):
The files to be uploaded.
repo_type (`str`):
Type of the repo to upload to: `"model"`, `"dataset"` or `"space"`.
repo_id (`str`):
A namespace (user or an organization) and a repo name separated
by a `/`.
headers (`Dict[str, str]`):
Headers to use for the request, including authorization headers and user agent.
endpoint: (`str`, *optional*):
The endpoint to use for the xetcas service. Defaults to `constants.ENDPOINT`.
revision (`str`, *optional*):
The git revision to upload to.
create_pr (`bool`, *optional*):
Whether or not to create a Pull Request with that commit.
Raises:
[`EnvironmentError`](https://docs.python.org/3/library/exceptions.html#EnvironmentError)
If an upload failed for any reason.
[`ValueError`](https://docs.python.org/3/library/exceptions.html#ValueError)
If the server returns malformed responses or if the user is unauthorized to upload to xet storage.
[`HTTPError`](https://requests.readthedocs.io/en/latest/api/#requests.HTTPError)
If the LFS batch endpoint returned an HTTP error.
**How it works:**
The file download system uses Xet storage, which is a content-addressable storage system that breaks files into chunks
for efficient storage and transfer.
`hf_xet.upload_files` manages uploading files by:
- Taking a list of file paths to upload
- Breaking files into smaller chunks for efficient storage
- Avoiding duplicate storage by recognizing identical chunks across files
- Connecting to a storage server (CAS server) that manages these chunks
The upload process works like this:
1. Create a local folder at ~/.cache/huggingface/xet/chunk-cache to store file chunks for reuse.
2. Process files in parallel (up to 8 files at once):
2.1. Read the file content.
2.2. Split the file content into smaller chunks based on content patterns: each chunk gets a unique ID based on what's in it.
2.3. For each chunk:
- Check if it already exists in storage.
- Skip uploading chunks that already exist.
2.4. Group chunks into larger blocks for efficient transfer.
2.5. Upload these blocks to the storage server.
2.6. Create and upload information about how the file is structured.
3. Return reference files that contain information about the uploaded files, which can be used later to download them.
"""
if len(additions) == 0:
return
# at this point, we know that hf_xet is installed
from hf_xet import upload_bytes, upload_files
try:
xet_connection_info = fetch_xet_connection_info_from_repo_info(
token_type=XetTokenType.WRITE,
repo_id=repo_id,
repo_type=repo_type,
revision=revision,
headers=headers,
endpoint=endpoint,
params={"create_pr": "1"} if create_pr else None,
)
except HfHubHTTPError as e:
if e.response.status_code == 401:
raise XetAuthorizationError(
f"You are unauthorized to upload to xet storage for {repo_type}/{repo_id}. "
f"Please check that you have configured your access token with write access to the repo."
) from e
raise
xet_endpoint = xet_connection_info.endpoint
access_token_info = (xet_connection_info.access_token, xet_connection_info.expiration_unix_epoch)
def token_refresher() -> Tuple[str, int]:
new_xet_connection = fetch_xet_connection_info_from_repo_info(
token_type=XetTokenType.WRITE,
repo_id=repo_id,
repo_type=repo_type,
revision=revision,
headers=headers,
endpoint=endpoint,
params={"create_pr": "1"} if create_pr else None,
)
if new_xet_connection is None:
raise XetRefreshTokenError("Failed to refresh xet token")
return new_xet_connection.access_token, new_xet_connection.expiration_unix_epoch
num_chunks = math.ceil(len(additions) / UPLOAD_BATCH_MAX_NUM_FILES)
num_chunks_num_digits = int(math.log10(num_chunks)) + 1
for i, chunk in enumerate(chunk_iterable(additions, chunk_size=UPLOAD_BATCH_MAX_NUM_FILES)):
_chunk = [op for op in chunk]
bytes_ops = [op for op in _chunk if isinstance(op.path_or_fileobj, bytes)]
paths_ops = [op for op in _chunk if isinstance(op.path_or_fileobj, (str, Path))]
expected_size = sum(op.upload_info.size for op in bytes_ops + paths_ops)
if num_chunks > 1:
description = f"Uploading Batch [{str(i + 1).zfill(num_chunks_num_digits)}/{num_chunks}]..."
else:
description = "Uploading..."
progress_cm = _get_progress_bar_context(
desc=description,
total=expected_size,
initial=0,
unit="B",
unit_scale=True,
name="huggingface_hub.xet_put",
log_level=logger.getEffectiveLevel(),
)
with progress_cm as progress:
def update_progress(increment: int):
progress.update(increment)
if len(paths_ops) > 0:
upload_files(
[str(op.path_or_fileobj) for op in paths_ops],
xet_endpoint,
access_token_info,
token_refresher,
update_progress,
repo_type,
)
if len(bytes_ops) > 0:
upload_bytes(
[op.path_or_fileobj for op in bytes_ops],
xet_endpoint,
access_token_info,
token_refresher,
update_progress,
repo_type,
)
return
def _validate_preupload_info(preupload_info: dict):
files = preupload_info.get("files")
if not isinstance(files, list):
raise ValueError("preupload_info is improperly formatted")
for file_info in files:
if not (
isinstance(file_info, dict)
and isinstance(file_info.get("path"), str)
and isinstance(file_info.get("uploadMode"), str)
and (file_info["uploadMode"] in ("lfs", "regular"))
):
raise ValueError("preupload_info is improperly formatted:")
return preupload_info
@validate_hf_hub_args
def _fetch_upload_modes(
additions: Iterable[CommitOperationAdd],
repo_type: str,
repo_id: str,
headers: Dict[str, str],
revision: str,
endpoint: Optional[str] = None,
create_pr: bool = False,
gitignore_content: Optional[str] = None,
) -> None:
"""
Requests the Hub "preupload" endpoint to determine whether each input file should be uploaded as a regular git blob,
as a git LFS blob, or as a XET file. Input `additions` are mutated in-place with the upload mode.
Args:
additions (`Iterable` of :class:`CommitOperationAdd`):
Iterable of :class:`CommitOperationAdd` describing the files to
upload to the Hub.
repo_type (`str`):
Type of the repo to upload to: `"model"`, `"dataset"` or `"space"`.
repo_id (`str`):
A namespace (user or an organization) and a repo name separated
by a `/`.
headers (`Dict[str, str]`):
Headers to use for the request, including authorization headers and user agent.
revision (`str`):
The git revision to upload the files to. Can be any valid git revision.
gitignore_content (`str`, *optional*):
The content of the `.gitignore` file to know which files should be ignored. The order of priority
is to first check if `gitignore_content` is passed, then check if the `.gitignore` file is present
in the list of files to commit and finally default to the `.gitignore` file already hosted on the Hub
(if any).
Raises:
[`~utils.HfHubHTTPError`]
If the Hub API returned an error.
[`ValueError`](https://docs.python.org/3/library/exceptions.html#ValueError)
If the Hub API response is improperly formatted.
"""
endpoint = endpoint if endpoint is not None else constants.ENDPOINT
# Fetch upload mode (LFS or regular) chunk by chunk.
upload_modes: Dict[str, UploadMode] = {}
should_ignore_info: Dict[str, bool] = {}
oid_info: Dict[str, Optional[str]] = {}
for chunk in chunk_iterable(additions, 256):
payload: Dict = {
"files": [
{
"path": op.path_in_repo,
"sample": base64.b64encode(op.upload_info.sample).decode("ascii"),
"size": op.upload_info.size,
}
for op in chunk
]
}
if gitignore_content is not None:
payload["gitIgnore"] = gitignore_content
resp = get_session().post(
f"{endpoint}/api/{repo_type}s/{repo_id}/preupload/{revision}",
json=payload,
headers=headers,
params={"create_pr": "1"} if create_pr else None,
)
hf_raise_for_status(resp)
preupload_info = _validate_preupload_info(resp.json())
upload_modes.update(**{file["path"]: file["uploadMode"] for file in preupload_info["files"]})
should_ignore_info.update(**{file["path"]: file["shouldIgnore"] for file in preupload_info["files"]})
oid_info.update(**{file["path"]: file.get("oid") for file in preupload_info["files"]})
# Set upload mode for each addition operation
for addition in additions:
addition._upload_mode = upload_modes[addition.path_in_repo]
addition._should_ignore = should_ignore_info[addition.path_in_repo]
addition._remote_oid = oid_info[addition.path_in_repo]
# Empty files cannot be uploaded as LFS (S3 would fail with a 501 Not Implemented)
# => empty files are uploaded as "regular" to still allow users to commit them.
for addition in additions:
if addition.upload_info.size == 0:
addition._upload_mode = "regular"
@validate_hf_hub_args
def _fetch_files_to_copy(
copies: Iterable[CommitOperationCopy],
repo_type: str,
repo_id: str,
headers: Dict[str, str],
revision: str,
endpoint: Optional[str] = None,
) -> Dict[Tuple[str, Optional[str]], Union["RepoFile", bytes]]:
"""
Fetch information about the files to copy.
For LFS files, we only need their metadata (file size and sha256) while for regular files
we need to download the raw content from the Hub.
Args:
copies (`Iterable` of :class:`CommitOperationCopy`):
Iterable of :class:`CommitOperationCopy` describing the files to
copy on the Hub.
repo_type (`str`):
Type of the repo to upload to: `"model"`, `"dataset"` or `"space"`.
repo_id (`str`):
A namespace (user or an organization) and a repo name separated
by a `/`.
headers (`Dict[str, str]`):
Headers to use for the request, including authorization headers and user agent.
revision (`str`):
The git revision to upload the files to. Can be any valid git revision.
Returns: `Dict[Tuple[str, Optional[str]], Union[RepoFile, bytes]]]`
Key is the file path and revision of the file to copy.
Value is the raw content as bytes (for regular files) or the file information as a RepoFile (for LFS files).
Raises:
[`~utils.HfHubHTTPError`]
If the Hub API returned an error.
[`ValueError`](https://docs.python.org/3/library/exceptions.html#ValueError)
If the Hub API response is improperly formatted.
"""
from .hf_api import HfApi, RepoFolder
hf_api = HfApi(endpoint=endpoint, headers=headers)
files_to_copy: Dict[Tuple[str, Optional[str]], Union["RepoFile", bytes]] = {}
# Store (path, revision) -> oid mapping
oid_info: Dict[Tuple[str, Optional[str]], Optional[str]] = {}
# 1. Fetch OIDs for destination paths in batches.
dest_paths = [op.path_in_repo for op in copies]
for offset in range(0, len(dest_paths), FETCH_LFS_BATCH_SIZE):
dest_repo_files = hf_api.get_paths_info(
repo_id=repo_id,
paths=dest_paths[offset : offset + FETCH_LFS_BATCH_SIZE],
revision=revision,
repo_type=repo_type,
)
for file in dest_repo_files:
if not isinstance(file, RepoFolder):
oid_info[(file.path, revision)] = file.blob_id
# 2. Group by source revision and fetch source file info in batches.
for src_revision, operations in groupby(copies, key=lambda op: op.src_revision):
operations = list(operations) # type: ignore
src_paths = [op.src_path_in_repo for op in operations]
for offset in range(0, len(src_paths), FETCH_LFS_BATCH_SIZE):
src_repo_files = hf_api.get_paths_info(
repo_id=repo_id,
paths=src_paths[offset : offset + FETCH_LFS_BATCH_SIZE],
revision=src_revision or revision,
repo_type=repo_type,
)
for src_repo_file in src_repo_files:
if isinstance(src_repo_file, RepoFolder):
raise NotImplementedError("Copying a folder is not implemented.")
oid_info[(src_repo_file.path, src_revision)] = src_repo_file.blob_id
# If it's an LFS file, store the RepoFile object. Otherwise, download raw bytes.
if src_repo_file.lfs:
files_to_copy[(src_repo_file.path, src_revision)] = src_repo_file
else:
# TODO: (optimization) download regular files to copy concurrently
url = hf_hub_url(
endpoint=endpoint,
repo_type=repo_type,
repo_id=repo_id,
revision=src_revision or revision,
filename=src_repo_file.path,
)
response = get_session().get(url, headers=headers)
hf_raise_for_status(response)
files_to_copy[(src_repo_file.path, src_revision)] = response.content
# 3. Ensure all operations found a corresponding file in the Hub
# and track src/dest OIDs for each operation.
for operation in operations:
if (operation.src_path_in_repo, src_revision) not in files_to_copy:
raise EntryNotFoundError(
f"Cannot copy {operation.src_path_in_repo} at revision "
f"{src_revision or revision}: file is missing on repo."
)
operation._src_oid = oid_info.get((operation.src_path_in_repo, operation.src_revision))
operation._dest_oid = oid_info.get((operation.path_in_repo, revision))
return files_to_copy
def _prepare_commit_payload(
operations: Iterable[CommitOperation],
files_to_copy: Dict[Tuple[str, Optional[str]], Union["RepoFile", bytes]],
commit_message: str,
commit_description: Optional[str] = None,
parent_commit: Optional[str] = None,
) -> Iterable[Dict[str, Any]]:
"""
Builds the payload to POST to the `/commit` API of the Hub.
Payload is returned as an iterator so that it can be streamed as a ndjson in the
POST request.
For more information, see:
- https://github.com/huggingface/huggingface_hub/issues/1085#issuecomment-1265208073
- http://ndjson.org/
"""
commit_description = commit_description if commit_description is not None else ""
# 1. Send a header item with the commit metadata
header_value = {"summary": commit_message, "description": commit_description}
if parent_commit is not None:
header_value["parentCommit"] = parent_commit
yield {"key": "header", "value": header_value}
nb_ignored_files = 0
# 2. Send operations, one per line
for operation in operations:
# Skip ignored files
if isinstance(operation, CommitOperationAdd) and operation._should_ignore:
logger.debug(f"Skipping file '{operation.path_in_repo}' in commit (ignored by gitignore file).")
nb_ignored_files += 1
continue
# 2.a. Case adding a regular file
if isinstance(operation, CommitOperationAdd) and operation._upload_mode == "regular":
yield {
"key": "file",
"value": {
"content": operation.b64content().decode(),
"path": operation.path_in_repo,
"encoding": "base64",
},
}
# 2.b. Case adding an LFS file
elif isinstance(operation, CommitOperationAdd) and operation._upload_mode == "lfs":
yield {
"key": "lfsFile",
"value": {
"path": operation.path_in_repo,
"algo": "sha256",
"oid": operation.upload_info.sha256.hex(),
"size": operation.upload_info.size,
},
}
# 2.c. Case deleting a file or folder
elif isinstance(operation, CommitOperationDelete):
yield {
"key": "deletedFolder" if operation.is_folder else "deletedFile",
"value": {"path": operation.path_in_repo},
}
# 2.d. Case copying a file or folder
elif isinstance(operation, CommitOperationCopy):
file_to_copy = files_to_copy[(operation.src_path_in_repo, operation.src_revision)]
if isinstance(file_to_copy, bytes):
yield {
"key": "file",
"value": {
"content": base64.b64encode(file_to_copy).decode(),
"path": operation.path_in_repo,
"encoding": "base64",
},
}
elif file_to_copy.lfs:
yield {
"key": "lfsFile",
"value": {
"path": operation.path_in_repo,
"algo": "sha256",
"oid": file_to_copy.lfs.sha256,
},
}
else:
raise ValueError(
"Malformed files_to_copy (should be raw file content as bytes or RepoFile objects with LFS info."
)
# 2.e. Never expected to happen
else:
raise ValueError(
f"Unknown operation to commit. Operation: {operation}. Upload mode:"
f" {getattr(operation, '_upload_mode', None)}"
)
if nb_ignored_files > 0:
logger.info(f"Skipped {nb_ignored_files} file(s) in commit (ignored by gitignore file).")
kronos-venv/lib/python3.12/site-packages/huggingface_hub/_commit_scheduler.py
+353
View File
@@ -0,0 +1,353 @@
import atexit
import logging
import os
import time
from concurrent.futures import Future
from dataclasses import dataclass
from io import SEEK_END, SEEK_SET, BytesIO
from pathlib import Path
from threading import Lock, Thread
from typing import Dict, List, Optional, Union
from .hf_api import DEFAULT_IGNORE_PATTERNS, CommitInfo, CommitOperationAdd, HfApi
from .utils import filter_repo_objects
logger = logging.getLogger(__name__)
@dataclass(frozen=True)
class _FileToUpload:
"""Temporary dataclass to store info about files to upload. Not meant to be used directly."""
local_path: Path
path_in_repo: str
size_limit: int
last_modified: float
class CommitScheduler:
"""
Scheduler to upload a local folder to the Hub at regular intervals (e.g. push to hub every 5 minutes).
The recommended way to use the scheduler is to use it as a context manager. This ensures that the scheduler is
properly stopped and the last commit is triggered when the script ends. The scheduler can also be stopped manually
with the `stop` method. Checkout the [upload guide](https://huggingface.co/docs/huggingface_hub/guides/upload#scheduled-uploads)
to learn more about how to use it.
Args:
repo_id (`str`):
The id of the repo to commit to.
folder_path (`str` or `Path`):
Path to the local folder to upload regularly.
every (`int` or `float`, *optional*):
The number of minutes between each commit. Defaults to 5 minutes.
path_in_repo (`str`, *optional*):
Relative path of the directory in the repo, for example: `"checkpoints/"`. Defaults to the root folder
of the repository.
repo_type (`str`, *optional*):
The type of the repo to commit to. Defaults to `model`.
revision (`str`, *optional*):
The revision of the repo to commit to. Defaults to `main`.
private (`bool`, *optional*):
Whether to make the repo private. If `None` (default), the repo will be public unless the organization's default is private. This value is ignored if the repo already exists.
token (`str`, *optional*):
The token to use to commit to the repo. Defaults to the token saved on the machine.
allow_patterns (`List[str]` or `str`, *optional*):
If provided, only files matching at least one pattern are uploaded.
ignore_patterns (`List[str]` or `str`, *optional*):
If provided, files matching any of the patterns are not uploaded.
squash_history (`bool`, *optional*):
Whether to squash the history of the repo after each commit. Defaults to `False`. Squashing commits is
useful to avoid degraded performances on the repo when it grows too large.
hf_api (`HfApi`, *optional*):
The [`HfApi`] client to use to commit to the Hub. Can be set with custom settings (user agent, token,...).
Example:
```py
>>> from pathlib import Path
>>> from huggingface_hub import CommitScheduler
# Scheduler uploads every 10 minutes
>>> csv_path = Path("watched_folder/data.csv")
>>> CommitScheduler(repo_id="test_scheduler", repo_type="dataset", folder_path=csv_path.parent, every=10)
>>> with csv_path.open("a") as f:
... f.write("first line")
# Some time later (...)
>>> with csv_path.open("a") as f:
... f.write("second line")
```
Example using a context manager:
```py
>>> from pathlib import Path
>>> from huggingface_hub import CommitScheduler
>>> with CommitScheduler(repo_id="test_scheduler", repo_type="dataset", folder_path="watched_folder", every=10) as scheduler:
... csv_path = Path("watched_folder/data.csv")
... with csv_path.open("a") as f:
... f.write("first line")
... (...)
... with csv_path.open("a") as f:
... f.write("second line")
# Scheduler is now stopped and last commit have been triggered
```
"""
def __init__(
self,
*,
repo_id: str,
folder_path: Union[str, Path],
every: Union[int, float] = 5,
path_in_repo: Optional[str] = None,
repo_type: Optional[str] = None,
revision: Optional[str] = None,
private: Optional[bool] = None,
token: Optional[str] = None,
allow_patterns: Optional[Union[List[str], str]] = None,
ignore_patterns: Optional[Union[List[str], str]] = None,
squash_history: bool = False,
hf_api: Optional["HfApi"] = None,
) -> None:
self.api = hf_api or HfApi(token=token)
# Folder
self.folder_path = Path(folder_path).expanduser().resolve()
self.path_in_repo = path_in_repo or ""
self.allow_patterns = allow_patterns
if ignore_patterns is None:
ignore_patterns = []
elif isinstance(ignore_patterns, str):
ignore_patterns = [ignore_patterns]
self.ignore_patterns = ignore_patterns + DEFAULT_IGNORE_PATTERNS
if self.folder_path.is_file():
raise ValueError(f"'folder_path' must be a directory, not a file: '{self.folder_path}'.")
self.folder_path.mkdir(parents=True, exist_ok=True)
# Repository
repo_url = self.api.create_repo(repo_id=repo_id, private=private, repo_type=repo_type, exist_ok=True)
self.repo_id = repo_url.repo_id
self.repo_type = repo_type
self.revision = revision
self.token = token
# Keep track of already uploaded files
self.last_uploaded: Dict[Path, float] = {} # key is local path, value is timestamp
# Scheduler
if not every > 0:
raise ValueError(f"'every' must be a positive integer, not '{every}'.")
self.lock = Lock()
self.every = every
self.squash_history = squash_history
logger.info(f"Scheduled job to push '{self.folder_path}' to '{self.repo_id}' every {self.every} minutes.")
self._scheduler_thread = Thread(target=self._run_scheduler, daemon=True)
self._scheduler_thread.start()
atexit.register(self._push_to_hub)
self.__stopped = False
def stop(self) -> None:
"""Stop the scheduler.
A stopped scheduler cannot be restarted. Mostly for tests purposes.
"""
self.__stopped = True
def __enter__(self) -> "CommitScheduler":
return self
def __exit__(self, exc_type, exc_value, traceback) -> None:
# Upload last changes before exiting
self.trigger().result()
self.stop()
return
def _run_scheduler(self) -> None:
"""Dumb thread waiting between each scheduled push to Hub."""
while True:
self.last_future = self.trigger()
time.sleep(self.every * 60)
if self.__stopped:
break
def trigger(self) -> Future:
"""Trigger a `push_to_hub` and return a future.
This method is automatically called every `every` minutes. You can also call it manually to trigger a commit
immediately, without waiting for the next scheduled commit.
"""
return self.api.run_as_future(self._push_to_hub)
def _push_to_hub(self) -> Optional[CommitInfo]:
if self.__stopped: # If stopped, already scheduled commits are ignored
return None
logger.info("(Background) scheduled commit triggered.")
try:
value = self.push_to_hub()
if self.squash_history:
logger.info("(Background) squashing repo history.")
self.api.super_squash_history(repo_id=self.repo_id, repo_type=self.repo_type, branch=self.revision)
return value
except Exception as e:
logger.error(f"Error while pushing to Hub: {e}") # Depending on the setup, error might be silenced
raise
def push_to_hub(self) -> Optional[CommitInfo]:
"""
Push folder to the Hub and return the commit info.
<Tip warning={true}>
This method is not meant to be called directly. It is run in the background by the scheduler, respecting a
queue mechanism to avoid concurrent commits. Making a direct call to the method might lead to concurrency
issues.
</Tip>
The default behavior of `push_to_hub` is to assume an append-only folder. It lists all files in the folder and
uploads only changed files. If no changes are found, the method returns without committing anything. If you want
to change this behavior, you can inherit from [`CommitScheduler`] and override this method. This can be useful
for example to compress data together in a single file before committing. For more details and examples, check
out our [integration guide](https://huggingface.co/docs/huggingface_hub/main/en/guides/upload#scheduled-uploads).
"""
# Check files to upload (with lock)
with self.lock:
logger.debug("Listing files to upload for scheduled commit.")
# List files from folder (taken from `_prepare_upload_folder_additions`)
relpath_to_abspath = {
path.relative_to(self.folder_path).as_posix(): path
for path in sorted(self.folder_path.glob("**/*")) # sorted to be deterministic
if path.is_file()
}
prefix = f"{self.path_in_repo.strip('/')}/" if self.path_in_repo else ""
# Filter with pattern + filter out unchanged files + retrieve current file size
files_to_upload: List[_FileToUpload] = []
for relpath in filter_repo_objects(
relpath_to_abspath.keys(), allow_patterns=self.allow_patterns, ignore_patterns=self.ignore_patterns
):
local_path = relpath_to_abspath[relpath]
stat = local_path.stat()
if self.last_uploaded.get(local_path) is None or self.last_uploaded[local_path] != stat.st_mtime:
files_to_upload.append(
_FileToUpload(
local_path=local_path,
path_in_repo=prefix + relpath,
size_limit=stat.st_size,
last_modified=stat.st_mtime,
)
)
# Return if nothing to upload
if len(files_to_upload) == 0:
logger.debug("Dropping schedule commit: no changed file to upload.")
return None
# Convert `_FileToUpload` as `CommitOperationAdd` (=> compute file shas + limit to file size)
logger.debug("Removing unchanged files since previous scheduled commit.")
add_operations = [
CommitOperationAdd(
# Cap the file to its current size, even if the user append data to it while a scheduled commit is happening
path_or_fileobj=PartialFileIO(file_to_upload.local_path, size_limit=file_to_upload.size_limit),
path_in_repo=file_to_upload.path_in_repo,
)
for file_to_upload in files_to_upload
]
# Upload files (append mode expected - no need for lock)
logger.debug("Uploading files for scheduled commit.")
commit_info = self.api.create_commit(
repo_id=self.repo_id,
repo_type=self.repo_type,
operations=add_operations,
commit_message="Scheduled Commit",
revision=self.revision,
)
# Successful commit: keep track of the latest "last_modified" for each file
for file in files_to_upload:
self.last_uploaded[file.local_path] = file.last_modified
return commit_info
class PartialFileIO(BytesIO):
"""A file-like object that reads only the first part of a file.
Useful to upload a file to the Hub when the user might still be appending data to it. Only the first part of the
file is uploaded (i.e. the part that was available when the filesystem was first scanned).
In practice, only used internally by the CommitScheduler to regularly push a folder to the Hub with minimal
disturbance for the user. The object is passed to `CommitOperationAdd`.
Only supports `read`, `tell` and `seek` methods.
Args:
file_path (`str` or `Path`):
Path to the file to read.
size_limit (`int`):
The maximum number of bytes to read from the file. If the file is larger than this, only the first part
will be read (and uploaded).
"""
def __init__(self, file_path: Union[str, Path], size_limit: int) -> None:
self._file_path = Path(file_path)
self._file = self._file_path.open("rb")
self._size_limit = min(size_limit, os.fstat(self._file.fileno()).st_size)
def __del__(self) -> None:
self._file.close()
return super().__del__()
def __repr__(self) -> str:
return f"<PartialFileIO file_path={self._file_path} size_limit={self._size_limit}>"
def __len__(self) -> int:
return self._size_limit
def __getattribute__(self, name: str):
if name.startswith("_") or name in ("read", "tell", "seek"): # only 3 public methods supported
return super().__getattribute__(name)
raise NotImplementedError(f"PartialFileIO does not support '{name}'.")
def tell(self) -> int:
"""Return the current file position."""
return self._file.tell()
def seek(self, __offset: int, __whence: int = SEEK_SET) -> int:
"""Change the stream position to the given offset.
Behavior is the same as a regular file, except that the position is capped to the size limit.
"""
if __whence == SEEK_END:
# SEEK_END => set from the truncated end
__offset = len(self) + __offset
__whence = SEEK_SET
pos = self._file.seek(__offset, __whence)
if pos > self._size_limit:
return self._file.seek(self._size_limit)
return pos
def read(self, __size: Optional[int] = -1) -> bytes:
"""Read at most `__size` bytes from the file.
Behavior is the same as a regular file, except that it is capped to the size limit.
"""
current = self._file.tell()
if __size is None or __size < 0:
# Read until file limit
truncated_size = self._size_limit - current
else:
# Read until file limit or __size
truncated_size = min(__size, self._size_limit - current)
return self._file.read(truncated_size)
kronos-venv/lib/python3.12/site-packages/huggingface_hub/_inference_endpoints.py
+413
View File
@@ -0,0 +1,413 @@
import time
from dataclasses import dataclass, field
from datetime import datetime
from enum import Enum
from typing import TYPE_CHECKING, Dict, Optional, Union
from huggingface_hub.errors import InferenceEndpointError, InferenceEndpointTimeoutError
from .utils import get_session, logging, parse_datetime
if TYPE_CHECKING:
from .hf_api import HfApi
from .inference._client import InferenceClient
from .inference._generated._async_client import AsyncInferenceClient
logger = logging.get_logger(__name__)
class InferenceEndpointStatus(str, Enum):
PENDING = "pending"
INITIALIZING = "initializing"
UPDATING = "updating"
UPDATE_FAILED = "updateFailed"
RUNNING = "running"
PAUSED = "paused"
FAILED = "failed"
SCALED_TO_ZERO = "scaledToZero"
class InferenceEndpointType(str, Enum):
PUBlIC = "public"
PROTECTED = "protected"
PRIVATE = "private"
@dataclass
class InferenceEndpoint:
"""
Contains information about a deployed Inference Endpoint.
Args:
name (`str`):
The unique name of the Inference Endpoint.
namespace (`str`):
The namespace where the Inference Endpoint is located.
repository (`str`):
The name of the model repository deployed on this Inference Endpoint.
status ([`InferenceEndpointStatus`]):
The current status of the Inference Endpoint.
url (`str`, *optional*):
The URL of the Inference Endpoint, if available. Only a deployed Inference Endpoint will have a URL.
framework (`str`):
The machine learning framework used for the model.
revision (`str`):
The specific model revision deployed on the Inference Endpoint.
task (`str`):
The task associated with the deployed model.
created_at (`datetime.datetime`):
The timestamp when the Inference Endpoint was created.
updated_at (`datetime.datetime`):
The timestamp of the last update of the Inference Endpoint.
type ([`InferenceEndpointType`]):
The type of the Inference Endpoint (public, protected, private).
raw (`Dict`):
The raw dictionary data returned from the API.
token (`str` or `bool`, *optional*):
Authentication token for the Inference Endpoint, if set when requesting the API. Will default to the
locally saved token if not provided. Pass `token=False` if you don't want to send your token to the server.
Example:
```python
>>> from huggingface_hub import get_inference_endpoint
>>> endpoint = get_inference_endpoint("my-text-to-image")
>>> endpoint
InferenceEndpoint(name='my-text-to-image', ...)
# Get status
>>> endpoint.status
'running'
>>> endpoint.url
'https://my-text-to-image.region.vendor.endpoints.huggingface.cloud'
# Run inference
>>> endpoint.client.text_to_image(...)
# Pause endpoint to save $$$
>>> endpoint.pause()
# ...
# Resume and wait for deployment
>>> endpoint.resume()
>>> endpoint.wait()
>>> endpoint.client.text_to_image(...)
```
"""
# Field in __repr__
name: str = field(init=False)
namespace: str
repository: str = field(init=False)
status: InferenceEndpointStatus = field(init=False)
health_route: str = field(init=False)
url: Optional[str] = field(init=False)
# Other fields
framework: str = field(repr=False, init=False)
revision: str = field(repr=False, init=False)
task: str = field(repr=False, init=False)
created_at: datetime = field(repr=False, init=False)
updated_at: datetime = field(repr=False, init=False)
type: InferenceEndpointType = field(repr=False, init=False)
# Raw dict from the API
raw: Dict = field(repr=False)
# Internal fields
_token: Union[str, bool, None] = field(repr=False, compare=False)
_api: "HfApi" = field(repr=False, compare=False)
@classmethod
def from_raw(
cls, raw: Dict, namespace: str, token: Union[str, bool, None] = None, api: Optional["HfApi"] = None
) -> "InferenceEndpoint":
"""Initialize object from raw dictionary."""
if api is None:
from .hf_api import HfApi
api = HfApi()
if token is None:
token = api.token
# All other fields are populated in __post_init__
return cls(raw=raw, namespace=namespace, _token=token, _api=api)
def __post_init__(self) -> None:
"""Populate fields from raw dictionary."""
self._populate_from_raw()
@property
def client(self) -> "InferenceClient":
"""Returns a client to make predictions on this Inference Endpoint.
Returns:
[`InferenceClient`]: an inference client pointing to the deployed endpoint.
Raises:
[`InferenceEndpointError`]: If the Inference Endpoint is not yet deployed.
"""
if self.url is None:
raise InferenceEndpointError(
"Cannot create a client for this Inference Endpoint as it is not yet deployed. "
"Please wait for the Inference Endpoint to be deployed using `endpoint.wait()` and try again."
)
from .inference._client import InferenceClient
return InferenceClient(
model=self.url,
token=self._token, # type: ignore[arg-type] # boolean token shouldn't be possible. In practice it's ok.
)
@property
def async_client(self) -> "AsyncInferenceClient":
"""Returns a client to make predictions on this Inference Endpoint.
Returns:
[`AsyncInferenceClient`]: an asyncio-compatible inference client pointing to the deployed endpoint.
Raises:
[`InferenceEndpointError`]: If the Inference Endpoint is not yet deployed.
"""
if self.url is None:
raise InferenceEndpointError(
"Cannot create a client for this Inference Endpoint as it is not yet deployed. "
"Please wait for the Inference Endpoint to be deployed using `endpoint.wait()` and try again."
)
from .inference._generated._async_client import AsyncInferenceClient
return AsyncInferenceClient(
model=self.url,
token=self._token, # type: ignore[arg-type] # boolean token shouldn't be possible. In practice it's ok.
)
def wait(self, timeout: Optional[int] = None, refresh_every: int = 5) -> "InferenceEndpoint":
"""Wait for the Inference Endpoint to be deployed.
Information from the server will be fetched every 1s. If the Inference Endpoint is not deployed after `timeout`
seconds, a [`InferenceEndpointTimeoutError`] will be raised. The [`InferenceEndpoint`] will be mutated in place with the latest
data.
Args:
timeout (`int`, *optional*):
The maximum time to wait for the Inference Endpoint to be deployed, in seconds. If `None`, will wait
indefinitely.
refresh_every (`int`, *optional*):
The time to wait between each fetch of the Inference Endpoint status, in seconds. Defaults to 5s.
Returns:
[`InferenceEndpoint`]: the same Inference Endpoint, mutated in place with the latest data.
Raises:
[`InferenceEndpointError`]
If the Inference Endpoint ended up in a failed state.
[`InferenceEndpointTimeoutError`]
If the Inference Endpoint is not deployed after `timeout` seconds.
"""
if timeout is not None and timeout < 0:
raise ValueError("`timeout` cannot be negative.")
if refresh_every <= 0:
raise ValueError("`refresh_every` must be positive.")
start = time.time()
while True:
if self.status == InferenceEndpointStatus.FAILED:
raise InferenceEndpointError(
f"Inference Endpoint {self.name} failed to deploy. Please check the logs for more information."
)
if self.status == InferenceEndpointStatus.UPDATE_FAILED:
raise InferenceEndpointError(
f"Inference Endpoint {self.name} failed to update. Please check the logs for more information."
)
if self.status == InferenceEndpointStatus.RUNNING and self.url is not None:
# Verify the endpoint is actually reachable
_health_url = f"{self.url.rstrip('/')}/{self.health_route.lstrip('/')}"
response = get_session().get(_health_url, headers=self._api._build_hf_headers(token=self._token))
if response.status_code == 200:
logger.info("Inference Endpoint is ready to be used.")
return self
if timeout is not None:
if time.time() - start > timeout:
raise InferenceEndpointTimeoutError("Timeout while waiting for Inference Endpoint to be deployed.")
logger.info(f"Inference Endpoint is not deployed yet ({self.status}). Waiting {refresh_every}s...")
time.sleep(refresh_every)
self.fetch()
def fetch(self) -> "InferenceEndpoint":
"""Fetch latest information about the Inference Endpoint.
Returns:
[`InferenceEndpoint`]: the same Inference Endpoint, mutated in place with the latest data.
"""
obj = self._api.get_inference_endpoint(name=self.name, namespace=self.namespace, token=self._token) # type: ignore [arg-type]
self.raw = obj.raw
self._populate_from_raw()
return self
def update(
self,
*,
# Compute update
accelerator: Optional[str] = None,
instance_size: Optional[str] = None,
instance_type: Optional[str] = None,
min_replica: Optional[int] = None,
max_replica: Optional[int] = None,
scale_to_zero_timeout: Optional[int] = None,
# Model update
repository: Optional[str] = None,
framework: Optional[str] = None,
revision: Optional[str] = None,
task: Optional[str] = None,
custom_image: Optional[Dict] = None,
secrets: Optional[Dict[str, str]] = None,
) -> "InferenceEndpoint":
"""Update the Inference Endpoint.
This method allows the update of either the compute configuration, the deployed model, or both. All arguments are
optional but at least one must be provided.
This is an alias for [`HfApi.update_inference_endpoint`]. The current object is mutated in place with the
latest data from the server.
Args:
accelerator (`str`, *optional*):
The hardware accelerator to be used for inference (e.g. `"cpu"`).
instance_size (`str`, *optional*):
The size or type of the instance to be used for hosting the model (e.g. `"x4"`).
instance_type (`str`, *optional*):
The cloud instance type where the Inference Endpoint will be deployed (e.g. `"intel-icl"`).
min_replica (`int`, *optional*):
The minimum number of replicas (instances) to keep running for the Inference Endpoint.
max_replica (`int`, *optional*):
The maximum number of replicas (instances) to scale to for the Inference Endpoint.
scale_to_zero_timeout (`int`, *optional*):
The duration in minutes before an inactive endpoint is scaled to zero.
repository (`str`, *optional*):
The name of the model repository associated with the Inference Endpoint (e.g. `"gpt2"`).
framework (`str`, *optional*):
The machine learning framework used for the model (e.g. `"custom"`).
revision (`str`, *optional*):
The specific model revision to deploy on the Inference Endpoint (e.g. `"6c0e6080953db56375760c0471a8c5f2929baf11"`).
task (`str`, *optional*):
The task on which to deploy the model (e.g. `"text-classification"`).
custom_image (`Dict`, *optional*):
A custom Docker image to use for the Inference Endpoint. This is useful if you want to deploy an
Inference Endpoint running on the `text-generation-inference` (TGI) framework (see examples).
secrets (`Dict[str, str]`, *optional*):
Secret values to inject in the container environment.
Returns:
[`InferenceEndpoint`]: the same Inference Endpoint, mutated in place with the latest data.
"""
# Make API call
obj = self._api.update_inference_endpoint(
name=self.name,
namespace=self.namespace,
accelerator=accelerator,
instance_size=instance_size,
instance_type=instance_type,
min_replica=min_replica,
max_replica=max_replica,
scale_to_zero_timeout=scale_to_zero_timeout,
repository=repository,
framework=framework,
revision=revision,
task=task,
custom_image=custom_image,
secrets=secrets,
token=self._token, # type: ignore [arg-type]
)
# Mutate current object
self.raw = obj.raw
self._populate_from_raw()
return self
def pause(self) -> "InferenceEndpoint":
"""Pause the Inference Endpoint.
A paused Inference Endpoint will not be charged. It can be resumed at any time using [`InferenceEndpoint.resume`].
This is different than scaling the Inference Endpoint to zero with [`InferenceEndpoint.scale_to_zero`], which
would be automatically restarted when a request is made to it.
This is an alias for [`HfApi.pause_inference_endpoint`]. The current object is mutated in place with the
latest data from the server.
Returns:
[`InferenceEndpoint`]: the same Inference Endpoint, mutated in place with the latest data.
"""
obj = self._api.pause_inference_endpoint(name=self.name, namespace=self.namespace, token=self._token) # type: ignore [arg-type]
self.raw = obj.raw
self._populate_from_raw()
return self
def resume(self, running_ok: bool = True) -> "InferenceEndpoint":
"""Resume the Inference Endpoint.
This is an alias for [`HfApi.resume_inference_endpoint`]. The current object is mutated in place with the
latest data from the server.
Args:
running_ok (`bool`, *optional*):
If `True`, the method will not raise an error if the Inference Endpoint is already running. Defaults to
`True`.
Returns:
[`InferenceEndpoint`]: the same Inference Endpoint, mutated in place with the latest data.
"""
obj = self._api.resume_inference_endpoint(
name=self.name, namespace=self.namespace, running_ok=running_ok, token=self._token
) # type: ignore [arg-type]
self.raw = obj.raw
self._populate_from_raw()
return self
def scale_to_zero(self) -> "InferenceEndpoint":
"""Scale Inference Endpoint to zero.
An Inference Endpoint scaled to zero will not be charged. It will be resume on the next request to it, with a
cold start delay. This is different than pausing the Inference Endpoint with [`InferenceEndpoint.pause`], which
would require a manual resume with [`InferenceEndpoint.resume`].
This is an alias for [`HfApi.scale_to_zero_inference_endpoint`]. The current object is mutated in place with the
latest data from the server.
Returns:
[`InferenceEndpoint`]: the same Inference Endpoint, mutated in place with the latest data.
"""
obj = self._api.scale_to_zero_inference_endpoint(name=self.name, namespace=self.namespace, token=self._token) # type: ignore [arg-type]
self.raw = obj.raw
self._populate_from_raw()
return self
def delete(self) -> None:
"""Delete the Inference Endpoint.
This operation is not reversible. If you don't want to be charged for an Inference Endpoint, it is preferable
to pause it with [`InferenceEndpoint.pause`] or scale it to zero with [`InferenceEndpoint.scale_to_zero`].
This is an alias for [`HfApi.delete_inference_endpoint`].
"""
self._api.delete_inference_endpoint(name=self.name, namespace=self.namespace, token=self._token) # type: ignore [arg-type]
def _populate_from_raw(self) -> None:
"""Populate fields from raw dictionary.
Called in __post_init__ + each time the Inference Endpoint is updated.
"""
# Repr fields
self.name = self.raw["name"]
self.repository = self.raw["model"]["repository"]
self.status = self.raw["status"]["state"]
self.url = self.raw["status"].get("url")
self.health_route = self.raw["healthRoute"]
# Other fields
self.framework = self.raw["model"]["framework"]
self.revision = self.raw["model"]["revision"]
self.task = self.raw["model"]["task"]
self.created_at = parse_datetime(self.raw["status"]["createdAt"])
self.updated_at = parse_datetime(self.raw["status"]["updatedAt"])
self.type = self.raw["type"]
kronos-venv/lib/python3.12/site-packages/huggingface_hub/_local_folder.py
+441
View File
@@ -0,0 +1,441 @@
# coding=utf-8
# Copyright 2024-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 to handle the `../.cache/huggingface` folder in local directories.
First discussed in https://github.com/huggingface/huggingface_hub/issues/1738 to store
download metadata when downloading files from the hub to a local directory (without
using the cache).
./.cache/huggingface folder structure:
[4.0K] data
├── [4.0K] .cache
│ └── [4.0K] huggingface
│ └── [4.0K] download
│ ├── [ 16] file.parquet.metadata
│ ├── [ 16] file.txt.metadata
│ └── [4.0K] folder
│ └── [ 16] file.parquet.metadata
│
├── [6.5G] file.parquet
├── [1.5K] file.txt
└── [4.0K] folder
└── [ 16] file.parquet
Download metadata file structure:
```
# file.txt.metadata
11c5a3d5811f50298f278a704980280950aedb10
a16a55fda99d2f2e7b69cce5cf93ff4ad3049930
1712656091.123
# file.parquet.metadata
11c5a3d5811f50298f278a704980280950aedb10
7c5d3f4b8b76583b422fcb9189ad6c89d5d97a094541ce8932dce3ecabde1421
1712656091.123
}
```
"""
import base64
import hashlib
import logging
import os
import time
from dataclasses import dataclass
from pathlib import Path
from typing import Optional
from .utils import WeakFileLock
logger = logging.getLogger(__name__)
@dataclass
class LocalDownloadFilePaths:
"""
Paths to the files related to a download process in a local dir.
Returned by [`get_local_download_paths`].
Attributes:
file_path (`Path`):
Path where the file will be saved.
lock_path (`Path`):
Path to the lock file used to ensure atomicity when reading/writing metadata.
metadata_path (`Path`):
Path to the metadata file.
"""
file_path: Path
lock_path: Path
metadata_path: Path
def incomplete_path(self, etag: str) -> Path:
"""Return the path where a file will be temporarily downloaded before being moved to `file_path`."""
return self.metadata_path.parent / f"{_short_hash(self.metadata_path.name)}.{etag}.incomplete"
@dataclass(frozen=True)
class LocalUploadFilePaths:
"""
Paths to the files related to an upload process in a local dir.
Returned by [`get_local_upload_paths`].
Attributes:
path_in_repo (`str`):
Path of the file in the repo.
file_path (`Path`):
Path where the file will be saved.
lock_path (`Path`):
Path to the lock file used to ensure atomicity when reading/writing metadata.
metadata_path (`Path`):
Path to the metadata file.
"""
path_in_repo: str
file_path: Path
lock_path: Path
metadata_path: Path
@dataclass
class LocalDownloadFileMetadata:
"""
Metadata about a file in the local directory related to a download process.
Attributes:
filename (`str`):
Path of the file in the repo.
commit_hash (`str`):
Commit hash of the file in the repo.
etag (`str`):
ETag of the file in the repo. Used to check if the file has changed.
For LFS files, this is the sha256 of the file. For regular files, it corresponds to the git hash.
timestamp (`int`):
Unix timestamp of when the metadata was saved i.e. when the metadata was accurate.
"""
filename: str
commit_hash: str
etag: str
timestamp: float
@dataclass
class LocalUploadFileMetadata:
"""
Metadata about a file in the local directory related to an upload process.
"""
size: int
# Default values correspond to "we don't know yet"
timestamp: Optional[float] = None
should_ignore: Optional[bool] = None
sha256: Optional[str] = None
upload_mode: Optional[str] = None
remote_oid: Optional[str] = None
is_uploaded: bool = False
is_committed: bool = False
def save(self, paths: LocalUploadFilePaths) -> None:
"""Save the metadata to disk."""
with WeakFileLock(paths.lock_path):
with paths.metadata_path.open("w") as f:
new_timestamp = time.time()
f.write(str(new_timestamp) + "\n")
f.write(str(self.size)) # never None
f.write("\n")
if self.should_ignore is not None:
f.write(str(int(self.should_ignore)))
f.write("\n")
if self.sha256 is not None:
f.write(self.sha256)
f.write("\n")
if self.upload_mode is not None:
f.write(self.upload_mode)
f.write("\n")
if self.remote_oid is not None:
f.write(self.remote_oid)
f.write("\n")
f.write(str(int(self.is_uploaded)) + "\n")
f.write(str(int(self.is_committed)) + "\n")
self.timestamp = new_timestamp
def get_local_download_paths(local_dir: Path, filename: str) -> LocalDownloadFilePaths:
"""Compute paths to the files related to a download process.
Folders containing the paths are all guaranteed to exist.
Args:
local_dir (`Path`):
Path to the local directory in which files are downloaded.
filename (`str`):
Path of the file in the repo.
Return:
[`LocalDownloadFilePaths`]: the paths to the files (file_path, lock_path, metadata_path, incomplete_path).
"""
# filename is the path in the Hub repository (separated by '/')
# make sure to have a cross platform transcription
sanitized_filename = os.path.join(*filename.split("/"))
if os.name == "nt":
if sanitized_filename.startswith("..\\") or "\\..\\" in sanitized_filename:
raise ValueError(
f"Invalid filename: cannot handle filename '{sanitized_filename}' on Windows. Please ask the repository"
" owner to rename this file."
)
file_path = local_dir / sanitized_filename
metadata_path = _huggingface_dir(local_dir) / "download" / f"{sanitized_filename}.metadata"
lock_path = metadata_path.with_suffix(".lock")
# Some Windows versions do not allow for paths longer than 255 characters.
# In this case, we must specify it as an extended path by using the "\\?\" prefix
if os.name == "nt":
if not str(local_dir).startswith("\\\\?\\") and len(os.path.abspath(lock_path)) > 255:
file_path = Path("\\\\?\\" + os.path.abspath(file_path))
lock_path = Path("\\\\?\\" + os.path.abspath(lock_path))
metadata_path = Path("\\\\?\\" + os.path.abspath(metadata_path))
file_path.parent.mkdir(parents=True, exist_ok=True)
metadata_path.parent.mkdir(parents=True, exist_ok=True)
return LocalDownloadFilePaths(file_path=file_path, lock_path=lock_path, metadata_path=metadata_path)
def get_local_upload_paths(local_dir: Path, filename: str) -> LocalUploadFilePaths:
"""Compute paths to the files related to an upload process.
Folders containing the paths are all guaranteed to exist.
Args:
local_dir (`Path`):
Path to the local directory that is uploaded.
filename (`str`):
Path of the file in the repo.
Return:
[`LocalUploadFilePaths`]: the paths to the files (file_path, lock_path, metadata_path).
"""
# filename is the path in the Hub repository (separated by '/')
# make sure to have a cross platform transcription
sanitized_filename = os.path.join(*filename.split("/"))
if os.name == "nt":
if sanitized_filename.startswith("..\\") or "\\..\\" in sanitized_filename:
raise ValueError(
f"Invalid filename: cannot handle filename '{sanitized_filename}' on Windows. Please ask the repository"
" owner to rename this file."
)
file_path = local_dir / sanitized_filename
metadata_path = _huggingface_dir(local_dir) / "upload" / f"{sanitized_filename}.metadata"
lock_path = metadata_path.with_suffix(".lock")
# Some Windows versions do not allow for paths longer than 255 characters.
# In this case, we must specify it as an extended path by using the "\\?\" prefix
if os.name == "nt":
if not str(local_dir).startswith("\\\\?\\") and len(os.path.abspath(lock_path)) > 255:
file_path = Path("\\\\?\\" + os.path.abspath(file_path))
lock_path = Path("\\\\?\\" + os.path.abspath(lock_path))
metadata_path = Path("\\\\?\\" + os.path.abspath(metadata_path))
file_path.parent.mkdir(parents=True, exist_ok=True)
metadata_path.parent.mkdir(parents=True, exist_ok=True)
return LocalUploadFilePaths(
path_in_repo=filename, file_path=file_path, lock_path=lock_path, metadata_path=metadata_path
)
def read_download_metadata(local_dir: Path, filename: str) -> Optional[LocalDownloadFileMetadata]:
"""Read metadata about a file in the local directory related to a download process.
Args:
local_dir (`Path`):
Path to the local directory in which files are downloaded.
filename (`str`):
Path of the file in the repo.
Return:
`[LocalDownloadFileMetadata]` or `None`: the metadata if it exists, `None` otherwise.
"""
paths = get_local_download_paths(local_dir, filename)
with WeakFileLock(paths.lock_path):
if paths.metadata_path.exists():
try:
with paths.metadata_path.open() as f:
commit_hash = f.readline().strip()
etag = f.readline().strip()
timestamp = float(f.readline().strip())
metadata = LocalDownloadFileMetadata(
filename=filename,
commit_hash=commit_hash,
etag=etag,
timestamp=timestamp,
)
except Exception as e:
# remove the metadata file if it is corrupted / not the right format
logger.warning(
f"Invalid metadata file {paths.metadata_path}: {e}. Removing it from disk and continue."
)
try:
paths.metadata_path.unlink()
except Exception as e:
logger.warning(f"Could not remove corrupted metadata file {paths.metadata_path}: {e}")
try:
# check if the file exists and hasn't been modified since the metadata was saved
stat = paths.file_path.stat()
if (
stat.st_mtime - 1 <= metadata.timestamp
): # allow 1s difference as stat.st_mtime might not be precise
return metadata
logger.info(f"Ignored metadata for '{filename}' (outdated). Will re-compute hash.")
except FileNotFoundError:
# file does not exist => metadata is outdated
return None
return None
def read_upload_metadata(local_dir: Path, filename: str) -> LocalUploadFileMetadata:
"""Read metadata about a file in the local directory related to an upload process.
TODO: factorize logic with `read_download_metadata`.
Args:
local_dir (`Path`):
Path to the local directory in which files are downloaded.
filename (`str`):
Path of the file in the repo.
Return:
`[LocalUploadFileMetadata]` or `None`: the metadata if it exists, `None` otherwise.
"""
paths = get_local_upload_paths(local_dir, filename)
with WeakFileLock(paths.lock_path):
if paths.metadata_path.exists():
try:
with paths.metadata_path.open() as f:
timestamp = float(f.readline().strip())
size = int(f.readline().strip()) # never None
_should_ignore = f.readline().strip()
should_ignore = None if _should_ignore == "" else bool(int(_should_ignore))
_sha256 = f.readline().strip()
sha256 = None if _sha256 == "" else _sha256
_upload_mode = f.readline().strip()
upload_mode = None if _upload_mode == "" else _upload_mode
if upload_mode not in (None, "regular", "lfs"):
raise ValueError(f"Invalid upload mode in metadata {paths.path_in_repo}: {upload_mode}")
_remote_oid = f.readline().strip()
remote_oid = None if _remote_oid == "" else _remote_oid
is_uploaded = bool(int(f.readline().strip()))
is_committed = bool(int(f.readline().strip()))
metadata = LocalUploadFileMetadata(
timestamp=timestamp,
size=size,
should_ignore=should_ignore,
sha256=sha256,
upload_mode=upload_mode,
remote_oid=remote_oid,
is_uploaded=is_uploaded,
is_committed=is_committed,
)
except Exception as e:
# remove the metadata file if it is corrupted / not the right format
logger.warning(
f"Invalid metadata file {paths.metadata_path}: {e}. Removing it from disk and continue."
)
try:
paths.metadata_path.unlink()
except Exception as e:
logger.warning(f"Could not remove corrupted metadata file {paths.metadata_path}: {e}")
# TODO: can we do better?
if (
metadata.timestamp is not None
and metadata.is_uploaded # file was uploaded
and not metadata.is_committed # but not committed
and time.time() - metadata.timestamp > 20 * 3600 # and it's been more than 20 hours
): # => we consider it as garbage-collected by S3
metadata.is_uploaded = False
# check if the file exists and hasn't been modified since the metadata was saved
try:
if metadata.timestamp is not None and paths.file_path.stat().st_mtime <= metadata.timestamp:
return metadata
logger.info(f"Ignored metadata for '{filename}' (outdated). Will re-compute hash.")
except FileNotFoundError:
# file does not exist => metadata is outdated
pass
# empty metadata => we don't know anything expect its size
return LocalUploadFileMetadata(size=paths.file_path.stat().st_size)
def write_download_metadata(local_dir: Path, filename: str, commit_hash: str, etag: str) -> None:
"""Write metadata about a file in the local directory related to a download process.
Args:
local_dir (`Path`):
Path to the local directory in which files are downloaded.
"""
paths = get_local_download_paths(local_dir, filename)
with WeakFileLock(paths.lock_path):
with paths.metadata_path.open("w") as f:
f.write(f"{commit_hash}\n{etag}\n{time.time()}\n")
def _huggingface_dir(local_dir: Path) -> Path:
"""Return the path to the `.cache/huggingface` directory in a local directory."""
# Wrap in lru_cache to avoid overwriting the .gitignore file if called multiple times
path = local_dir / ".cache" / "huggingface"
path.mkdir(exist_ok=True, parents=True)
# Create a .gitignore file in the .cache/huggingface directory if it doesn't exist
# Should be thread-safe enough like this.
gitignore = path / ".gitignore"
gitignore_lock = path / ".gitignore.lock"
if not gitignore.exists():
try:
with WeakFileLock(gitignore_lock, timeout=0.1):
gitignore.write_text("*")
except IndexError:
pass
except OSError: # TimeoutError, FileNotFoundError, PermissionError, etc.
pass
try:
gitignore_lock.unlink()
except OSError:
pass
return path
def _short_hash(filename: str) -> str:
return base64.urlsafe_b64encode(hashlib.sha1(filename.encode()).digest()).decode()
kronos-venv/lib/python3.12/site-packages/huggingface_hub/_login.py
+520
View File
@@ -0,0 +1,520 @@
# Copyright 2020 The HuggingFace Team. All rights reserved.
#
# 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 methods to log in to the Hub."""
import os
import subprocess
from getpass import getpass
from pathlib import Path
from typing import Optional
from . import constants
from .commands._cli_utils import ANSI
from .utils import (
capture_output,
get_token,
is_google_colab,
is_notebook,
list_credential_helpers,
logging,
run_subprocess,
set_git_credential,
unset_git_credential,
)
from .utils._auth import (
_get_token_by_name,
_get_token_from_environment,
_get_token_from_file,
_get_token_from_google_colab,
_save_stored_tokens,
_save_token,
get_stored_tokens,
)
from .utils._deprecation import _deprecate_arguments, _deprecate_positional_args
logger = logging.get_logger(__name__)
_HF_LOGO_ASCII = """
_| _| _| _| _|_|_| _|_|_| _|_|_| _| _| _|_|_| _|_|_|_| _|_| _|_|_| _|_|_|_|
_| _| _| _| _| _| _| _|_| _| _| _| _| _| _| _|
_|_|_|_| _| _| _| _|_| _| _|_| _| _| _| _| _| _|_| _|_|_| _|_|_|_| _| _|_|_|
_| _| _| _| _| _| _| _| _| _| _|_| _| _| _| _| _| _| _|
_| _| _|_| _|_|_| _|_|_| _|_|_| _| _| _|_|_| _| _| _| _|_|_| _|_|_|_|
"""
@_deprecate_arguments(
version="1.0",
deprecated_args="write_permission",
custom_message="Fine-grained tokens added complexity to the permissions, making it irrelevant to check if a token has 'write' access.",
)
@_deprecate_positional_args(version="1.0")
def login(
token: Optional[str] = None,
*,
add_to_git_credential: bool = False,
new_session: bool = True,
write_permission: bool = False,
) -> None:
"""Login the machine to access the Hub.
The `token` is persisted in cache and set as a git credential. Once done, the machine
is logged in and the access token will be available across all `huggingface_hub`
components. If `token` is not provided, it will be prompted to the user either with
a widget (in a notebook) or via the terminal.
To log in from outside of a script, one can also use `huggingface-cli login` which is
a cli command that wraps [`login`].
<Tip>
[`login`] is a drop-in replacement method for [`notebook_login`] as it wraps and
extends its capabilities.
</Tip>
<Tip>
When the token is not passed, [`login`] will automatically detect if the script runs
in a notebook or not. However, this detection might not be accurate due to the
variety of notebooks that exists nowadays. If that is the case, you can always force
the UI by using [`notebook_login`] or [`interpreter_login`].
</Tip>
Args:
token (`str`, *optional*):
User access token to generate from https://huggingface.co/settings/token.
add_to_git_credential (`bool`, defaults to `False`):
If `True`, token will be set as git credential. If no git credential helper
is configured, a warning will be displayed to the user. If `token` is `None`,
the value of `add_to_git_credential` is ignored and will be prompted again
to the end user.
new_session (`bool`, defaults to `True`):
If `True`, will request a token even if one is already saved on the machine.
write_permission (`bool`):
Ignored and deprecated argument.
Raises:
[`ValueError`](https://docs.python.org/3/library/exceptions.html#ValueError)
If an organization token is passed. Only personal account tokens are valid
to log in.
[`ValueError`](https://docs.python.org/3/library/exceptions.html#ValueError)
If token is invalid.
[`ImportError`](https://docs.python.org/3/library/exceptions.html#ImportError)
If running in a notebook but `ipywidgets` is not installed.
"""
if token is not None:
if not add_to_git_credential:
logger.info(
"The token has not been saved to the git credentials helper. Pass "
"`add_to_git_credential=True` in this function directly or "
"`--add-to-git-credential` if using via `huggingface-cli` if "
"you want to set the git credential as well."
)
_login(token, add_to_git_credential=add_to_git_credential)
elif is_notebook():
notebook_login(new_session=new_session)
else:
interpreter_login(new_session=new_session)
def logout(token_name: Optional[str] = None) -> None:
"""Logout the machine from the Hub.
Token is deleted from the machine and removed from git credential.
Args:
token_name (`str`, *optional*):
Name of the access token to logout from. If `None`, will logout from all saved access tokens.
Raises:
[`ValueError`](https://docs.python.org/3/library/exceptions.html#ValueError):
If the access token name is not found.
"""
if get_token() is None and not get_stored_tokens(): # No active token and no saved access tokens
logger.warning("Not logged in!")
return
if not token_name:
# Delete all saved access tokens and token
for file_path in (constants.HF_TOKEN_PATH, constants.HF_STORED_TOKENS_PATH):
try:
Path(file_path).unlink()
except FileNotFoundError:
pass
logger.info("Successfully logged out from all access tokens.")
else:
_logout_from_token(token_name)
logger.info(f"Successfully logged out from access token: {token_name}.")
unset_git_credential()
# Check if still logged in
if _get_token_from_google_colab() is not None:
raise EnvironmentError(
"You are automatically logged in using a Google Colab secret.\n"
"To log out, you must unset the `HF_TOKEN` secret in your Colab settings."
)
if _get_token_from_environment() is not None:
raise EnvironmentError(
"Token has been deleted from your machine but you are still logged in.\n"
"To log out, you must clear out both `HF_TOKEN` and `HUGGING_FACE_HUB_TOKEN` environment variables."
)
def auth_switch(token_name: str, add_to_git_credential: bool = False) -> None:
"""Switch to a different access token.
Args:
token_name (`str`):
Name of the access token to switch to.
add_to_git_credential (`bool`, defaults to `False`):
If `True`, token will be set as git credential. If no git credential helper
is configured, a warning will be displayed to the user. If `token` is `None`,
the value of `add_to_git_credential` is ignored and will be prompted again
to the end user.
Raises:
[`ValueError`](https://docs.python.org/3/library/exceptions.html#ValueError):
If the access token name is not found.
"""
token = _get_token_by_name(token_name)
if not token:
raise ValueError(f"Access token {token_name} not found in {constants.HF_STORED_TOKENS_PATH}")
# Write token to HF_TOKEN_PATH
_set_active_token(token_name, add_to_git_credential)
logger.info(f"The current active token is: {token_name}")
token_from_environment = _get_token_from_environment()
if token_from_environment is not None and token_from_environment != token:
logger.warning(
"The environment variable `HF_TOKEN` is set and will override the access token you've just switched to."
)
def auth_list() -> None:
"""List all stored access tokens."""
tokens = get_stored_tokens()
if not tokens:
logger.info("No access tokens found.")
return
# Find current token
current_token = get_token()
current_token_name = None
for token_name in tokens:
if tokens.get(token_name) == current_token:
current_token_name = token_name
# Print header
max_offset = max(len("token"), max(len(token) for token in tokens)) + 2
print(f" {{:<{max_offset}}}| {{:<15}}".format("name", "token"))
print("-" * (max_offset + 2) + "|" + "-" * 15)
# Print saved access tokens
for token_name in tokens:
token = tokens.get(token_name, "<not set>")
masked_token = f"{token[:3]}****{token[-4:]}" if token != "<not set>" else token
is_current = "*" if token == current_token else " "
print(f"{is_current} {{:<{max_offset}}}| {{:<15}}".format(token_name, masked_token))
if _get_token_from_environment():
logger.warning(
"\nNote: Environment variable `HF_TOKEN` is set and is the current active token independently from the stored tokens listed above."
)
elif current_token_name is None:
logger.warning(
"\nNote: No active token is set and no environment variable `HF_TOKEN` is found. Use `huggingface-cli login` to log in."
)
###
# Interpreter-based login (text)
###
@_deprecate_arguments(
version="1.0",
deprecated_args="write_permission",
custom_message="Fine-grained tokens added complexity to the permissions, making it irrelevant to check if a token has 'write' access.",
)
@_deprecate_positional_args(version="1.0")
def interpreter_login(*, new_session: bool = True, write_permission: bool = False) -> None:
"""
Displays a prompt to log in to the HF website and store the token.
This is equivalent to [`login`] without passing a token when not run in a notebook.
[`interpreter_login`] is useful if you want to force the use of the terminal prompt
instead of a notebook widget.
For more details, see [`login`].
Args:
new_session (`bool`, defaults to `True`):
If `True`, will request a token even if one is already saved on the machine.
write_permission (`bool`):
Ignored and deprecated argument.
"""
if not new_session and get_token() is not None:
logger.info("User is already logged in.")
return
from .commands.delete_cache import _ask_for_confirmation_no_tui
print(_HF_LOGO_ASCII)
if get_token() is not None:
logger.info(
" A token is already saved on your machine. Run `huggingface-cli"
" whoami` to get more information or `huggingface-cli logout` if you want"
" to log out."
)
logger.info(" Setting a new token will erase the existing one.")
logger.info(
" To log in, `huggingface_hub` requires a token generated from https://huggingface.co/settings/tokens ."
)
if os.name == "nt":
logger.info("Token can be pasted using 'Right-Click'.")
token = getpass("Enter your token (input will not be visible): ")
add_to_git_credential = _ask_for_confirmation_no_tui("Add token as git credential?")
_login(token=token, add_to_git_credential=add_to_git_credential)
###
# Notebook-based login (widget)
###
NOTEBOOK_LOGIN_PASSWORD_HTML = """<center> <img
src=https://huggingface.co/front/assets/huggingface_logo-noborder.svg
alt='Hugging Face'> <br> Immediately click login after typing your password or
it might be stored in plain text in this notebook file. </center>"""
NOTEBOOK_LOGIN_TOKEN_HTML_START = """<center> <img
src=https://huggingface.co/front/assets/huggingface_logo-noborder.svg
alt='Hugging Face'> <br> Copy a token from <a
href="https://huggingface.co/settings/tokens" target="_blank">your Hugging Face
tokens page</a> and paste it below. <br> Immediately click login after copying
your token or it might be stored in plain text in this notebook file. </center>"""
NOTEBOOK_LOGIN_TOKEN_HTML_END = """
<b>Pro Tip:</b> If you don't already have one, you can create a dedicated
'notebooks' token with 'write' access, that you can then easily reuse for all
notebooks. </center>"""
@_deprecate_arguments(
version="1.0",
deprecated_args="write_permission",
custom_message="Fine-grained tokens added complexity to the permissions, making it irrelevant to check if a token has 'write' access.",
)
@_deprecate_positional_args(version="1.0")
def notebook_login(*, new_session: bool = True, write_permission: bool = False) -> None:
"""
Displays a widget to log in to the HF website and store the token.
This is equivalent to [`login`] without passing a token when run in a notebook.
[`notebook_login`] is useful if you want to force the use of the notebook widget
instead of a prompt in the terminal.
For more details, see [`login`].
Args:
new_session (`bool`, defaults to `True`):
If `True`, will request a token even if one is already saved on the machine.
write_permission (`bool`):
Ignored and deprecated argument.
"""
try:
import ipywidgets.widgets as widgets # type: ignore
from IPython.display import display # type: ignore
except ImportError:
raise ImportError(
"The `notebook_login` function can only be used in a notebook (Jupyter or"
" Colab) and you need the `ipywidgets` module: `pip install ipywidgets`."
)
if not new_session and get_token() is not None:
logger.info("User is already logged in.")
return
box_layout = widgets.Layout(display="flex", flex_flow="column", align_items="center", width="50%")
token_widget = widgets.Password(description="Token:")
git_checkbox_widget = widgets.Checkbox(value=True, description="Add token as git credential?")
token_finish_button = widgets.Button(description="Login")
login_token_widget = widgets.VBox(
[
widgets.HTML(NOTEBOOK_LOGIN_TOKEN_HTML_START),
token_widget,
git_checkbox_widget,
token_finish_button,
widgets.HTML(NOTEBOOK_LOGIN_TOKEN_HTML_END),
],
layout=box_layout,
)
display(login_token_widget)
# On click events
def login_token_event(t):
"""Event handler for the login button."""
token = token_widget.value
add_to_git_credential = git_checkbox_widget.value
# Erase token and clear value to make sure it's not saved in the notebook.
token_widget.value = ""
# Hide inputs
login_token_widget.children = [widgets.Label("Connecting...")]
try:
with capture_output() as captured:
_login(token, add_to_git_credential=add_to_git_credential)
message = captured.getvalue()
except Exception as error:
message = str(error)
# Print result (success message or error)
login_token_widget.children = [widgets.Label(line) for line in message.split("\n") if line.strip()]
token_finish_button.on_click(login_token_event)
###
# Login private helpers
###
def _login(
token: str,
add_to_git_credential: bool,
) -> None:
from .hf_api import whoami # avoid circular import
if token.startswith("api_org"):
raise ValueError("You must use your personal account token, not an organization token.")
token_info = whoami(token)
permission = token_info["auth"]["accessToken"]["role"]
logger.info(f"Token is valid (permission: {permission}).")
token_name = token_info["auth"]["accessToken"]["displayName"]
# Store token locally
_save_token(token=token, token_name=token_name)
# Set active token
_set_active_token(token_name=token_name, add_to_git_credential=add_to_git_credential)
logger.info("Login successful.")
if _get_token_from_environment():
logger.warning(
"Note: Environment variable`HF_TOKEN` is set and is the current active token independently from the token you've just configured."
)
else:
logger.info(f"The current active token is: `{token_name}`")
def _logout_from_token(token_name: str) -> None:
"""Logout from a specific access token.
Args:
token_name (`str`):
The name of the access token to logout from.
Raises:
[`ValueError`](https://docs.python.org/3/library/exceptions.html#ValueError):
If the access token name is not found.
"""
stored_tokens = get_stored_tokens()
# If there is no access tokens saved or the access token name is not found, do nothing
if not stored_tokens or token_name not in stored_tokens:
return
token = stored_tokens.pop(token_name)
_save_stored_tokens(stored_tokens)
if token == _get_token_from_file():
logger.warning(f"Active token '{token_name}' has been deleted.")
Path(constants.HF_TOKEN_PATH).unlink(missing_ok=True)
def _set_active_token(
token_name: str,
add_to_git_credential: bool,
) -> None:
"""Set the active access token.
Args:
token_name (`str`):
The name of the token to set as active.
"""
token = _get_token_by_name(token_name)
if not token:
raise ValueError(f"Token {token_name} not found in {constants.HF_STORED_TOKENS_PATH}")
if add_to_git_credential:
if _is_git_credential_helper_configured():
set_git_credential(token)
logger.info(
"Your token has been saved in your configured git credential helpers"
+ f" ({','.join(list_credential_helpers())})."
)
else:
logger.warning("Token has not been saved to git credential helper.")
# Write token to HF_TOKEN_PATH
path = Path(constants.HF_TOKEN_PATH)
path.parent.mkdir(parents=True, exist_ok=True)
path.write_text(token)
logger.info(f"Your token has been saved to {constants.HF_TOKEN_PATH}")
def _is_git_credential_helper_configured() -> bool:
"""Check if a git credential helper is configured.
Warns user if not the case (except for Google Colab where "store" is set by default
by `huggingface_hub`).
"""
helpers = list_credential_helpers()
if len(helpers) > 0:
return True # Do not warn: at least 1 helper is set
# Only in Google Colab to avoid the warning message
# See https://github.com/huggingface/huggingface_hub/issues/1043#issuecomment-1247010710
if is_google_colab():
_set_store_as_git_credential_helper_globally()
return True # Do not warn: "store" is used by default in Google Colab
# Otherwise, warn user
print(
ANSI.red(
"Cannot authenticate through git-credential as no helper is defined on your"
" machine.\nYou might have to re-authenticate when pushing to the Hugging"
" Face Hub.\nRun the following command in your terminal in case you want to"
" set the 'store' credential helper as default.\n\ngit config --global"
" credential.helper store\n\nRead"
" https://git-scm.com/book/en/v2/Git-Tools-Credential-Storage for more"
" details."
)
)
return False
def _set_store_as_git_credential_helper_globally() -> None:
"""Set globally the credential.helper to `store`.
To be used only in Google Colab as we assume the user doesn't care about the git
credential config. It is the only particular case where we don't want to display the
warning message in [`notebook_login()`].
Related:
- https://github.com/huggingface/huggingface_hub/issues/1043
- https://github.com/huggingface/huggingface_hub/issues/1051
- https://git-scm.com/docs/git-credential-store
"""
try:
run_subprocess("git config --global credential.helper store")
except subprocess.CalledProcessError as exc:
raise EnvironmentError(exc.stderr)
kronos-venv/lib/python3.12/site-packages/huggingface_hub/_oauth.py
+464
View File
@@ -0,0 +1,464 @@
import datetime
import hashlib
import logging
import os
import time
import urllib.parse
import warnings
from dataclasses import dataclass
from typing import TYPE_CHECKING, Dict, List, Optional, Tuple, Union
from . import constants
from .hf_api import whoami
from .utils import experimental, get_token
logger = logging.getLogger(__name__)
if TYPE_CHECKING:
import fastapi
@dataclass
class OAuthOrgInfo:
"""
Information about an organization linked to a user logged in with OAuth.
Attributes:
sub (`str`):
Unique identifier for the org. OpenID Connect field.
name (`str`):
The org's full name. OpenID Connect field.
preferred_username (`str`):
The org's username. OpenID Connect field.
picture (`str`):
The org's profile picture URL. OpenID Connect field.
is_enterprise (`bool`):
Whether the org is an enterprise org. Hugging Face field.
can_pay (`Optional[bool]`, *optional*):
Whether the org has a payment method set up. Hugging Face field.
role_in_org (`Optional[str]`, *optional*):
The user's role in the org. Hugging Face field.
pending_sso (`Optional[bool]`, *optional*):
Indicates if the user granted the OAuth app access to the org but didn't complete SSO. Hugging Face field.
missing_mfa (`Optional[bool]`, *optional*):
Indicates if the user granted the OAuth app access to the org but didn't complete MFA. Hugging Face field.
"""
sub: str
name: str
preferred_username: str
picture: str
is_enterprise: bool
can_pay: Optional[bool] = None
role_in_org: Optional[str] = None
pending_sso: Optional[bool] = None
missing_mfa: Optional[bool] = None
@dataclass
class OAuthUserInfo:
"""
Information about a user logged in with OAuth.
Attributes:
sub (`str`):
Unique identifier for the user, even in case of rename. OpenID Connect field.
name (`str`):
The user's full name. OpenID Connect field.
preferred_username (`str`):
The user's username. OpenID Connect field.
email_verified (`Optional[bool]`, *optional*):
Indicates if the user's email is verified. OpenID Connect field.
email (`Optional[str]`, *optional*):
The user's email address. OpenID Connect field.
picture (`str`):
The user's profile picture URL. OpenID Connect field.
profile (`str`):
The user's profile URL. OpenID Connect field.
website (`Optional[str]`, *optional*):
The user's website URL. OpenID Connect field.
is_pro (`bool`):
Whether the user is a pro user. Hugging Face field.
can_pay (`Optional[bool]`, *optional*):
Whether the user has a payment method set up. Hugging Face field.
orgs (`Optional[List[OrgInfo]]`, *optional*):
List of organizations the user is part of. Hugging Face field.
"""
sub: str
name: str
preferred_username: str
email_verified: Optional[bool]
email: Optional[str]
picture: str
profile: str
website: Optional[str]
is_pro: bool
can_pay: Optional[bool]
orgs: Optional[List[OAuthOrgInfo]]
@dataclass
class OAuthInfo:
"""
Information about the OAuth login.
Attributes:
access_token (`str`):
The access token.
access_token_expires_at (`datetime.datetime`):
The expiration date of the access token.
user_info ([`OAuthUserInfo`]):
The user information.
state (`str`, *optional*):
State passed to the OAuth provider in the original request to the OAuth provider.
scope (`str`):
Granted scope.
"""
access_token: str
access_token_expires_at: datetime.datetime
user_info: OAuthUserInfo
state: Optional[str]
scope: str
@experimental
def attach_huggingface_oauth(app: "fastapi.FastAPI", route_prefix: str = "/"):
"""
Add OAuth endpoints to a FastAPI app to enable OAuth login with Hugging Face.
How to use:
- Call this method on your FastAPI app to add the OAuth endpoints.
- Inside your route handlers, call `parse_huggingface_oauth(request)` to retrieve the OAuth info.
- If user is logged in, an [`OAuthInfo`] object is returned with the user's info. If not, `None` is returned.
- In your app, make sure to add links to `/oauth/huggingface/login` and `/oauth/huggingface/logout` for the user to log in and out.
Example:
```py
from huggingface_hub import attach_huggingface_oauth, parse_huggingface_oauth
# Create a FastAPI app
app = FastAPI()
# Add OAuth endpoints to the FastAPI app
attach_huggingface_oauth(app)
# Add a route that greets the user if they are logged in
@app.get("/")
def greet_json(request: Request):
# Retrieve the OAuth info from the request
oauth_info = parse_huggingface_oauth(request) # e.g. OAuthInfo dataclass
if oauth_info is None:
return {"msg": "Not logged in!"}
return {"msg": f"Hello, {oauth_info.user_info.preferred_username}!"}
```
"""
# TODO: handle generic case (handling OAuth in a non-Space environment with custom dev values) (low priority)
# Add SessionMiddleware to the FastAPI app to store the OAuth info in the session.
# Session Middleware requires a secret key to sign the cookies. Let's use a hash
# of the OAuth secret key to make it unique to the Space + updated in case OAuth
# config gets updated. When ran locally, we use an empty string as a secret key.
try:
from starlette.middleware.sessions import SessionMiddleware
except ImportError as e:
raise ImportError(
"Cannot initialize OAuth to due a missing library. Please run `pip install huggingface_hub[oauth]` or add "
"`huggingface_hub[oauth]` to your requirements.txt file in order to install the required dependencies."
) from e
session_secret = (constants.OAUTH_CLIENT_SECRET or "") + "-v1"
app.add_middleware(
SessionMiddleware, # type: ignore[arg-type]
secret_key=hashlib.sha256(session_secret.encode()).hexdigest(),
same_site="none",
https_only=True,
) # type: ignore
# Add OAuth endpoints to the FastAPI app:
# - {route_prefix}/oauth/huggingface/login
# - {route_prefix}/oauth/huggingface/callback
# - {route_prefix}/oauth/huggingface/logout
# If the app is running in a Space, OAuth is enabled normally.
# Otherwise, we mock the endpoints to make the user log in with a fake user profile - without any calls to hf.co.
route_prefix = route_prefix.strip("/")
if os.getenv("SPACE_ID") is not None:
logger.info("OAuth is enabled in the Space. Adding OAuth routes.")
_add_oauth_routes(app, route_prefix=route_prefix)
else:
logger.info("App is not running in a Space. Adding mocked OAuth routes.")
_add_mocked_oauth_routes(app, route_prefix=route_prefix)
def parse_huggingface_oauth(request: "fastapi.Request") -> Optional[OAuthInfo]:
"""
Returns the information from a logged in user as a [`OAuthInfo`] object.
For flexibility and future-proofing, this method is very lax in its parsing and does not raise errors.
Missing fields are set to `None` without a warning.
Return `None`, if the user is not logged in (no info in session cookie).
See [`attach_huggingface_oauth`] for an example on how to use this method.
"""
if "oauth_info" not in request.session:
logger.debug("No OAuth info in session.")
return None
logger.debug("Parsing OAuth info from session.")
oauth_data = request.session["oauth_info"]
user_data = oauth_data.get("userinfo", {})
orgs_data = user_data.get("orgs", [])
orgs = (
[
OAuthOrgInfo(
sub=org.get("sub"),
name=org.get("name"),
preferred_username=org.get("preferred_username"),
picture=org.get("picture"),
is_enterprise=org.get("isEnterprise"),
can_pay=org.get("canPay"),
role_in_org=org.get("roleInOrg"),
pending_sso=org.get("pendingSSO"),
missing_mfa=org.get("missingMFA"),
)
for org in orgs_data
]
if orgs_data
else None
)
user_info = OAuthUserInfo(
sub=user_data.get("sub"),
name=user_data.get("name"),
preferred_username=user_data.get("preferred_username"),
email_verified=user_data.get("email_verified"),
email=user_data.get("email"),
picture=user_data.get("picture"),
profile=user_data.get("profile"),
website=user_data.get("website"),
is_pro=user_data.get("isPro"),
can_pay=user_data.get("canPay"),
orgs=orgs,
)
return OAuthInfo(
access_token=oauth_data.get("access_token"),
access_token_expires_at=datetime.datetime.fromtimestamp(oauth_data.get("expires_at")),
user_info=user_info,
state=oauth_data.get("state"),
scope=oauth_data.get("scope"),
)
def _add_oauth_routes(app: "fastapi.FastAPI", route_prefix: str) -> None:
"""Add OAuth routes to the FastAPI app (login, callback handler and logout)."""
try:
import fastapi
from authlib.integrations.base_client.errors import MismatchingStateError
from authlib.integrations.starlette_client import OAuth
from fastapi.responses import RedirectResponse
except ImportError as e:
raise ImportError(
"Cannot initialize OAuth to due a missing library. Please run `pip install huggingface_hub[oauth]` or add "
"`huggingface_hub[oauth]` to your requirements.txt file."
) from e
# Check environment variables
msg = (
"OAuth is required but '{}' environment variable is not set. Make sure you've enabled OAuth in your Space by"
" setting `hf_oauth: true` in the Space metadata."
)
if constants.OAUTH_CLIENT_ID is None:
raise ValueError(msg.format("OAUTH_CLIENT_ID"))
if constants.OAUTH_CLIENT_SECRET is None:
raise ValueError(msg.format("OAUTH_CLIENT_SECRET"))
if constants.OAUTH_SCOPES is None:
raise ValueError(msg.format("OAUTH_SCOPES"))
if constants.OPENID_PROVIDER_URL is None:
raise ValueError(msg.format("OPENID_PROVIDER_URL"))
# Register OAuth server
oauth = OAuth()
oauth.register(
name="huggingface",
client_id=constants.OAUTH_CLIENT_ID,
client_secret=constants.OAUTH_CLIENT_SECRET,
client_kwargs={"scope": constants.OAUTH_SCOPES},
server_metadata_url=constants.OPENID_PROVIDER_URL + "/.well-known/openid-configuration",
)
login_uri, callback_uri, logout_uri = _get_oauth_uris(route_prefix)
# Register OAuth endpoints
@app.get(login_uri)
async def oauth_login(request: fastapi.Request) -> RedirectResponse:
"""Endpoint that redirects to HF OAuth page."""
redirect_uri = _generate_redirect_uri(request)
return await oauth.huggingface.authorize_redirect(request, redirect_uri) # type: ignore
@app.get(callback_uri)
async def oauth_redirect_callback(request: fastapi.Request) -> RedirectResponse:
"""Endpoint that handles the OAuth callback."""
try:
oauth_info = await oauth.huggingface.authorize_access_token(request) # type: ignore
except MismatchingStateError:
# Parse query params
nb_redirects = int(request.query_params.get("_nb_redirects", 0))
target_url = request.query_params.get("_target_url")
# Build redirect URI with the same query params as before and bump nb_redirects count
query_params: Dict[str, Union[int, str]] = {"_nb_redirects": nb_redirects + 1}
if target_url:
query_params["_target_url"] = target_url
redirect_uri = f"{login_uri}?{urllib.parse.urlencode(query_params)}"
# If the user is redirected more than 3 times, it is very likely that the cookie is not working properly.
# (e.g. browser is blocking third-party cookies in iframe). In this case, redirect the user in the
# non-iframe view.
if nb_redirects > constants.OAUTH_MAX_REDIRECTS:
host = os.environ.get("SPACE_HOST")
if host is None: # cannot happen in a Space
raise RuntimeError(
"App is not running in a Space (SPACE_HOST environment variable is not set). Cannot redirect to non-iframe view."
) from None
host_url = "https://" + host.rstrip("/")
return RedirectResponse(host_url + redirect_uri)
# Redirect the user to the login page again
return RedirectResponse(redirect_uri)
# OAuth login worked => store the user info in the session and redirect
logger.debug("Successfully logged in with OAuth. Storing user info in session.")
request.session["oauth_info"] = oauth_info
return RedirectResponse(_get_redirect_target(request))
@app.get(logout_uri)
async def oauth_logout(request: fastapi.Request) -> RedirectResponse:
"""Endpoint that logs out the user (e.g. delete info from cookie session)."""
logger.debug("Logged out with OAuth. Removing user info from session.")
request.session.pop("oauth_info", None)
return RedirectResponse(_get_redirect_target(request))
def _add_mocked_oauth_routes(app: "fastapi.FastAPI", route_prefix: str = "/") -> None:
"""Add fake oauth routes if app is run locally and OAuth is enabled.
Using OAuth will have the same behavior as in a Space but instead of authenticating with HF, a mocked user profile
is added to the session.
"""
try:
import fastapi
from fastapi.responses import RedirectResponse
from starlette.datastructures import URL
except ImportError as e:
raise ImportError(
"Cannot initialize OAuth to due a missing library. Please run `pip install huggingface_hub[oauth]` or add "
"`huggingface_hub[oauth]` to your requirements.txt file."
) from e
warnings.warn(
"OAuth is not supported outside of a Space environment. To help you debug your app locally, the oauth endpoints"
" are mocked to return your profile and token. To make it work, your machine must be logged in to Huggingface."
)
mocked_oauth_info = _get_mocked_oauth_info()
login_uri, callback_uri, logout_uri = _get_oauth_uris(route_prefix)
# Define OAuth routes
@app.get(login_uri)
async def oauth_login(request: fastapi.Request) -> RedirectResponse:
"""Fake endpoint that redirects to HF OAuth page."""
# Define target (where to redirect after login)
redirect_uri = _generate_redirect_uri(request)
return RedirectResponse(callback_uri + "?" + urllib.parse.urlencode({"_target_url": redirect_uri}))
@app.get(callback_uri)
async def oauth_redirect_callback(request: fastapi.Request) -> RedirectResponse:
"""Endpoint that handles the OAuth callback."""
request.session["oauth_info"] = mocked_oauth_info
return RedirectResponse(_get_redirect_target(request))
@app.get(logout_uri)
async def oauth_logout(request: fastapi.Request) -> RedirectResponse:
"""Endpoint that logs out the user (e.g. delete cookie session)."""
request.session.pop("oauth_info", None)
logout_url = URL("/").include_query_params(**request.query_params)
return RedirectResponse(url=logout_url, status_code=302) # see https://github.com/gradio-app/gradio/pull/9659
def _generate_redirect_uri(request: "fastapi.Request") -> str:
if "_target_url" in request.query_params:
# if `_target_url` already in query params => respect it
target = request.query_params["_target_url"]
else:
# otherwise => keep query params
target = "/?" + urllib.parse.urlencode(request.query_params)
redirect_uri = request.url_for("oauth_redirect_callback").include_query_params(_target_url=target)
redirect_uri_as_str = str(redirect_uri)
if redirect_uri.netloc.endswith(".hf.space"):
# In Space, FastAPI redirect as http but we want https
redirect_uri_as_str = redirect_uri_as_str.replace("http://", "https://")
return redirect_uri_as_str
def _get_redirect_target(request: "fastapi.Request", default_target: str = "/") -> str:
return request.query_params.get("_target_url", default_target)
def _get_mocked_oauth_info() -> Dict:
token = get_token()
if token is None:
raise ValueError(
"Your machine must be logged in to HF to debug an OAuth app locally. Please"
" run `huggingface-cli login` or set `HF_TOKEN` as environment variable "
"with one of your access token. You can generate a new token in your "
"settings page (https://huggingface.co/settings/tokens)."
)
user = whoami()
if user["type"] != "user":
raise ValueError(
"Your machine is not logged in with a personal account. Please use a "
"personal access token. You can generate a new token in your settings page"
" (https://huggingface.co/settings/tokens)."
)
return {
"access_token": token,
"token_type": "bearer",
"expires_in": 8 * 60 * 60, # 8 hours
"id_token": "FOOBAR",
"scope": "openid profile",
"refresh_token": "hf_oauth__refresh_token",
"expires_at": int(time.time()) + 8 * 60 * 60, # 8 hours
"userinfo": {
"sub": "0123456789",
"name": user["fullname"],
"preferred_username": user["name"],
"profile": f"https://huggingface.co/{user['name']}",
"picture": user["avatarUrl"],
"website": "",
"aud": "00000000-0000-0000-0000-000000000000",
"auth_time": 1691672844,
"nonce": "aaaaaaaaaaaaaaaaaaa",
"iat": 1691672844,
"exp": 1691676444,
"iss": "https://huggingface.co",
},
}
def _get_oauth_uris(route_prefix: str = "/") -> Tuple[str, str, str]:
route_prefix = route_prefix.strip("/")
if route_prefix:
route_prefix = f"/{route_prefix}"
return (
f"{route_prefix}/oauth/huggingface/login",
f"{route_prefix}/oauth/huggingface/callback",
f"{route_prefix}/oauth/huggingface/logout",
)
kronos-venv/lib/python3.12/site-packages/huggingface_hub/_snapshot_download.py
+338
View File
@@ -0,0 +1,338 @@
import os
from pathlib import Path
from typing import Dict, Iterable, List, Literal, Optional, Type, Union
import requests
from tqdm.auto import tqdm as base_tqdm
from tqdm.contrib.concurrent import thread_map
from . import constants
from .errors import (
GatedRepoError,
HfHubHTTPError,
LocalEntryNotFoundError,
RepositoryNotFoundError,
RevisionNotFoundError,
)
from .file_download import REGEX_COMMIT_HASH, hf_hub_download, repo_folder_name
from .hf_api import DatasetInfo, HfApi, ModelInfo, RepoFile, SpaceInfo
from .utils import OfflineModeIsEnabled, filter_repo_objects, logging, validate_hf_hub_args
from .utils import tqdm as hf_tqdm
logger = logging.get_logger(__name__)
VERY_LARGE_REPO_THRESHOLD = 50000 # After this limit, we don't consider `repo_info.siblings` to be reliable enough
@validate_hf_hub_args
def snapshot_download(
repo_id: str,
*,
repo_type: Optional[str] = None,
revision: Optional[str] = None,
cache_dir: Union[str, Path, None] = None,
local_dir: Union[str, Path, None] = None,
library_name: Optional[str] = None,
library_version: Optional[str] = None,
user_agent: Optional[Union[Dict, str]] = None,
proxies: Optional[Dict] = None,
etag_timeout: float = constants.DEFAULT_ETAG_TIMEOUT,
force_download: bool = False,
token: Optional[Union[bool, str]] = None,
local_files_only: bool = False,
allow_patterns: Optional[Union[List[str], str]] = None,
ignore_patterns: Optional[Union[List[str], str]] = None,
max_workers: int = 8,
tqdm_class: Optional[Type[base_tqdm]] = None,
headers: Optional[Dict[str, str]] = None,
endpoint: Optional[str] = None,
# Deprecated args
local_dir_use_symlinks: Union[bool, Literal["auto"]] = "auto",
resume_download: Optional[bool] = None,
) -> str:
"""Download repo files.
Download a whole snapshot of a repo's files at the specified revision. This is useful when you want all files from
a repo, because you don't know which ones you will need a priori. All files are nested inside a folder in order
to keep their actual filename relative to that folder. You can also filter which files to download using
`allow_patterns` and `ignore_patterns`.
If `local_dir` is provided, the file structure from the repo will be replicated in this location. When using this
option, the `cache_dir` will not be used and a `.cache/huggingface/` folder will be created at the root of `local_dir`
to store some metadata related to the downloaded files. While this mechanism is not as robust as the main
cache-system, it's optimized for regularly pulling the latest version of a repository.
An alternative would be to clone the repo but this requires git and git-lfs to be installed and properly
configured. It is also not possible to filter which files to download when cloning a repository using git.
Args:
repo_id (`str`):
A user or an organization name and a repo name separated by a `/`.
repo_type (`str`, *optional*):
Set to `"dataset"` or `"space"` if downloading from a dataset or space,
`None` or `"model"` if downloading from a model. Default is `None`.
revision (`str`, *optional*):
An optional Git revision id which can be a branch name, a tag, or a
commit hash.
cache_dir (`str`, `Path`, *optional*):
Path to the folder where cached files are stored.
local_dir (`str` or `Path`, *optional*):
If provided, the downloaded files will be placed under this directory.
library_name (`str`, *optional*):
The name of the library to which the object corresponds.
library_version (`str`, *optional*):
The version of the library.
user_agent (`str`, `dict`, *optional*):
The user-agent info in the form of a dictionary or a string.
proxies (`dict`, *optional*):
Dictionary mapping protocol to the URL of the proxy passed to
`requests.request`.
etag_timeout (`float`, *optional*, defaults to `10`):
When fetching ETag, how many seconds to wait for the server to send
data before giving up which is passed to `requests.request`.
force_download (`bool`, *optional*, defaults to `False`):
Whether the file should be downloaded even if it already exists in the local cache.
token (`str`, `bool`, *optional*):
A token to be used for the download.
- If `True`, the token is read from the HuggingFace config
folder.
- If a string, it's used as the authentication token.
headers (`dict`, *optional*):
Additional headers to include in the request. Those headers take precedence over the others.
local_files_only (`bool`, *optional*, defaults to `False`):
If `True`, avoid downloading the file and return the path to the
local cached file if it exists.
allow_patterns (`List[str]` or `str`, *optional*):
If provided, only files matching at least one pattern are downloaded.
ignore_patterns (`List[str]` or `str`, *optional*):
If provided, files matching any of the patterns are not downloaded.
max_workers (`int`, *optional*):
Number of concurrent threads to download files (1 thread = 1 file download).
Defaults to 8.
tqdm_class (`tqdm`, *optional*):
If provided, overwrites the default behavior for the progress bar. Passed
argument must inherit from `tqdm.auto.tqdm` or at least mimic its behavior.
Note that the `tqdm_class` is not passed to each individual download.
Defaults to the custom HF progress bar that can be disabled by setting
`HF_HUB_DISABLE_PROGRESS_BARS` environment variable.
Returns:
`str`: folder path of the repo snapshot.
Raises:
[`~utils.RepositoryNotFoundError`]
If the repository to download from cannot be found. This may be because it doesn't exist,
or because it is set to `private` and you do not have access.
[`~utils.RevisionNotFoundError`]
If the revision to download from cannot be found.
[`EnvironmentError`](https://docs.python.org/3/library/exceptions.html#EnvironmentError)
If `token=True` and the token cannot be found.
[`OSError`](https://docs.python.org/3/library/exceptions.html#OSError) if
ETag cannot be determined.
[`ValueError`](https://docs.python.org/3/library/exceptions.html#ValueError)
if some parameter value is invalid.
"""
if cache_dir is None:
cache_dir = constants.HF_HUB_CACHE
if revision is None:
revision = constants.DEFAULT_REVISION
if isinstance(cache_dir, Path):
cache_dir = str(cache_dir)
if repo_type is None:
repo_type = "model"
if repo_type not in constants.REPO_TYPES:
raise ValueError(f"Invalid repo type: {repo_type}. Accepted repo types are: {str(constants.REPO_TYPES)}")
storage_folder = os.path.join(cache_dir, repo_folder_name(repo_id=repo_id, repo_type=repo_type))
api = HfApi(
library_name=library_name,
library_version=library_version,
user_agent=user_agent,
endpoint=endpoint,
headers=headers,
token=token,
)
repo_info: Union[ModelInfo, DatasetInfo, SpaceInfo, None] = None
api_call_error: Optional[Exception] = None
if not local_files_only:
# try/except logic to handle different errors => taken from `hf_hub_download`
try:
# if we have internet connection we want to list files to download
repo_info = api.repo_info(repo_id=repo_id, repo_type=repo_type, revision=revision)
except (requests.exceptions.SSLError, requests.exceptions.ProxyError):
# Actually raise for those subclasses of ConnectionError
raise
except (
requests.exceptions.ConnectionError,
requests.exceptions.Timeout,
OfflineModeIsEnabled,
) as error:
# Internet connection is down
# => will try to use local files only
api_call_error = error
pass
except RevisionNotFoundError:
# The repo was found but the revision doesn't exist on the Hub (never existed or got deleted)
raise
except requests.HTTPError as error:
# Multiple reasons for an http error:
# - Repository is private and invalid/missing token sent
# - Repository is gated and invalid/missing token sent
# - Hub is down (error 500 or 504)
# => let's switch to 'local_files_only=True' to check if the files are already cached.
# (if it's not the case, the error will be re-raised)
api_call_error = error
pass
# At this stage, if `repo_info` is None it means either:
# - internet connection is down
# - internet connection is deactivated (local_files_only=True or HF_HUB_OFFLINE=True)
# - repo is private/gated and invalid/missing token sent
# - Hub is down
# => let's look if we can find the appropriate folder in the cache:
# - if the specified revision is a commit hash, look inside "snapshots".
# - f the specified revision is a branch or tag, look inside "refs".
# => if local_dir is not None, we will return the path to the local folder if it exists.
if repo_info is None:
# Try to get which commit hash corresponds to the specified revision
commit_hash = None
if REGEX_COMMIT_HASH.match(revision):
commit_hash = revision
else:
ref_path = os.path.join(storage_folder, "refs", revision)
if os.path.exists(ref_path):
# retrieve commit_hash from refs file
with open(ref_path) as f:
commit_hash = f.read()
# Try to locate snapshot folder for this commit hash
if commit_hash is not None and local_dir is None:
snapshot_folder = os.path.join(storage_folder, "snapshots", commit_hash)
if os.path.exists(snapshot_folder):
# Snapshot folder exists => let's return it
# (but we can't check if all the files are actually there)
return snapshot_folder
# If local_dir is not None, return it if it exists and is not empty
if local_dir is not None:
local_dir = Path(local_dir)
if local_dir.is_dir() and any(local_dir.iterdir()):
logger.warning(
f"Returning existing local_dir `{local_dir}` as remote repo cannot be accessed in `snapshot_download` ({api_call_error})."
)
return str(local_dir.resolve())
# If we couldn't find the appropriate folder on disk, raise an error.
if local_files_only:
raise LocalEntryNotFoundError(
"Cannot find an appropriate cached snapshot folder for the specified revision on the local disk and "
"outgoing traffic has been disabled. To enable repo look-ups and downloads online, pass "
"'local_files_only=False' as input."
)
elif isinstance(api_call_error, OfflineModeIsEnabled):
raise LocalEntryNotFoundError(
"Cannot find an appropriate cached snapshot folder for the specified revision on the local disk and "
"outgoing traffic has been disabled. To enable repo look-ups and downloads online, set "
"'HF_HUB_OFFLINE=0' as environment variable."
) from api_call_error
elif isinstance(api_call_error, (RepositoryNotFoundError, GatedRepoError)) or (
isinstance(api_call_error, HfHubHTTPError) and api_call_error.response.status_code == 401
):
# Repo not found, gated, or specific authentication error => let's raise the actual error
raise api_call_error
else:
# Otherwise: most likely a connection issue or Hub downtime => let's warn the user
raise LocalEntryNotFoundError(
"An error happened while trying to locate the files on the Hub and we cannot find the appropriate"
" snapshot folder for the specified revision on the local disk. Please check your internet connection"
" and try again."
) from api_call_error
# At this stage, internet connection is up and running
# => let's download the files!
assert repo_info.sha is not None, "Repo info returned from server must have a revision sha."
assert repo_info.siblings is not None, "Repo info returned from server must have a siblings list."
# Corner case: on very large repos, the siblings list in `repo_info` might not contain all files.
# In that case, we need to use the `list_repo_tree` method to prevent caching issues.
repo_files: Iterable[str] = [f.rfilename for f in repo_info.siblings]
has_many_files = len(repo_info.siblings) > VERY_LARGE_REPO_THRESHOLD
if has_many_files:
logger.info("The repo has more than 50,000 files. Using `list_repo_tree` to ensure all files are listed.")
repo_files = (
f.rfilename
for f in api.list_repo_tree(repo_id=repo_id, recursive=True, revision=revision, repo_type=repo_type)
if isinstance(f, RepoFile)
)
filtered_repo_files: Iterable[str] = filter_repo_objects(
items=repo_files,
allow_patterns=allow_patterns,
ignore_patterns=ignore_patterns,
)
if not has_many_files:
filtered_repo_files = list(filtered_repo_files)
tqdm_desc = f"Fetching {len(filtered_repo_files)} files"
else:
tqdm_desc = "Fetching ... files"
commit_hash = repo_info.sha
snapshot_folder = os.path.join(storage_folder, "snapshots", commit_hash)
# if passed revision is not identical to commit_hash
# then revision has to be a branch name or tag name.
# In that case store a ref.
if revision != commit_hash:
ref_path = os.path.join(storage_folder, "refs", revision)
try:
os.makedirs(os.path.dirname(ref_path), exist_ok=True)
with open(ref_path, "w") as f:
f.write(commit_hash)
except OSError as e:
logger.warning(f"Ignored error while writing commit hash to {ref_path}: {e}.")
# we pass the commit_hash to hf_hub_download
# so no network call happens if we already
# have the file locally.
def _inner_hf_hub_download(repo_file: str):
return hf_hub_download(
repo_id,
filename=repo_file,
repo_type=repo_type,
revision=commit_hash,
endpoint=endpoint,
cache_dir=cache_dir,
local_dir=local_dir,
local_dir_use_symlinks=local_dir_use_symlinks,
library_name=library_name,
library_version=library_version,
user_agent=user_agent,
proxies=proxies,
etag_timeout=etag_timeout,
resume_download=resume_download,
force_download=force_download,
token=token,
headers=headers,
)
if constants.HF_HUB_ENABLE_HF_TRANSFER:
# when using hf_transfer we don't want extra parallelism
# from the one hf_transfer provides
for file in filtered_repo_files:
_inner_hf_hub_download(file)
else:
thread_map(
_inner_hf_hub_download,
filtered_repo_files,
desc=tqdm_desc,
max_workers=max_workers,
# User can use its own tqdm class or the default one from `huggingface_hub.utils`
tqdm_class=tqdm_class or hf_tqdm,
)
if local_dir is not None:
return str(os.path.realpath(local_dir))
return snapshot_folder
kronos-venv/lib/python3.12/site-packages/huggingface_hub/_space_api.py
+168
View File
@@ -0,0 +1,168 @@
# coding=utf-8
# Copyright 2019-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.
from dataclasses import dataclass
from datetime import datetime
from enum import Enum
from typing import Dict, Optional
from huggingface_hub.utils import parse_datetime
class SpaceStage(str, Enum):
"""
Enumeration of possible stage of a Space on the Hub.
Value can be compared to a string:
```py
assert SpaceStage.BUILDING == "BUILDING"
```
Taken from https://github.com/huggingface/moon-landing/blob/main/server/repo_types/SpaceInfo.ts#L61 (private url).
"""
# Copied from moon-landing > server > repo_types > SpaceInfo.ts (private repo)
NO_APP_FILE = "NO_APP_FILE"
CONFIG_ERROR = "CONFIG_ERROR"
BUILDING = "BUILDING"
BUILD_ERROR = "BUILD_ERROR"
RUNNING = "RUNNING"
RUNNING_BUILDING = "RUNNING_BUILDING"
RUNTIME_ERROR = "RUNTIME_ERROR"
DELETING = "DELETING"
STOPPED = "STOPPED"
PAUSED = "PAUSED"
class SpaceHardware(str, Enum):
"""
Enumeration of hardwares available to run your Space on the Hub.
Value can be compared to a string:
```py
assert SpaceHardware.CPU_BASIC == "cpu-basic"
```
Taken from https://github.com/huggingface-internal/moon-landing/blob/main/server/repo_types/SpaceHardwareFlavor.ts (private url).
"""
# CPU
CPU_BASIC = "cpu-basic"
CPU_UPGRADE = "cpu-upgrade"
CPU_XL = "cpu-xl"
# ZeroGPU
ZERO_A10G = "zero-a10g"
# GPU
T4_SMALL = "t4-small"
T4_MEDIUM = "t4-medium"
L4X1 = "l4x1"
L4X4 = "l4x4"
L40SX1 = "l40sx1"
L40SX4 = "l40sx4"
L40SX8 = "l40sx8"
A10G_SMALL = "a10g-small"
A10G_LARGE = "a10g-large"
A10G_LARGEX2 = "a10g-largex2"
A10G_LARGEX4 = "a10g-largex4"
A100_LARGE = "a100-large"
H100 = "h100"
H100X8 = "h100x8"
class SpaceStorage(str, Enum):
"""
Enumeration of persistent storage available for your Space on the Hub.
Value can be compared to a string:
```py
assert SpaceStorage.SMALL == "small"
```
Taken from https://github.com/huggingface/moon-landing/blob/main/server/repo_types/SpaceHardwareFlavor.ts#L24 (private url).
"""
SMALL = "small"
MEDIUM = "medium"
LARGE = "large"
@dataclass
class SpaceRuntime:
"""
Contains information about the current runtime of a Space.
Args:
stage (`str`):
Current stage of the space. Example: RUNNING.
hardware (`str` or `None`):
Current hardware of the space. Example: "cpu-basic". Can be `None` if Space
is `BUILDING` for the first time.
requested_hardware (`str` or `None`):
Requested hardware. Can be different than `hardware` especially if the request
has just been made. Example: "t4-medium". Can be `None` if no hardware has
been requested yet.
sleep_time (`int` or `None`):
Number of seconds the Space will be kept alive after the last request. By default (if value is `None`), the
Space will never go to sleep if it's running on an upgraded hardware, while it will go to sleep after 48
hours on a free 'cpu-basic' hardware. For more details, see https://huggingface.co/docs/hub/spaces-gpus#sleep-time.
raw (`dict`):
Raw response from the server. Contains more information about the Space
runtime like number of replicas, number of cpu, memory size,...
"""
stage: SpaceStage
hardware: Optional[SpaceHardware]
requested_hardware: Optional[SpaceHardware]
sleep_time: Optional[int]
storage: Optional[SpaceStorage]
raw: Dict
def __init__(self, data: Dict) -> None:
self.stage = data["stage"]
self.hardware = data.get("hardware", {}).get("current")
self.requested_hardware = data.get("hardware", {}).get("requested")
self.sleep_time = data.get("gcTimeout")
self.storage = data.get("storage")
self.raw = data
@dataclass
class SpaceVariable:
"""
Contains information about the current variables of a Space.
Args:
key (`str`):
Variable key. Example: `"MODEL_REPO_ID"`
value (`str`):
Variable value. Example: `"the_model_repo_id"`.
description (`str` or None):
Description of the variable. Example: `"Model Repo ID of the implemented model"`.
updatedAt (`datetime` or None):
datetime of the last update of the variable (if the variable has been updated at least once).
"""
key: str
value: str
description: Optional[str]
updated_at: Optional[datetime]
def __init__(self, key: str, values: Dict) -> None:
self.key = key
self.value = values["value"]
self.description = values.get("description")
updated_at = values.get("updatedAt")
self.updated_at = parse_datetime(updated_at) if updated_at is not None else None
kronos-venv/lib/python3.12/site-packages/huggingface_hub/_tensorboard_logger.py
+194
View File
@@ -0,0 +1,194 @@
# Copyright 2023 The HuggingFace Team. All rights reserved.
#
# 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 a logger to push training logs to the Hub, using Tensorboard."""
from pathlib import Path
from typing import TYPE_CHECKING, List, Optional, Union
from ._commit_scheduler import CommitScheduler
from .errors import EntryNotFoundError
from .repocard import ModelCard
from .utils import experimental
# Depending on user's setup, SummaryWriter can come either from 'tensorboardX'
# or from 'torch.utils.tensorboard'. Both are compatible so let's try to load
# from either of them.
try:
from tensorboardX import SummaryWriter
is_summary_writer_available = True
except ImportError:
try:
from torch.utils.tensorboard import SummaryWriter
is_summary_writer_available = False
except ImportError:
# Dummy class to avoid failing at import. Will raise on instance creation.
SummaryWriter = object
is_summary_writer_available = False
if TYPE_CHECKING:
from tensorboardX import SummaryWriter
class HFSummaryWriter(SummaryWriter):
"""
Wrapper around the tensorboard's `SummaryWriter` to push training logs to the Hub.
Data is logged locally and then pushed to the Hub asynchronously. Pushing data to the Hub is done in a separate
thread to avoid blocking the training script. In particular, if the upload fails for any reason (e.g. a connection
issue), the main script will not be interrupted. Data is automatically pushed to the Hub every `commit_every`
minutes (default to every 5 minutes).
<Tip warning={true}>
`HFSummaryWriter` is experimental. Its API is subject to change in the future without prior notice.
</Tip>
Args:
repo_id (`str`):
The id of the repo to which the logs will be pushed.
logdir (`str`, *optional*):
The directory where the logs will be written. If not specified, a local directory will be created by the
underlying `SummaryWriter` object.
commit_every (`int` or `float`, *optional*):
The frequency (in minutes) at which the logs will be pushed to the Hub. Defaults to 5 minutes.
squash_history (`bool`, *optional*):
Whether to squash the history of the repo after each commit. Defaults to `False`. Squashing commits is
useful to avoid degraded performances on the repo when it grows too large.
repo_type (`str`, *optional*):
The type of the repo to which the logs will be pushed. Defaults to "model".
repo_revision (`str`, *optional*):
The revision of the repo to which the logs will be pushed. Defaults to "main".
repo_private (`bool`, *optional*):
Whether to make the repo private. If `None` (default), the repo will be public unless the organization's default is private. This value is ignored if the repo already exists.
path_in_repo (`str`, *optional*):
The path to the folder in the repo where the logs will be pushed. Defaults to "tensorboard/".
repo_allow_patterns (`List[str]` or `str`, *optional*):
A list of patterns to include in the upload. Defaults to `"*.tfevents.*"`. Check out the
[upload guide](https://huggingface.co/docs/huggingface_hub/guides/upload#upload-a-folder) for more details.
repo_ignore_patterns (`List[str]` or `str`, *optional*):
A list of patterns to exclude in the upload. Check out the
[upload guide](https://huggingface.co/docs/huggingface_hub/guides/upload#upload-a-folder) for more details.
token (`str`, *optional*):
Authentication token. Will default to the stored token. See https://huggingface.co/settings/token for more
details
kwargs:
Additional keyword arguments passed to `SummaryWriter`.
Examples:
```diff
# Taken from https://pytorch.org/docs/stable/tensorboard.html
- from torch.utils.tensorboard import SummaryWriter
+ from huggingface_hub import HFSummaryWriter
import numpy as np
- writer = SummaryWriter()
+ writer = HFSummaryWriter(repo_id="username/my-trained-model")
for n_iter in range(100):
writer.add_scalar('Loss/train', np.random.random(), n_iter)
writer.add_scalar('Loss/test', np.random.random(), n_iter)
writer.add_scalar('Accuracy/train', np.random.random(), n_iter)
writer.add_scalar('Accuracy/test', np.random.random(), n_iter)
```
```py
>>> from huggingface_hub import HFSummaryWriter
# Logs are automatically pushed every 15 minutes (5 by default) + when exiting the context manager
>>> with HFSummaryWriter(repo_id="test_hf_logger", commit_every=15) as logger:
... logger.add_scalar("a", 1)
... logger.add_scalar("b", 2)
```
"""
@experimental
def __new__(cls, *args, **kwargs) -> "HFSummaryWriter":
if not is_summary_writer_available:
raise ImportError(
"You must have `tensorboard` installed to use `HFSummaryWriter`. Please run `pip install --upgrade"
" tensorboardX` first."
)
return super().__new__(cls)
def __init__(
self,
repo_id: str,
*,
logdir: Optional[str] = None,
commit_every: Union[int, float] = 5,
squash_history: bool = False,
repo_type: Optional[str] = None,
repo_revision: Optional[str] = None,
repo_private: Optional[bool] = None,
path_in_repo: Optional[str] = "tensorboard",
repo_allow_patterns: Optional[Union[List[str], str]] = "*.tfevents.*",
repo_ignore_patterns: Optional[Union[List[str], str]] = None,
token: Optional[str] = None,
**kwargs,
):
# Initialize SummaryWriter
super().__init__(logdir=logdir, **kwargs)
# Check logdir has been correctly initialized and fail early otherwise. In practice, SummaryWriter takes care of it.
if not isinstance(self.logdir, str):
raise ValueError(f"`self.logdir` must be a string. Got '{self.logdir}' of type {type(self.logdir)}.")
# Append logdir name to `path_in_repo`
if path_in_repo is None or path_in_repo == "":
path_in_repo = Path(self.logdir).name
else:
path_in_repo = path_in_repo.strip("/") + "/" + Path(self.logdir).name
# Initialize scheduler
self.scheduler = CommitScheduler(
folder_path=self.logdir,
path_in_repo=path_in_repo,
repo_id=repo_id,
repo_type=repo_type,
revision=repo_revision,
private=repo_private,
token=token,
allow_patterns=repo_allow_patterns,
ignore_patterns=repo_ignore_patterns,
every=commit_every,
squash_history=squash_history,
)
# Exposing some high-level info at root level
self.repo_id = self.scheduler.repo_id
self.repo_type = self.scheduler.repo_type
self.repo_revision = self.scheduler.revision
# Add `hf-summary-writer` tag to the model card metadata
try:
card = ModelCard.load(repo_id_or_path=self.repo_id, repo_type=self.repo_type)
except EntryNotFoundError:
card = ModelCard("")
tags = card.data.get("tags", [])
if "hf-summary-writer" not in tags:
tags.append("hf-summary-writer")
card.data["tags"] = tags
card.push_to_hub(repo_id=self.repo_id, repo_type=self.repo_type)
def __exit__(self, exc_type, exc_val, exc_tb):
"""Push to hub in a non-blocking way when exiting the logger's context manager."""
super().__exit__(exc_type, exc_val, exc_tb)
future = self.scheduler.trigger()
future.result()
kronos-venv/lib/python3.12/site-packages/huggingface_hub/_upload_large_folder.py
+625
View File
@@ -0,0 +1,625 @@
# coding=utf-8
# Copyright 2024-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.
import enum
import logging
import os
import queue
import shutil
import sys
import threading
import time
import traceback
from datetime import datetime
from pathlib import Path
from threading import Lock
from typing import TYPE_CHECKING, List, Optional, Tuple, Union
from urllib.parse import quote
from . import constants
from ._commit_api import CommitOperationAdd, UploadInfo, _fetch_upload_modes
from ._local_folder import LocalUploadFileMetadata, LocalUploadFilePaths, get_local_upload_paths, read_upload_metadata
from .constants import DEFAULT_REVISION, REPO_TYPES
from .utils import DEFAULT_IGNORE_PATTERNS, filter_repo_objects, tqdm
from .utils._cache_manager import _format_size
from .utils.sha import sha_fileobj
if TYPE_CHECKING:
from .hf_api import HfApi
logger = logging.getLogger(__name__)
WAITING_TIME_IF_NO_TASKS = 10 # seconds
MAX_NB_FILES_FETCH_UPLOAD_MODE = 100
COMMIT_SIZE_SCALE: List[int] = [20, 50, 75, 100, 125, 200, 250, 400, 600, 1000]
def upload_large_folder_internal(
api: "HfApi",
repo_id: str,
folder_path: Union[str, Path],
*,
repo_type: str, # Repo type is required!
revision: Optional[str] = None,
private: Optional[bool] = None,
allow_patterns: Optional[Union[List[str], str]] = None,
ignore_patterns: Optional[Union[List[str], str]] = None,
num_workers: Optional[int] = None,
print_report: bool = True,
print_report_every: int = 60,
):
"""Upload a large folder to the Hub in the most resilient way possible.
See [`HfApi.upload_large_folder`] for the full documentation.
"""
# 1. Check args and setup
if repo_type is None:
raise ValueError(
"For large uploads, `repo_type` is explicitly required. Please set it to `model`, `dataset` or `space`."
" If you are using the CLI, pass it as `--repo-type=model`."
)
if repo_type not in REPO_TYPES:
raise ValueError(f"Invalid repo type, must be one of {REPO_TYPES}")
if revision is None:
revision = DEFAULT_REVISION
folder_path = Path(folder_path).expanduser().resolve()
if not folder_path.is_dir():
raise ValueError(f"Provided path: '{folder_path}' is not a directory")
if ignore_patterns is None:
ignore_patterns = []
elif isinstance(ignore_patterns, str):
ignore_patterns = [ignore_patterns]
ignore_patterns += DEFAULT_IGNORE_PATTERNS
if num_workers is None:
nb_cores = os.cpu_count() or 1
num_workers = max(nb_cores - 2, 2) # Use all but 2 cores, or at least 2 cores
# 2. Create repo if missing
repo_url = api.create_repo(repo_id=repo_id, repo_type=repo_type, private=private, exist_ok=True)
logger.info(f"Repo created: {repo_url}")
repo_id = repo_url.repo_id
# 3. List files to upload
filtered_paths_list = filter_repo_objects(
(path.relative_to(folder_path).as_posix() for path in folder_path.glob("**/*") if path.is_file()),
allow_patterns=allow_patterns,
ignore_patterns=ignore_patterns,
)
paths_list = [get_local_upload_paths(folder_path, relpath) for relpath in filtered_paths_list]
logger.info(f"Found {len(paths_list)} candidate files to upload")
# Read metadata for each file
items = [
(paths, read_upload_metadata(folder_path, paths.path_in_repo))
for paths in tqdm(paths_list, desc="Recovering from metadata files")
]
# 4. Start workers
status = LargeUploadStatus(items)
threads = [
threading.Thread(
target=_worker_job,
kwargs={
"status": status,
"api": api,
"repo_id": repo_id,
"repo_type": repo_type,
"revision": revision,
},
)
for _ in range(num_workers)
]
for thread in threads:
thread.start()
# 5. Print regular reports
if print_report:
print("\n\n" + status.current_report())
last_report_ts = time.time()
while True:
time.sleep(1)
if time.time() - last_report_ts >= print_report_every:
if print_report:
_print_overwrite(status.current_report())
last_report_ts = time.time()
if status.is_done():
logging.info("Is done: exiting main loop")
break
for thread in threads:
thread.join()
logger.info(status.current_report())
logging.info("Upload is complete!")
####################
# Logic to manage workers and synchronize tasks
####################
class WorkerJob(enum.Enum):
SHA256 = enum.auto()
GET_UPLOAD_MODE = enum.auto()
PREUPLOAD_LFS = enum.auto()
COMMIT = enum.auto()
WAIT = enum.auto() # if no tasks are available but we don't want to exit
JOB_ITEM_T = Tuple[LocalUploadFilePaths, LocalUploadFileMetadata]
class LargeUploadStatus:
"""Contains information, queues and tasks for a large upload process."""
def __init__(self, items: List[JOB_ITEM_T]):
self.items = items
self.queue_sha256: "queue.Queue[JOB_ITEM_T]" = queue.Queue()
self.queue_get_upload_mode: "queue.Queue[JOB_ITEM_T]" = queue.Queue()
self.queue_preupload_lfs: "queue.Queue[JOB_ITEM_T]" = queue.Queue()
self.queue_commit: "queue.Queue[JOB_ITEM_T]" = queue.Queue()
self.lock = Lock()
self.nb_workers_sha256: int = 0
self.nb_workers_get_upload_mode: int = 0
self.nb_workers_preupload_lfs: int = 0
self.nb_workers_commit: int = 0
self.nb_workers_waiting: int = 0
self.last_commit_attempt: Optional[float] = None
self._started_at = datetime.now()
self._chunk_idx: int = 1
self._chunk_lock: Lock = Lock()
# Setup queues
for item in self.items:
paths, metadata = item
if metadata.sha256 is None:
self.queue_sha256.put(item)
elif metadata.upload_mode is None:
self.queue_get_upload_mode.put(item)
elif metadata.upload_mode == "lfs" and not metadata.is_uploaded:
self.queue_preupload_lfs.put(item)
elif not metadata.is_committed:
self.queue_commit.put(item)
else:
logger.debug(f"Skipping file {paths.path_in_repo} (already uploaded and committed)")
def target_chunk(self) -> int:
with self._chunk_lock:
return COMMIT_SIZE_SCALE[self._chunk_idx]
def update_chunk(self, success: bool, nb_items: int, duration: float) -> None:
with self._chunk_lock:
if not success:
logger.warning(f"Failed to commit {nb_items} files at once. Will retry with less files in next batch.")
self._chunk_idx -= 1
elif nb_items >= COMMIT_SIZE_SCALE[self._chunk_idx] and duration < 40:
logger.info(f"Successfully committed {nb_items} at once. Increasing the limit for next batch.")
self._chunk_idx += 1
self._chunk_idx = max(0, min(self._chunk_idx, len(COMMIT_SIZE_SCALE) - 1))
def current_report(self) -> str:
"""Generate a report of the current status of the large upload."""
nb_hashed = 0
size_hashed = 0
nb_preuploaded = 0
nb_lfs = 0
nb_lfs_unsure = 0
size_preuploaded = 0
nb_committed = 0
size_committed = 0
total_size = 0
ignored_files = 0
total_files = 0
with self.lock:
for _, metadata in self.items:
if metadata.should_ignore:
ignored_files += 1
continue
total_size += metadata.size
total_files += 1
if metadata.sha256 is not None:
nb_hashed += 1
size_hashed += metadata.size
if metadata.upload_mode == "lfs":
nb_lfs += 1
if metadata.upload_mode is None:
nb_lfs_unsure += 1
if metadata.is_uploaded:
nb_preuploaded += 1
size_preuploaded += metadata.size
if metadata.is_committed:
nb_committed += 1
size_committed += metadata.size
total_size_str = _format_size(total_size)
now = datetime.now()
now_str = now.strftime("%Y-%m-%d %H:%M:%S")
elapsed = now - self._started_at
elapsed_str = str(elapsed).split(".")[0] # remove milliseconds
message = "\n" + "-" * 10
message += f" {now_str} ({elapsed_str}) "
message += "-" * 10 + "\n"
message += "Files: "
message += f"hashed {nb_hashed}/{total_files} ({_format_size(size_hashed)}/{total_size_str}) | "
message += f"pre-uploaded: {nb_preuploaded}/{nb_lfs} ({_format_size(size_preuploaded)}/{total_size_str})"
if nb_lfs_unsure > 0:
message += f" (+{nb_lfs_unsure} unsure)"
message += f" | committed: {nb_committed}/{total_files} ({_format_size(size_committed)}/{total_size_str})"
message += f" | ignored: {ignored_files}\n"
message += "Workers: "
message += f"hashing: {self.nb_workers_sha256} | "
message += f"get upload mode: {self.nb_workers_get_upload_mode} | "
message += f"pre-uploading: {self.nb_workers_preupload_lfs} | "
message += f"committing: {self.nb_workers_commit} | "
message += f"waiting: {self.nb_workers_waiting}\n"
message += "-" * 51
return message
def is_done(self) -> bool:
with self.lock:
return all(metadata.is_committed or metadata.should_ignore for _, metadata in self.items)
def _worker_job(
status: LargeUploadStatus,
api: "HfApi",
repo_id: str,
repo_type: str,
revision: str,
):
"""
Main process for a worker. The worker will perform tasks based on the priority list until all files are uploaded
and committed. If no tasks are available, the worker will wait for 10 seconds before checking again.
If a task fails for any reason, the item(s) are put back in the queue for another worker to pick up.
Read `upload_large_folder` docstring for more information on how tasks are prioritized.
"""
while True:
next_job: Optional[Tuple[WorkerJob, List[JOB_ITEM_T]]] = None
# Determine next task
next_job = _determine_next_job(status)
if next_job is None:
return
job, items = next_job
# Perform task
if job == WorkerJob.SHA256:
item = items[0] # single item
try:
_compute_sha256(item)
status.queue_get_upload_mode.put(item)
except KeyboardInterrupt:
raise
except Exception as e:
logger.error(f"Failed to compute sha256: {e}")
traceback.format_exc()
status.queue_sha256.put(item)
with status.lock:
status.nb_workers_sha256 -= 1
elif job == WorkerJob.GET_UPLOAD_MODE:
try:
_get_upload_mode(items, api=api, repo_id=repo_id, repo_type=repo_type, revision=revision)
except KeyboardInterrupt:
raise
except Exception as e:
logger.error(f"Failed to get upload mode: {e}")
traceback.format_exc()
# Items are either:
# - dropped (if should_ignore)
# - put in LFS queue (if LFS)
# - put in commit queue (if regular)
# - or put back (if error occurred).
for item in items:
_, metadata = item
if metadata.should_ignore:
continue
if metadata.upload_mode == "lfs":
status.queue_preupload_lfs.put(item)
elif metadata.upload_mode == "regular":
status.queue_commit.put(item)
else:
status.queue_get_upload_mode.put(item)
with status.lock:
status.nb_workers_get_upload_mode -= 1
elif job == WorkerJob.PREUPLOAD_LFS:
item = items[0] # single item
try:
_preupload_lfs(item, api=api, repo_id=repo_id, repo_type=repo_type, revision=revision)
status.queue_commit.put(item)
except KeyboardInterrupt:
raise
except Exception as e:
logger.error(f"Failed to preupload LFS: {e}")
traceback.format_exc()
status.queue_preupload_lfs.put(item)
with status.lock:
status.nb_workers_preupload_lfs -= 1
elif job == WorkerJob.COMMIT:
start_ts = time.time()
success = True
try:
_commit(items, api=api, repo_id=repo_id, repo_type=repo_type, revision=revision)
except KeyboardInterrupt:
raise
except Exception as e:
logger.error(f"Failed to commit: {e}")
traceback.format_exc()
for item in items:
status.queue_commit.put(item)
success = False
duration = time.time() - start_ts
status.update_chunk(success, len(items), duration)
with status.lock:
status.last_commit_attempt = time.time()
status.nb_workers_commit -= 1
elif job == WorkerJob.WAIT:
time.sleep(WAITING_TIME_IF_NO_TASKS)
with status.lock:
status.nb_workers_waiting -= 1
def _determine_next_job(status: LargeUploadStatus) -> Optional[Tuple[WorkerJob, List[JOB_ITEM_T]]]:
with status.lock:
# 1. Commit if more than 5 minutes since last commit attempt (and at least 1 file)
if (
status.nb_workers_commit == 0
and status.queue_commit.qsize() > 0
and status.last_commit_attempt is not None
and time.time() - status.last_commit_attempt > 5 * 60
):
status.nb_workers_commit += 1
logger.debug("Job: commit (more than 5 minutes since last commit attempt)")
return (WorkerJob.COMMIT, _get_n(status.queue_commit, status.target_chunk()))
# 2. Commit if at least 100 files are ready to commit
elif status.nb_workers_commit == 0 and status.queue_commit.qsize() >= 150:
status.nb_workers_commit += 1
logger.debug("Job: commit (>100 files ready)")
return (WorkerJob.COMMIT, _get_n(status.queue_commit, status.target_chunk()))
# 3. Get upload mode if at least 100 files
elif status.queue_get_upload_mode.qsize() >= MAX_NB_FILES_FETCH_UPLOAD_MODE:
status.nb_workers_get_upload_mode += 1
logger.debug(f"Job: get upload mode (>{MAX_NB_FILES_FETCH_UPLOAD_MODE} files ready)")
return (WorkerJob.GET_UPLOAD_MODE, _get_n(status.queue_get_upload_mode, MAX_NB_FILES_FETCH_UPLOAD_MODE))
# 4. Preupload LFS file if at least 1 file and no worker is preuploading LFS
elif status.queue_preupload_lfs.qsize() > 0 and status.nb_workers_preupload_lfs == 0:
status.nb_workers_preupload_lfs += 1
logger.debug("Job: preupload LFS (no other worker preuploading LFS)")
return (WorkerJob.PREUPLOAD_LFS, _get_one(status.queue_preupload_lfs))
# 5. Compute sha256 if at least 1 file and no worker is computing sha256
elif status.queue_sha256.qsize() > 0 and status.nb_workers_sha256 == 0:
status.nb_workers_sha256 += 1
logger.debug("Job: sha256 (no other worker computing sha256)")
return (WorkerJob.SHA256, _get_one(status.queue_sha256))
# 6. Get upload mode if at least 1 file and no worker is getting upload mode
elif status.queue_get_upload_mode.qsize() > 0 and status.nb_workers_get_upload_mode == 0:
status.nb_workers_get_upload_mode += 1
logger.debug("Job: get upload mode (no other worker getting upload mode)")
return (WorkerJob.GET_UPLOAD_MODE, _get_n(status.queue_get_upload_mode, MAX_NB_FILES_FETCH_UPLOAD_MODE))
# 7. Preupload LFS file if at least 1 file
# Skip if hf_transfer is enabled and there is already a worker preuploading LFS
elif status.queue_preupload_lfs.qsize() > 0 and (
status.nb_workers_preupload_lfs == 0 or not constants.HF_HUB_ENABLE_HF_TRANSFER
):
status.nb_workers_preupload_lfs += 1
logger.debug("Job: preupload LFS")
return (WorkerJob.PREUPLOAD_LFS, _get_one(status.queue_preupload_lfs))
# 8. Compute sha256 if at least 1 file
elif status.queue_sha256.qsize() > 0:
status.nb_workers_sha256 += 1
logger.debug("Job: sha256")
return (WorkerJob.SHA256, _get_one(status.queue_sha256))
# 9. Get upload mode if at least 1 file
elif status.queue_get_upload_mode.qsize() > 0:
status.nb_workers_get_upload_mode += 1
logger.debug("Job: get upload mode")
return (WorkerJob.GET_UPLOAD_MODE, _get_n(status.queue_get_upload_mode, MAX_NB_FILES_FETCH_UPLOAD_MODE))
# 10. Commit if at least 1 file and 1 min since last commit attempt
elif (
status.nb_workers_commit == 0
and status.queue_commit.qsize() > 0
and status.last_commit_attempt is not None
and time.time() - status.last_commit_attempt > 1 * 60
):
status.nb_workers_commit += 1
logger.debug("Job: commit (1 min since last commit attempt)")
return (WorkerJob.COMMIT, _get_n(status.queue_commit, status.target_chunk()))
# 11. Commit if at least 1 file all other queues are empty and all workers are waiting
# e.g. when it's the last commit
elif (
status.nb_workers_commit == 0
and status.queue_commit.qsize() > 0
and status.queue_sha256.qsize() == 0
and status.queue_get_upload_mode.qsize() == 0
and status.queue_preupload_lfs.qsize() == 0
and status.nb_workers_sha256 == 0
and status.nb_workers_get_upload_mode == 0
and status.nb_workers_preupload_lfs == 0
):
status.nb_workers_commit += 1
logger.debug("Job: commit")
return (WorkerJob.COMMIT, _get_n(status.queue_commit, status.target_chunk()))
# 12. If all queues are empty, exit
elif all(metadata.is_committed or metadata.should_ignore for _, metadata in status.items):
logger.info("All files have been processed! Exiting worker.")
return None
# 13. If no task is available, wait
else:
status.nb_workers_waiting += 1
logger.debug(f"No task available, waiting... ({WAITING_TIME_IF_NO_TASKS}s)")
return (WorkerJob.WAIT, [])
####################
# Atomic jobs (sha256, get_upload_mode, preupload_lfs, commit)
####################
def _compute_sha256(item: JOB_ITEM_T) -> None:
"""Compute sha256 of a file and save it in metadata."""
paths, metadata = item
if metadata.sha256 is None:
with paths.file_path.open("rb") as f:
metadata.sha256 = sha_fileobj(f).hex()
metadata.save(paths)
def _get_upload_mode(items: List[JOB_ITEM_T], api: "HfApi", repo_id: str, repo_type: str, revision: str) -> None:
"""Get upload mode for each file and update metadata.
Also receive info if the file should be ignored.
"""
additions = [_build_hacky_operation(item) for item in items]
_fetch_upload_modes(
additions=additions,
repo_type=repo_type,
repo_id=repo_id,
headers=api._build_hf_headers(),
revision=quote(revision, safe=""),
endpoint=api.endpoint,
)
for item, addition in zip(items, additions):
paths, metadata = item
metadata.upload_mode = addition._upload_mode
metadata.should_ignore = addition._should_ignore
metadata.remote_oid = addition._remote_oid
metadata.save(paths)
def _preupload_lfs(item: JOB_ITEM_T, api: "HfApi", repo_id: str, repo_type: str, revision: str) -> None:
"""Preupload LFS file and update metadata."""
paths, metadata = item
addition = _build_hacky_operation(item)
api.preupload_lfs_files(
repo_id=repo_id,
repo_type=repo_type,
revision=revision,
additions=[addition],
)
metadata.is_uploaded = True
metadata.save(paths)
def _commit(items: List[JOB_ITEM_T], api: "HfApi", repo_id: str, repo_type: str, revision: str) -> None:
"""Commit files to the repo."""
additions = [_build_hacky_operation(item) for item in items]
api.create_commit(
repo_id=repo_id,
repo_type=repo_type,
revision=revision,
operations=additions,
commit_message="Add files using upload-large-folder tool",
)
for paths, metadata in items:
metadata.is_committed = True
metadata.save(paths)
####################
# Hacks with CommitOperationAdd to bypass checks/sha256 calculation
####################
class HackyCommitOperationAdd(CommitOperationAdd):
def __post_init__(self) -> None:
if isinstance(self.path_or_fileobj, Path):
self.path_or_fileobj = str(self.path_or_fileobj)
def _build_hacky_operation(item: JOB_ITEM_T) -> HackyCommitOperationAdd:
paths, metadata = item
operation = HackyCommitOperationAdd(path_in_repo=paths.path_in_repo, path_or_fileobj=paths.file_path)
with paths.file_path.open("rb") as file:
sample = file.peek(512)[:512]
if metadata.sha256 is None:
raise ValueError("sha256 must have been computed by now!")
operation.upload_info = UploadInfo(sha256=bytes.fromhex(metadata.sha256), size=metadata.size, sample=sample)
operation._upload_mode = metadata.upload_mode # type: ignore[assignment]
operation._should_ignore = metadata.should_ignore
operation._remote_oid = metadata.remote_oid
return operation
####################
# Misc helpers
####################
def _get_one(queue: "queue.Queue[JOB_ITEM_T]") -> List[JOB_ITEM_T]:
return [queue.get()]
def _get_n(queue: "queue.Queue[JOB_ITEM_T]", n: int) -> List[JOB_ITEM_T]:
return [queue.get() for _ in range(min(queue.qsize(), n))]
def _print_overwrite(report: str) -> None:
"""Print a report, overwriting the previous lines.
Since tqdm in using `sys.stderr` to (re-)write progress bars, we need to use `sys.stdout`
to print the report.
Note: works well only if no other process is writing to `sys.stdout`!
"""
report += "\n"
# Get terminal width
terminal_width = shutil.get_terminal_size().columns
# Count number of lines that should be cleared
nb_lines = sum(len(line) // terminal_width + 1 for line in report.splitlines())
# Clear previous lines based on the number of lines in the report
for _ in range(nb_lines):
sys.stdout.write("\r\033[K") # Clear line
sys.stdout.write("\033[F") # Move cursor up one line
# Print the new report, filling remaining space with whitespace
sys.stdout.write(report)
sys.stdout.write(" " * (terminal_width - len(report.splitlines()[-1])))
sys.stdout.flush()
kronos-venv/lib/python3.12/site-packages/huggingface_hub/_webhooks_payload.py
+137
View File
@@ -0,0 +1,137 @@
# 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 data structures to parse the webhooks payload."""
from typing import List, Literal, Optional
from .utils import is_pydantic_available
if is_pydantic_available():
from pydantic import BaseModel
else:
# Define a dummy BaseModel to avoid import errors when pydantic is not installed
# Import error will be raised when trying to use the class
class BaseModel: # type: ignore [no-redef]
def __init__(self, *args, **kwargs) -> None:
raise ImportError(
"You must have `pydantic` installed to use `WebhookPayload`. This is an optional dependency that"
" should be installed separately. Please run `pip install --upgrade pydantic` and retry."
)
# This is an adaptation of the ReportV3 interface implemented in moon-landing. V0, V1 and V2 have been ignored as they
# are not in used anymore. To keep in sync when format is updated in
# https://github.com/huggingface/moon-landing/blob/main/server/lib/HFWebhooks.ts (internal link).
WebhookEvent_T = Literal[
"create",
"delete",
"move",
"update",
]
RepoChangeEvent_T = Literal[
"add",
"move",
"remove",
"update",
]
RepoType_T = Literal[
"dataset",
"model",
"space",
]
DiscussionStatus_T = Literal[
"closed",
"draft",
"open",
"merged",
]
SupportedWebhookVersion = Literal[3]
class ObjectId(BaseModel):
id: str
class WebhookPayloadUrl(BaseModel):
web: str
api: Optional[str] = None
class WebhookPayloadMovedTo(BaseModel):
name: str
owner: ObjectId
class WebhookPayloadWebhook(ObjectId):
version: SupportedWebhookVersion
class WebhookPayloadEvent(BaseModel):
action: WebhookEvent_T
scope: str
class WebhookPayloadDiscussionChanges(BaseModel):
base: str
mergeCommitId: Optional[str] = None
class WebhookPayloadComment(ObjectId):
author: ObjectId
hidden: bool
content: Optional[str] = None
url: WebhookPayloadUrl
class WebhookPayloadDiscussion(ObjectId):
num: int
author: ObjectId
url: WebhookPayloadUrl
title: str
isPullRequest: bool
status: DiscussionStatus_T
changes: Optional[WebhookPayloadDiscussionChanges] = None
pinned: Optional[bool] = None
class WebhookPayloadRepo(ObjectId):
owner: ObjectId
head_sha: Optional[str] = None
name: str
private: bool
subdomain: Optional[str] = None
tags: Optional[List[str]] = None
type: Literal["dataset", "model", "space"]
url: WebhookPayloadUrl
class WebhookPayloadUpdatedRef(BaseModel):
ref: str
oldSha: Optional[str] = None
newSha: Optional[str] = None
class WebhookPayload(BaseModel):
event: WebhookPayloadEvent
repo: WebhookPayloadRepo
discussion: Optional[WebhookPayloadDiscussion] = None
comment: Optional[WebhookPayloadComment] = None
webhook: WebhookPayloadWebhook
movedTo: Optional[WebhookPayloadMovedTo] = None
updatedRefs: Optional[List[WebhookPayloadUpdatedRef]] = None
kronos-venv/lib/python3.12/site-packages/huggingface_hub/_webhooks_server.py
+388
View File
@@ -0,0 +1,388 @@
# 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 `WebhooksServer` and `webhook_endpoint` to create a webhook server easily."""
import atexit
import inspect
import os
from functools import wraps
from typing import TYPE_CHECKING, Any, Callable, Dict, Optional
from .utils import experimental, is_fastapi_available, is_gradio_available
if TYPE_CHECKING:
import gradio as gr
from fastapi import Request
if is_fastapi_available():
from fastapi import FastAPI, Request
from fastapi.responses import JSONResponse
else:
# Will fail at runtime if FastAPI is not available
FastAPI = Request = JSONResponse = None # type: ignore [misc, assignment]
_global_app: Optional["WebhooksServer"] = None
_is_local = os.environ.get("SPACE_ID") is None
@experimental
class WebhooksServer:
"""
The [`WebhooksServer`] class lets you create an instance of a Gradio app that can receive Huggingface webhooks.
These webhooks can be registered using the [`~WebhooksServer.add_webhook`] decorator. Webhook endpoints are added to
the app as a POST endpoint to the FastAPI router. Once all the webhooks are registered, the `launch` method has to be
called to start the app.
It is recommended to accept [`WebhookPayload`] as the first argument of the webhook function. It is a Pydantic
model that contains all the information about the webhook event. The data will be parsed automatically for you.
Check out the [webhooks guide](../guides/webhooks_server) for a step-by-step tutorial on how to setup your
WebhooksServer and deploy it on a Space.
<Tip warning={true}>
`WebhooksServer` is experimental. Its API is subject to change in the future.
</Tip>
<Tip warning={true}>
You must have `gradio` installed to use `WebhooksServer` (`pip install --upgrade gradio`).
</Tip>
Args:
ui (`gradio.Blocks`, optional):
A Gradio UI instance to be used as the Space landing page. If `None`, a UI displaying instructions
about the configured webhooks is created.
webhook_secret (`str`, optional):
A secret key to verify incoming webhook requests. You can set this value to any secret you want as long as
you also configure it in your [webhooks settings panel](https://huggingface.co/settings/webhooks). You
can also set this value as the `WEBHOOK_SECRET` environment variable. If no secret is provided, the
webhook endpoints are opened without any security.
Example:
```python
import gradio as gr
from huggingface_hub import WebhooksServer, WebhookPayload
with gr.Blocks() as ui:
...
app = WebhooksServer(ui=ui, webhook_secret="my_secret_key")
@app.add_webhook("/say_hello")
async def hello(payload: WebhookPayload):
return {"message": "hello"}
app.launch()
```
"""
def __new__(cls, *args, **kwargs) -> "WebhooksServer":
if not is_gradio_available():
raise ImportError(
"You must have `gradio` installed to use `WebhooksServer`. Please run `pip install --upgrade gradio`"
" first."
)
if not is_fastapi_available():
raise ImportError(
"You must have `fastapi` installed to use `WebhooksServer`. Please run `pip install --upgrade fastapi`"
" first."
)
return super().__new__(cls)
def __init__(
self,
ui: Optional["gr.Blocks"] = None,
webhook_secret: Optional[str] = None,
) -> None:
self._ui = ui
self.webhook_secret = webhook_secret or os.getenv("WEBHOOK_SECRET")
self.registered_webhooks: Dict[str, Callable] = {}
_warn_on_empty_secret(self.webhook_secret)
def add_webhook(self, path: Optional[str] = None) -> Callable:
"""
Decorator to add a webhook to the [`WebhooksServer`] server.
Args:
path (`str`, optional):
The URL path to register the webhook function. If not provided, the function name will be used as the
path. In any case, all webhooks are registered under `/webhooks`.
Raises:
ValueError: If the provided path is already registered as a webhook.
Example:
```python
from huggingface_hub import WebhooksServer, WebhookPayload
app = WebhooksServer()
@app.add_webhook
async def trigger_training(payload: WebhookPayload):
if payload.repo.type == "dataset" and payload.event.action == "update":
# Trigger a training job if a dataset is updated
...
app.launch()
```
"""
# Usage: directly as decorator. Example: `@app.add_webhook`
if callable(path):
# If path is a function, it means it was used as a decorator without arguments
return self.add_webhook()(path)
# Usage: provide a path. Example: `@app.add_webhook(...)`
@wraps(FastAPI.post)
def _inner_post(*args, **kwargs):
func = args[0]
abs_path = f"/webhooks/{(path or func.__name__).strip('/')}"
if abs_path in self.registered_webhooks:
raise ValueError(f"Webhook {abs_path} already exists.")
self.registered_webhooks[abs_path] = func
return _inner_post
def launch(self, prevent_thread_lock: bool = False, **launch_kwargs: Any) -> None:
"""Launch the Gradio app and register webhooks to the underlying FastAPI server.
Input parameters are forwarded to Gradio when launching the app.
"""
ui = self._ui or self._get_default_ui()
# Start Gradio App
# - as non-blocking so that webhooks can be added afterwards
# - as shared if launch locally (to debug webhooks)
launch_kwargs.setdefault("share", _is_local)
self.fastapi_app, _, _ = ui.launch(prevent_thread_lock=True, **launch_kwargs)
# Register webhooks to FastAPI app
for path, func in self.registered_webhooks.items():
# Add secret check if required
if self.webhook_secret is not None:
func = _wrap_webhook_to_check_secret(func, webhook_secret=self.webhook_secret)
# Add route to FastAPI app
self.fastapi_app.post(path)(func)
# Print instructions and block main thread
space_host = os.environ.get("SPACE_HOST")
url = "https://" + space_host if space_host is not None else (ui.share_url or ui.local_url)
if url is None:
raise ValueError("Cannot find the URL of the app. Please provide a valid `ui` or update `gradio` version.")
url = url.strip("/")
message = "\nWebhooks are correctly setup and ready to use:"
message += "\n" + "\n".join(f" - POST {url}{webhook}" for webhook in self.registered_webhooks)
message += "\nGo to https://huggingface.co/settings/webhooks to setup your webhooks."
print(message)
if not prevent_thread_lock:
ui.block_thread()
def _get_default_ui(self) -> "gr.Blocks":
"""Default UI if not provided (lists webhooks and provides basic instructions)."""
import gradio as gr
with gr.Blocks() as ui:
gr.Markdown("# This is an app to process 🤗 Webhooks")
gr.Markdown(
"Webhooks are a foundation for MLOps-related features. They allow you to listen for new changes on"
" specific repos or to all repos belonging to particular set of users/organizations (not just your"
" repos, but any repo). Check out this [guide](https://huggingface.co/docs/hub/webhooks) to get to"
" know more about webhooks on the Huggingface Hub."
)
gr.Markdown(
f"{len(self.registered_webhooks)} webhook(s) are registered:"
+ "\n\n"
+ "\n ".join(
f"- [{webhook_path}]({_get_webhook_doc_url(webhook.__name__, webhook_path)})"
for webhook_path, webhook in self.registered_webhooks.items()
)
)
gr.Markdown(
"Go to https://huggingface.co/settings/webhooks to setup your webhooks."
+ "\nYou app is running locally. Please look at the logs to check the full URL you need to set."
if _is_local
else (
"\nThis app is running on a Space. You can find the corresponding URL in the options menu"
" (top-right) > 'Embed the Space'. The URL looks like 'https://{username}-{repo_name}.hf.space'."
)
)
return ui
@experimental
def webhook_endpoint(path: Optional[str] = None) -> Callable:
"""Decorator to start a [`WebhooksServer`] and register the decorated function as a webhook endpoint.
This is a helper to get started quickly. If you need more flexibility (custom landing page or webhook secret),
you can use [`WebhooksServer`] directly. You can register multiple webhook endpoints (to the same server) by using
this decorator multiple times.
Check out the [webhooks guide](../guides/webhooks_server) for a step-by-step tutorial on how to setup your
server and deploy it on a Space.
<Tip warning={true}>
`webhook_endpoint` is experimental. Its API is subject to change in the future.
</Tip>
<Tip warning={true}>
You must have `gradio` installed to use `webhook_endpoint` (`pip install --upgrade gradio`).
</Tip>
Args:
path (`str`, optional):
The URL path to register the webhook function. If not provided, the function name will be used as the path.
In any case, all webhooks are registered under `/webhooks`.
Examples:
The default usage is to register a function as a webhook endpoint. The function name will be used as the path.
The server will be started automatically at exit (i.e. at the end of the script).
```python
from huggingface_hub import webhook_endpoint, WebhookPayload
@webhook_endpoint
async def trigger_training(payload: WebhookPayload):
if payload.repo.type == "dataset" and payload.event.action == "update":
# Trigger a training job if a dataset is updated
...
# Server is automatically started at the end of the script.
```
Advanced usage: register a function as a webhook endpoint and start the server manually. This is useful if you
are running it in a notebook.
```python
from huggingface_hub import webhook_endpoint, WebhookPayload
@webhook_endpoint
async def trigger_training(payload: WebhookPayload):
if payload.repo.type == "dataset" and payload.event.action == "update":
# Trigger a training job if a dataset is updated
...
# Start the server manually
trigger_training.launch()
```
"""
if callable(path):
# If path is a function, it means it was used as a decorator without arguments
return webhook_endpoint()(path)
@wraps(WebhooksServer.add_webhook)
def _inner(func: Callable) -> Callable:
app = _get_global_app()
app.add_webhook(path)(func)
if len(app.registered_webhooks) == 1:
# Register `app.launch` to run at exit (only once)
atexit.register(app.launch)
@wraps(app.launch)
def _launch_now():
# Run the app directly (without waiting atexit)
atexit.unregister(app.launch)
app.launch()
func.launch = _launch_now # type: ignore
return func
return _inner
def _get_global_app() -> WebhooksServer:
global _global_app
if _global_app is None:
_global_app = WebhooksServer()
return _global_app
def _warn_on_empty_secret(webhook_secret: Optional[str]) -> None:
if webhook_secret is None:
print("Webhook secret is not defined. This means your webhook endpoints will be open to everyone.")
print(
"To add a secret, set `WEBHOOK_SECRET` as environment variable or pass it at initialization: "
"\n\t`app = WebhooksServer(webhook_secret='my_secret', ...)`"
)
print(
"For more details about webhook secrets, please refer to"
" https://huggingface.co/docs/hub/webhooks#webhook-secret."
)
else:
print("Webhook secret is correctly defined.")
def _get_webhook_doc_url(webhook_name: str, webhook_path: str) -> str:
"""Returns the anchor to a given webhook in the docs (experimental)"""
return "/docs#/default/" + webhook_name + webhook_path.replace("/", "_") + "_post"
def _wrap_webhook_to_check_secret(func: Callable, webhook_secret: str) -> Callable:
"""Wraps a webhook function to check the webhook secret before calling the function.
This is a hacky way to add the `request` parameter to the function signature. Since FastAPI based itself on route
parameters to inject the values to the function, we need to hack the function signature to retrieve the `Request`
object (and hence the headers). A far cleaner solution would be to use a middleware. However, since
`fastapi==0.90.1`, a middleware cannot be added once the app has started. And since the FastAPI app is started by
Gradio internals (and not by us), we cannot add a middleware.
This method is called only when a secret has been defined by the user. If a request is sent without the
"x-webhook-secret", the function will return a 401 error (unauthorized). If the header is sent but is incorrect,
the function will return a 403 error (forbidden).
Inspired by https://stackoverflow.com/a/33112180.
"""
initial_sig = inspect.signature(func)
@wraps(func)
async def _protected_func(request: Request, **kwargs):
request_secret = request.headers.get("x-webhook-secret")
if request_secret is None:
return JSONResponse({"error": "x-webhook-secret header not set."}, status_code=401)
if request_secret != webhook_secret:
return JSONResponse({"error": "Invalid webhook secret."}, status_code=403)
# Inject `request` in kwargs if required
if "request" in initial_sig.parameters:
kwargs["request"] = request
# Handle both sync and async routes
if inspect.iscoroutinefunction(func):
return await func(**kwargs)
else:
return func(**kwargs)
# Update signature to include request
if "request" not in initial_sig.parameters:
_protected_func.__signature__ = initial_sig.replace( # type: ignore
parameters=(
inspect.Parameter(name="request", kind=inspect.Parameter.POSITIONAL_OR_KEYWORD, annotation=Request),
)
+ tuple(initial_sig.parameters.values())
)
# Return protected route
return _protected_func
kronos-venv/lib/python3.12/site-packages/huggingface_hub/commands/__init__.py
+27
View File
@@ -0,0 +1,27 @@
# Copyright 2020 The HuggingFace Team. All rights reserved.
#
# 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.
from abc import ABC, abstractmethod
from argparse import _SubParsersAction
class BaseHuggingfaceCLICommand(ABC):
@staticmethod
@abstractmethod
def register_subcommand(parser: _SubParsersAction):
raise NotImplementedError()
@abstractmethod
def run(self):
raise NotImplementedError()
kronos-venv/lib/python3.12/site-packages/huggingface_hub/commands/_cli_utils.py
+69
View File
@@ -0,0 +1,69 @@
# Copyright 2022 The HuggingFace Team. All rights reserved.
#
# 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 a utility for good-looking prints."""
import os
from typing import List, Union
class ANSI:
"""
Helper for en.wikipedia.org/wiki/ANSI_escape_code
"""
_bold = "\u001b[1m"
_gray = "\u001b[90m"
_red = "\u001b[31m"
_reset = "\u001b[0m"
_yellow = "\u001b[33m"
@classmethod
def bold(cls, s: str) -> str:
return cls._format(s, cls._bold)
@classmethod
def gray(cls, s: str) -> str:
return cls._format(s, cls._gray)
@classmethod
def red(cls, s: str) -> str:
return cls._format(s, cls._bold + cls._red)
@classmethod
def yellow(cls, s: str) -> str:
return cls._format(s, cls._yellow)
@classmethod
def _format(cls, s: str, code: str) -> str:
if os.environ.get("NO_COLOR"):
# See https://no-color.org/
return s
return f"{code}{s}{cls._reset}"
def tabulate(rows: List[List[Union[str, int]]], headers: List[str]) -> str:
"""
Inspired by:
- stackoverflow.com/a/8356620/593036
- stackoverflow.com/questions/9535954/printing-lists-as-tabular-data
"""
col_widths = [max(len(str(x)) for x in col) for col in zip(*rows, headers)]
row_format = ("{{:{}}} " * len(headers)).format(*col_widths)
lines = []
lines.append(row_format.format(*headers))
lines.append(row_format.format(*["-" * w for w in col_widths]))
for row in rows:
lines.append(row_format.format(*row))
return "\n".join(lines)
kronos-venv/lib/python3.12/site-packages/huggingface_hub/commands/delete_cache.py
+474
View File
@@ -0,0 +1,474 @@
# coding=utf-8
# Copyright 2022-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 command to delete some revisions from the HF cache directory.
Usage:
huggingface-cli delete-cache
huggingface-cli delete-cache --disable-tui
huggingface-cli delete-cache --dir ~/.cache/huggingface/hub
huggingface-cli delete-cache --sort=size
NOTE:
This command is based on `InquirerPy` to build the multiselect menu in the terminal.
This dependency has to be installed with `pip install huggingface_hub[cli]`. Since
we want to avoid as much as possible cross-platform issues, I chose a library that
is built on top of `python-prompt-toolkit` which seems to be a reference in terminal
GUI (actively maintained on both Unix and Windows, 7.9k stars).
For the moment, the TUI feature is in beta.
See:
- https://github.com/kazhala/InquirerPy
- https://inquirerpy.readthedocs.io/en/latest/
- https://github.com/prompt-toolkit/python-prompt-toolkit
Other solutions could have been:
- `simple_term_menu`: would be good as well for our use case but some issues suggest
that Windows is less supported.
See: https://github.com/IngoMeyer441/simple-term-menu
- `PyInquirer`: very similar to `InquirerPy` but older and not maintained anymore.
In particular, no support of Python3.10.
See: https://github.com/CITGuru/PyInquirer
- `pick` (or `pickpack`): easy to use and flexible but built on top of Python's
standard library `curses` that is specific to Unix (not implemented on Windows).
See https://github.com/wong2/pick and https://github.com/anafvana/pickpack.
- `inquirer`: lot of traction (700 stars) but explicitly states "experimental
support of Windows". Not built on top of `python-prompt-toolkit`.
See https://github.com/magmax/python-inquirer
TODO: add support for `huggingface-cli delete-cache aaaaaa bbbbbb cccccc (...)` ?
TODO: add "--keep-last" arg to delete revisions that are not on `main` ref
TODO: add "--filter" arg to filter repositories by name ?
TODO: add "--limit" arg to limit to X repos ?
TODO: add "-y" arg for immediate deletion ?
See discussions in https://github.com/huggingface/huggingface_hub/issues/1025.
"""
import os
from argparse import Namespace, _SubParsersAction
from functools import wraps
from tempfile import mkstemp
from typing import Any, Callable, Iterable, List, Literal, Optional, Union
from ..utils import CachedRepoInfo, CachedRevisionInfo, HFCacheInfo, scan_cache_dir
from . import BaseHuggingfaceCLICommand
from ._cli_utils import ANSI
try:
from InquirerPy import inquirer
from InquirerPy.base.control import Choice
from InquirerPy.separator import Separator
_inquirer_py_available = True
except ImportError:
_inquirer_py_available = False
SortingOption_T = Literal["alphabetical", "lastUpdated", "lastUsed", "size"]
def require_inquirer_py(fn: Callable) -> Callable:
"""Decorator to flag methods that require `InquirerPy`."""
# TODO: refactor this + imports in a unified pattern across codebase
@wraps(fn)
def _inner(*args, **kwargs):
if not _inquirer_py_available:
raise ImportError(
"The `delete-cache` command requires extra dependencies to work with"
" the TUI.\nPlease run `pip install huggingface_hub[cli]` to install"
" them.\nOtherwise, disable TUI using the `--disable-tui` flag."
)
return fn(*args, **kwargs)
return _inner
# Possibility for the user to cancel deletion
_CANCEL_DELETION_STR = "CANCEL_DELETION"
class DeleteCacheCommand(BaseHuggingfaceCLICommand):
@staticmethod
def register_subcommand(parser: _SubParsersAction):
delete_cache_parser = parser.add_parser("delete-cache", help="Delete revisions from the cache directory.")
delete_cache_parser.add_argument(
"--dir",
type=str,
default=None,
help="cache directory (optional). Default to the default HuggingFace cache.",
)
delete_cache_parser.add_argument(
"--disable-tui",
action="store_true",
help=(
"Disable Terminal User Interface (TUI) mode. Useful if your"
" platform/terminal doesn't support the multiselect menu."
),
)
delete_cache_parser.add_argument(
"--sort",
nargs="?",
choices=["alphabetical", "lastUpdated", "lastUsed", "size"],
help=(
"Sort repositories by the specified criteria. Options: "
"'alphabetical' (A-Z), "
"'lastUpdated' (newest first), "
"'lastUsed' (most recent first), "
"'size' (largest first)."
),
)
delete_cache_parser.set_defaults(func=DeleteCacheCommand)
def __init__(self, args: Namespace) -> None:
self.cache_dir: Optional[str] = args.dir
self.disable_tui: bool = args.disable_tui
self.sort_by: Optional[SortingOption_T] = args.sort
def run(self):
"""Run `delete-cache` command with or without TUI."""
# Scan cache directory
hf_cache_info = scan_cache_dir(self.cache_dir)
# Manual review from the user
if self.disable_tui:
selected_hashes = _manual_review_no_tui(hf_cache_info, preselected=[], sort_by=self.sort_by)
else:
selected_hashes = _manual_review_tui(hf_cache_info, preselected=[], sort_by=self.sort_by)
# If deletion is not cancelled
if len(selected_hashes) > 0 and _CANCEL_DELETION_STR not in selected_hashes:
confirm_message = _get_expectations_str(hf_cache_info, selected_hashes) + " Confirm deletion ?"
# Confirm deletion
if self.disable_tui:
confirmed = _ask_for_confirmation_no_tui(confirm_message)
else:
confirmed = _ask_for_confirmation_tui(confirm_message)
# Deletion is confirmed
if confirmed:
strategy = hf_cache_info.delete_revisions(*selected_hashes)
print("Start deletion.")
strategy.execute()
print(
f"Done. Deleted {len(strategy.repos)} repo(s) and"
f" {len(strategy.snapshots)} revision(s) for a total of"
f" {strategy.expected_freed_size_str}."
)
return
# Deletion is cancelled
print("Deletion is cancelled. Do nothing.")
def _get_repo_sorting_key(repo: CachedRepoInfo, sort_by: Optional[SortingOption_T] = None):
if sort_by == "alphabetical":
return (repo.repo_type, repo.repo_id.lower()) # by type then name
elif sort_by == "lastUpdated":
return -max(rev.last_modified for rev in repo.revisions) # newest first
elif sort_by == "lastUsed":
return -repo.last_accessed # most recently used first
elif sort_by == "size":
return -repo.size_on_disk # largest first
else:
return (repo.repo_type, repo.repo_id) # default stable order
@require_inquirer_py
def _manual_review_tui(
hf_cache_info: HFCacheInfo,
preselected: List[str],
sort_by: Optional[SortingOption_T] = None,
) -> List[str]:
"""Ask the user for a manual review of the revisions to delete.
Displays a multi-select menu in the terminal (TUI).
"""
# Define multiselect list
choices = _get_tui_choices_from_scan(
repos=hf_cache_info.repos,
preselected=preselected,
sort_by=sort_by,
)
checkbox = inquirer.checkbox(
message="Select revisions to delete:",
choices=choices, # List of revisions with some pre-selection
cycle=False, # No loop between top and bottom
height=100, # Large list if possible
# We use the instruction to display to the user the expected effect of the
# deletion.
instruction=_get_expectations_str(
hf_cache_info,
selected_hashes=[c.value for c in choices if isinstance(c, Choice) and c.enabled],
),
# We use the long instruction to should keybindings instructions to the user
long_instruction="Press <space> to select, <enter> to validate and <ctrl+c> to quit without modification.",
# Message that is displayed once the user validates its selection.
transformer=lambda result: f"{len(result)} revision(s) selected.",
)
# Add a callback to update the information line when a revision is
# selected/unselected
def _update_expectations(_) -> None:
# Hacky way to dynamically set an instruction message to the checkbox when
# a revision hash is selected/unselected.
checkbox._instruction = _get_expectations_str(
hf_cache_info,
selected_hashes=[choice["value"] for choice in checkbox.content_control.choices if choice["enabled"]],
)
checkbox.kb_func_lookup["toggle"].append({"func": _update_expectations})
# Finally display the form to the user.
try:
return checkbox.execute()
except KeyboardInterrupt:
return [] # Quit without deletion
@require_inquirer_py
def _ask_for_confirmation_tui(message: str, default: bool = True) -> bool:
"""Ask for confirmation using Inquirer."""
return inquirer.confirm(message, default=default).execute()
def _get_tui_choices_from_scan(
repos: Iterable[CachedRepoInfo],
preselected: List[str],
sort_by: Optional[SortingOption_T] = None,
) -> List:
"""Build a list of choices from the scanned repos.
Args:
repos (*Iterable[`CachedRepoInfo`]*):
List of scanned repos on which we want to delete revisions.
preselected (*List[`str`]*):
List of revision hashes that will be preselected.
sort_by (*Optional[SortingOption_T]*):
Sorting direction. Choices: "alphabetical", "lastUpdated", "lastUsed", "size".
Return:
The list of choices to pass to `inquirer.checkbox`.
"""
choices: List[Union[Choice, Separator]] = []
# First choice is to cancel the deletion
choices.append(
Choice(
_CANCEL_DELETION_STR,
name="None of the following (if selected, nothing will be deleted).",
enabled=False,
)
)
# Sort repos based on specified criteria
sorted_repos = sorted(repos, key=lambda repo: _get_repo_sorting_key(repo, sort_by))
for repo in sorted_repos:
# Repo as separator
choices.append(
Separator(
f"\n{repo.repo_type.capitalize()} {repo.repo_id} ({repo.size_on_disk_str},"
f" used {repo.last_accessed_str})"
)
)
for revision in sorted(repo.revisions, key=_revision_sorting_order):
# Revision as choice
choices.append(
Choice(
revision.commit_hash,
name=(
f"{revision.commit_hash[:8]}:"
f" {', '.join(sorted(revision.refs)) or '(detached)'} #"
f" modified {revision.last_modified_str}"
),
enabled=revision.commit_hash in preselected,
)
)
# Return choices
return choices
def _manual_review_no_tui(
hf_cache_info: HFCacheInfo,
preselected: List[str],
sort_by: Optional[SortingOption_T] = None,
) -> List[str]:
"""Ask the user for a manual review of the revisions to delete.
Used when TUI is disabled. Manual review happens in a separate tmp file that the
user can manually edit.
"""
# 1. Generate temporary file with delete commands.
fd, tmp_path = mkstemp(suffix=".txt") # suffix to make it easier to find by editors
os.close(fd)
lines = []
sorted_repos = sorted(hf_cache_info.repos, key=lambda repo: _get_repo_sorting_key(repo, sort_by))
for repo in sorted_repos:
lines.append(
f"\n# {repo.repo_type.capitalize()} {repo.repo_id} ({repo.size_on_disk_str},"
f" used {repo.last_accessed_str})"
)
for revision in sorted(repo.revisions, key=_revision_sorting_order):
lines.append(
# Deselect by prepending a '#'
f"{'' if revision.commit_hash in preselected else '#'} "
f" {revision.commit_hash} # Refs:"
# Print `refs` as comment on same line
f" {', '.join(sorted(revision.refs)) or '(detached)'} # modified"
# Print `last_modified` as comment on same line
f" {revision.last_modified_str}"
)
with open(tmp_path, "w") as f:
f.write(_MANUAL_REVIEW_NO_TUI_INSTRUCTIONS)
f.write("\n".join(lines))
# 2. Prompt instructions to user.
instructions = f"""
TUI is disabled. In order to select which revisions you want to delete, please edit
the following file using the text editor of your choice. Instructions for manual
editing are located at the beginning of the file. Edit the file, save it and confirm
to continue.
File to edit: {ANSI.bold(tmp_path)}
"""
print("\n".join(line.strip() for line in instructions.strip().split("\n")))
# 3. Wait for user confirmation.
while True:
selected_hashes = _read_manual_review_tmp_file(tmp_path)
if _ask_for_confirmation_no_tui(
_get_expectations_str(hf_cache_info, selected_hashes) + " Continue ?",
default=False,
):
break
# 4. Return selected_hashes sorted to maintain stable order
os.remove(tmp_path)
return sorted(selected_hashes) # Sort to maintain stable order
def _ask_for_confirmation_no_tui(message: str, default: bool = True) -> bool:
"""Ask for confirmation using pure-python."""
YES = ("y", "yes", "1")
NO = ("n", "no", "0")
DEFAULT = ""
ALL = YES + NO + (DEFAULT,)
full_message = message + (" (Y/n) " if default else " (y/N) ")
while True:
answer = input(full_message).lower()
if answer == DEFAULT:
return default
if answer in YES:
return True
if answer in NO:
return False
print(f"Invalid input. Must be one of {ALL}")
def _get_expectations_str(hf_cache_info: HFCacheInfo, selected_hashes: List[str]) -> str:
"""Format a string to display to the user how much space would be saved.
Example:
```
>>> _get_expectations_str(hf_cache_info, selected_hashes)
'7 revisions selected counting for 4.3G.'
```
"""
if _CANCEL_DELETION_STR in selected_hashes:
return "Nothing will be deleted."
strategy = hf_cache_info.delete_revisions(*selected_hashes)
return f"{len(selected_hashes)} revisions selected counting for {strategy.expected_freed_size_str}."
def _read_manual_review_tmp_file(tmp_path: str) -> List[str]:
"""Read the manually reviewed instruction file and return a list of revision hash.
Example:
```txt
# This is the tmp file content
###
# Commented out line
123456789 # revision hash
# Something else
# a_newer_hash # 2 days ago
an_older_hash # 3 days ago
```
```py
>>> _read_manual_review_tmp_file(tmp_path)
['123456789', 'an_older_hash']
```
"""
with open(tmp_path) as f:
content = f.read()
# Split lines
lines = [line.strip() for line in content.split("\n")]
# Filter commented lines
selected_lines = [line for line in lines if not line.startswith("#")]
# Select only before comment
selected_hashes = [line.split("#")[0].strip() for line in selected_lines]
# Return revision hashes
return [hash for hash in selected_hashes if len(hash) > 0]
_MANUAL_REVIEW_NO_TUI_INSTRUCTIONS = f"""
# INSTRUCTIONS
# ------------
# This is a temporary file created by running `huggingface-cli delete-cache` with the
# `--disable-tui` option. It contains a set of revisions that can be deleted from your
# local cache directory.
#
# Please manually review the revisions you want to delete:
# - Revision hashes can be commented out with '#'.
# - Only non-commented revisions in this file will be deleted.
# - Revision hashes that are removed from this file are ignored as well.
# - If `{_CANCEL_DELETION_STR}` line is uncommented, the all cache deletion is cancelled and
# no changes will be applied.
#
# Once you've manually reviewed this file, please confirm deletion in the terminal. This
# file will be automatically removed once done.
# ------------
# KILL SWITCH
# ------------
# Un-comment following line to completely cancel the deletion process
# {_CANCEL_DELETION_STR}
# ------------
# REVISIONS
# ------------
""".strip()
def _revision_sorting_order(revision: CachedRevisionInfo) -> Any:
# Sort by last modified (oldest first)
return revision.last_modified
kronos-venv/lib/python3.12/site-packages/huggingface_hub/commands/download.py
+200
View File
@@ -0,0 +1,200 @@
# 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 command to download files from the Hub with the CLI.
Usage:
huggingface-cli download --help
# Download file
huggingface-cli download gpt2 config.json
# Download entire repo
huggingface-cli download fffiloni/zeroscope --repo-type=space --revision=refs/pr/78
# Download repo with filters
huggingface-cli download gpt2 --include="*.safetensors"
# Download with token
huggingface-cli download Wauplin/private-model --token=hf_***
# Download quietly (no progress bar, no warnings, only the returned path)
huggingface-cli download gpt2 config.json --quiet
# Download to local dir
huggingface-cli download gpt2 --local-dir=./models/gpt2
"""
import warnings
from argparse import Namespace, _SubParsersAction
from typing import List, Optional
from huggingface_hub import logging
from huggingface_hub._snapshot_download import snapshot_download
from huggingface_hub.commands import BaseHuggingfaceCLICommand
from huggingface_hub.file_download import hf_hub_download
from huggingface_hub.utils import disable_progress_bars, enable_progress_bars
logger = logging.get_logger(__name__)
class DownloadCommand(BaseHuggingfaceCLICommand):
@staticmethod
def register_subcommand(parser: _SubParsersAction):
download_parser = parser.add_parser("download", help="Download files from the Hub")
download_parser.add_argument(
"repo_id", type=str, help="ID of the repo to download from (e.g. `username/repo-name`)."
)
download_parser.add_argument(
"filenames", type=str, nargs="*", help="Files to download (e.g. `config.json`, `data/metadata.jsonl`)."
)
download_parser.add_argument(
"--repo-type",
choices=["model", "dataset", "space"],
default="model",
help="Type of repo to download from (defaults to 'model').",
)
download_parser.add_argument(
"--revision",
type=str,
help="An optional Git revision id which can be a branch name, a tag, or a commit hash.",
)
download_parser.add_argument(
"--include", nargs="*", type=str, help="Glob patterns to match files to download."
)
download_parser.add_argument(
"--exclude", nargs="*", type=str, help="Glob patterns to exclude from files to download."
)
download_parser.add_argument(
"--cache-dir", type=str, help="Path to the directory where to save the downloaded files."
)
download_parser.add_argument(
"--local-dir",
type=str,
help=(
"If set, the downloaded file will be placed under this directory. Check out"
" https://huggingface.co/docs/huggingface_hub/guides/download#download-files-to-local-folder for more"
" details."
),
)
download_parser.add_argument(
"--local-dir-use-symlinks",
choices=["auto", "True", "False"],
help=("Deprecated and ignored. Downloading to a local directory does not use symlinks anymore."),
)
download_parser.add_argument(
"--force-download",
action="store_true",
help="If True, the files will be downloaded even if they are already cached.",
)
download_parser.add_argument(
"--resume-download",
action="store_true",
help="Deprecated and ignored. Downloading a file to local dir always attempts to resume previously interrupted downloads (unless hf-transfer is enabled).",
)
download_parser.add_argument(
"--token", type=str, help="A User Access Token generated from https://huggingface.co/settings/tokens"
)
download_parser.add_argument(
"--quiet",
action="store_true",
help="If True, progress bars are disabled and only the path to the download files is printed.",
)
download_parser.add_argument(
"--max-workers",
type=int,
default=8,
help="Maximum number of workers to use for downloading files. Default is 8.",
)
download_parser.set_defaults(func=DownloadCommand)
def __init__(self, args: Namespace) -> None:
self.token = args.token
self.repo_id: str = args.repo_id
self.filenames: List[str] = args.filenames
self.repo_type: str = args.repo_type
self.revision: Optional[str] = args.revision
self.include: Optional[List[str]] = args.include
self.exclude: Optional[List[str]] = args.exclude
self.cache_dir: Optional[str] = args.cache_dir
self.local_dir: Optional[str] = args.local_dir
self.force_download: bool = args.force_download
self.resume_download: Optional[bool] = args.resume_download or None
self.quiet: bool = args.quiet
self.max_workers: int = args.max_workers
if args.local_dir_use_symlinks is not None:
warnings.warn(
"Ignoring --local-dir-use-symlinks. Downloading to a local directory does not use symlinks anymore.",
FutureWarning,
)
def run(self) -> None:
if self.quiet:
disable_progress_bars()
with warnings.catch_warnings():
warnings.simplefilter("ignore")
print(self._download()) # Print path to downloaded files
enable_progress_bars()
else:
logging.set_verbosity_info()
print(self._download()) # Print path to downloaded files
logging.set_verbosity_warning()
def _download(self) -> str:
# Warn user if patterns are ignored
if len(self.filenames) > 0:
if self.include is not None and len(self.include) > 0:
warnings.warn("Ignoring `--include` since filenames have being explicitly set.")
if self.exclude is not None and len(self.exclude) > 0:
warnings.warn("Ignoring `--exclude` since filenames have being explicitly set.")
# Single file to download: use `hf_hub_download`
if len(self.filenames) == 1:
return hf_hub_download(
repo_id=self.repo_id,
repo_type=self.repo_type,
revision=self.revision,
filename=self.filenames[0],
cache_dir=self.cache_dir,
resume_download=self.resume_download,
force_download=self.force_download,
token=self.token,
local_dir=self.local_dir,
library_name="huggingface-cli",
)
# Otherwise: use `snapshot_download` to ensure all files comes from same revision
elif len(self.filenames) == 0:
allow_patterns = self.include
ignore_patterns = self.exclude
else:
allow_patterns = self.filenames
ignore_patterns = None
return snapshot_download(
repo_id=self.repo_id,
repo_type=self.repo_type,
revision=self.revision,
allow_patterns=allow_patterns,
ignore_patterns=ignore_patterns,
resume_download=self.resume_download,
force_download=self.force_download,
cache_dir=self.cache_dir,
token=self.token,
local_dir=self.local_dir,
library_name="huggingface-cli",
max_workers=self.max_workers,
)
kronos-venv/lib/python3.12/site-packages/huggingface_hub/commands/env.py
+36
View File
@@ -0,0 +1,36 @@
# Copyright 2022 The HuggingFace Team. All rights reserved.
#
# 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 command to print information about the environment.
Usage:
huggingface-cli env
"""
from argparse import _SubParsersAction
from ..utils import dump_environment_info
from . import BaseHuggingfaceCLICommand
class EnvironmentCommand(BaseHuggingfaceCLICommand):
def __init__(self, args):
self.args = args
@staticmethod
def register_subcommand(parser: _SubParsersAction):
env_parser = parser.add_parser("env", help="Print information about the environment.")
env_parser.set_defaults(func=EnvironmentCommand)
def run(self) -> None:
dump_environment_info()
kronos-venv/lib/python3.12/site-packages/huggingface_hub/commands/huggingface_cli.py
+63
View File
@@ -0,0 +1,63 @@
# Copyright 2020 The HuggingFace Team. All rights reserved.
#
# 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.
from argparse import ArgumentParser
from huggingface_hub.commands.delete_cache import DeleteCacheCommand
from huggingface_hub.commands.download import DownloadCommand
from huggingface_hub.commands.env import EnvironmentCommand
from huggingface_hub.commands.lfs import LfsCommands
from huggingface_hub.commands.repo import RepoCommands
from huggingface_hub.commands.repo_files import RepoFilesCommand
from huggingface_hub.commands.scan_cache import ScanCacheCommand
from huggingface_hub.commands.tag import TagCommands
from huggingface_hub.commands.upload import UploadCommand
from huggingface_hub.commands.upload_large_folder import UploadLargeFolderCommand
from huggingface_hub.commands.user import UserCommands
from huggingface_hub.commands.version import VersionCommand
def main():
parser = ArgumentParser("huggingface-cli", usage="huggingface-cli <command> [<args>]")
commands_parser = parser.add_subparsers(help="huggingface-cli command helpers")
# Register commands
DownloadCommand.register_subcommand(commands_parser)
UploadCommand.register_subcommand(commands_parser)
RepoFilesCommand.register_subcommand(commands_parser)
EnvironmentCommand.register_subcommand(commands_parser)
UserCommands.register_subcommand(commands_parser)
RepoCommands.register_subcommand(commands_parser)
LfsCommands.register_subcommand(commands_parser)
ScanCacheCommand.register_subcommand(commands_parser)
DeleteCacheCommand.register_subcommand(commands_parser)
TagCommands.register_subcommand(commands_parser)
VersionCommand.register_subcommand(commands_parser)
# Experimental
UploadLargeFolderCommand.register_subcommand(commands_parser)
# Let's go
args = parser.parse_args()
if not hasattr(args, "func"):
parser.print_help()
exit(1)
# Run
service = args.func(args)
service.run()
if __name__ == "__main__":
main()
kronos-venv/lib/python3.12/site-packages/huggingface_hub/commands/lfs.py
+200
View File
@@ -0,0 +1,200 @@
"""
Implementation of a custom transfer agent for the transfer type "multipart" for
git-lfs.
Inspired by:
github.com/cbartz/git-lfs-swift-transfer-agent/blob/master/git_lfs_swift_transfer.py
Spec is: github.com/git-lfs/git-lfs/blob/master/docs/custom-transfers.md
To launch debugger while developing:
``` [lfs "customtransfer.multipart"]
path = /path/to/huggingface_hub/.env/bin/python args = -m debugpy --listen 5678
--wait-for-client
/path/to/huggingface_hub/src/huggingface_hub/commands/huggingface_cli.py
lfs-multipart-upload ```"""
import json
import os
import subprocess
import sys
from argparse import _SubParsersAction
from typing import Dict, List, Optional
from huggingface_hub.commands import BaseHuggingfaceCLICommand
from huggingface_hub.lfs import LFS_MULTIPART_UPLOAD_COMMAND
from ..utils import get_session, hf_raise_for_status, logging
from ..utils._lfs import SliceFileObj
logger = logging.get_logger(__name__)
class LfsCommands(BaseHuggingfaceCLICommand):
"""
Implementation of a custom transfer agent for the transfer type "multipart"
for git-lfs. This lets users upload large files >5GB 🔥. Spec for LFS custom
transfer agent is:
https://github.com/git-lfs/git-lfs/blob/master/docs/custom-transfers.md
This introduces two commands to the CLI:
1. $ huggingface-cli lfs-enable-largefiles
This should be executed once for each model repo that contains a model file
>5GB. It's documented in the error message you get if you just try to git
push a 5GB file without having enabled it before.
2. $ huggingface-cli lfs-multipart-upload
This command is called by lfs directly and is not meant to be called by the
user.
"""
@staticmethod
def register_subcommand(parser: _SubParsersAction):
enable_parser = parser.add_parser(
"lfs-enable-largefiles", help="Configure your repository to enable upload of files > 5GB."
)
enable_parser.add_argument("path", type=str, help="Local path to repository you want to configure.")
enable_parser.set_defaults(func=lambda args: LfsEnableCommand(args))
# Command will get called by git-lfs, do not call it directly.
upload_parser = parser.add_parser(LFS_MULTIPART_UPLOAD_COMMAND, add_help=False)
upload_parser.set_defaults(func=lambda args: LfsUploadCommand(args))
class LfsEnableCommand:
def __init__(self, args):
self.args = args
def run(self):
local_path = os.path.abspath(self.args.path)
if not os.path.isdir(local_path):
print("This does not look like a valid git repo.")
exit(1)
subprocess.run(
"git config lfs.customtransfer.multipart.path huggingface-cli".split(),
check=True,
cwd=local_path,
)
subprocess.run(
f"git config lfs.customtransfer.multipart.args {LFS_MULTIPART_UPLOAD_COMMAND}".split(),
check=True,
cwd=local_path,
)
print("Local repo set up for largefiles")
def write_msg(msg: Dict):
"""Write out the message in Line delimited JSON."""
msg_str = json.dumps(msg) + "\n"
sys.stdout.write(msg_str)
sys.stdout.flush()
def read_msg() -> Optional[Dict]:
"""Read Line delimited JSON from stdin."""
msg = json.loads(sys.stdin.readline().strip())
if "terminate" in (msg.get("type"), msg.get("event")):
# terminate message received
return None
if msg.get("event") not in ("download", "upload"):
logger.critical("Received unexpected message")
sys.exit(1)
return msg
class LfsUploadCommand:
def __init__(self, args) -> None:
self.args = args
def run(self) -> None:
# Immediately after invoking a custom transfer process, git-lfs
# sends initiation data to the process over stdin.
# This tells the process useful information about the configuration.
init_msg = json.loads(sys.stdin.readline().strip())
if not (init_msg.get("event") == "init" and init_msg.get("operation") == "upload"):
write_msg({"error": {"code": 32, "message": "Wrong lfs init operation"}})
sys.exit(1)
# The transfer process should use the information it needs from the
# initiation structure, and also perform any one-off setup tasks it
# needs to do. It should then respond on stdout with a simple empty
# confirmation structure, as follows:
write_msg({})
# After the initiation exchange, git-lfs will send any number of
# transfer requests to the stdin of the transfer process, in a serial sequence.
while True:
msg = read_msg()
if msg is None:
# When all transfers have been processed, git-lfs will send
# a terminate event to the stdin of the transfer process.
# On receiving this message the transfer process should
# clean up and terminate. No response is expected.
sys.exit(0)
oid = msg["oid"]
filepath = msg["path"]
completion_url = msg["action"]["href"]
header = msg["action"]["header"]
chunk_size = int(header.pop("chunk_size"))
presigned_urls: List[str] = list(header.values())
# Send a "started" progress event to allow other workers to start.
# Otherwise they're delayed until first "progress" event is reported,
# i.e. after the first 5GB by default (!)
write_msg(
{
"event": "progress",
"oid": oid,
"bytesSoFar": 1,
"bytesSinceLast": 0,
}
)
parts = []
with open(filepath, "rb") as file:
for i, presigned_url in enumerate(presigned_urls):
with SliceFileObj(
file,
seek_from=i * chunk_size,
read_limit=chunk_size,
) as data:
r = get_session().put(presigned_url, data=data)
hf_raise_for_status(r)
parts.append(
{
"etag": r.headers.get("etag"),
"partNumber": i + 1,
}
)
# In order to support progress reporting while data is uploading / downloading,
# the transfer process should post messages to stdout
write_msg(
{
"event": "progress",
"oid": oid,
"bytesSoFar": (i + 1) * chunk_size,
"bytesSinceLast": chunk_size,
}
)
# Not precise but that's ok.
r = get_session().post(
completion_url,
json={
"oid": oid,
"parts": parts,
},
)
hf_raise_for_status(r)
write_msg({"event": "complete", "oid": oid})
kronos-venv/lib/python3.12/site-packages/huggingface_hub/commands/repo.py
+147
View File
@@ -0,0 +1,147 @@
# Copyright 2025 The HuggingFace Team. All rights reserved.
#
# 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 commands to interact with repositories on the Hugging Face Hub.
Usage:
# create a new dataset repo on the Hub
huggingface-cli repo create my-cool-dataset --repo-type=dataset
# create a private model repo on the Hub
huggingface-cli repo create my-cool-model --private
"""
import argparse
from argparse import _SubParsersAction
from typing import Optional
from huggingface_hub.commands import BaseHuggingfaceCLICommand
from huggingface_hub.commands._cli_utils import ANSI
from huggingface_hub.constants import SPACES_SDK_TYPES
from huggingface_hub.hf_api import HfApi
from huggingface_hub.utils import logging
logger = logging.get_logger(__name__)
class RepoCommands(BaseHuggingfaceCLICommand):
@staticmethod
def register_subcommand(parser: _SubParsersAction):
repo_parser = parser.add_parser("repo", help="{create} Commands to interact with your huggingface.co repos.")
repo_subparsers = repo_parser.add_subparsers(help="huggingface.co repos related commands")
repo_create_parser = repo_subparsers.add_parser("create", help="Create a new repo on huggingface.co")
repo_create_parser.add_argument(
"repo_id",
type=str,
help="The ID of the repo to create to (e.g. `username/repo-name`). The username is optional and will be set to your username if not provided.",
)
repo_create_parser.add_argument(
"--repo-type",
type=str,
help='Optional: set to "dataset" or "space" if creating a dataset or space, default is model.',
)
repo_create_parser.add_argument(
"--space_sdk",
type=str,
help='Optional: Hugging Face Spaces SDK type. Required when --type is set to "space".',
choices=SPACES_SDK_TYPES,
)
repo_create_parser.add_argument(
"--private",
action="store_true",
help="Whether to create a private repository. Defaults to public unless the organization's default is private.",
)
repo_create_parser.add_argument(
"--token",
type=str,
help="Hugging Face token. Will default to the locally saved token if not provided.",
)
repo_create_parser.add_argument(
"--exist-ok",
action="store_true",
help="Do not raise an error if repo already exists.",
)
repo_create_parser.add_argument(
"--resource-group-id",
type=str,
help="Resource group in which to create the repo. Resource groups is only available for Enterprise Hub organizations.",
)
repo_create_parser.add_argument(
"--type",
type=str,
help="[Deprecated]: use --repo-type instead.",
)
repo_create_parser.add_argument(
"-y",
"--yes",
action="store_true",
help="[Deprecated] no effect.",
)
repo_create_parser.add_argument(
"--organization", type=str, help="[Deprecated] Pass the organization namespace directly in the repo_id."
)
repo_create_parser.set_defaults(func=lambda args: RepoCreateCommand(args))
class RepoCreateCommand:
def __init__(self, args: argparse.Namespace):
self.repo_id: str = args.repo_id
self.repo_type: Optional[str] = args.repo_type or args.type
self.space_sdk: Optional[str] = args.space_sdk
self.organization: Optional[str] = args.organization
self.yes: bool = args.yes
self.private: bool = args.private
self.token: Optional[str] = args.token
self.exist_ok: bool = args.exist_ok
self.resource_group_id: Optional[str] = args.resource_group_id
if args.type is not None:
print(
ANSI.yellow(
"The --type argument is deprecated and will be removed in a future version. Use --repo-type instead."
)
)
if self.organization is not None:
print(
ANSI.yellow(
"The --organization argument is deprecated and will be removed in a future version. Pass the organization namespace directly in the repo_id."
)
)
if self.yes:
print(
ANSI.yellow(
"The --yes argument is deprecated and will be removed in a future version. It does not have any effect."
)
)
self._api = HfApi()
def run(self):
if self.organization is not None:
if "/" in self.repo_id:
print(ANSI.red("You cannot pass both --organization and a repo_id with a namespace."))
exit(1)
self.repo_id = f"{self.organization}/{self.repo_id}"
repo_url = self._api.create_repo(
repo_id=self.repo_id,
repo_type=self.repo_type,
private=self.private,
token=self.token,
exist_ok=self.exist_ok,
resource_group_id=self.resource_group_id,
space_sdk=self.space_sdk,
)
print(f"Successfully created {ANSI.bold(repo_url.repo_id)} on the Hub.")
print(f"Your repo is now available at {ANSI.bold(repo_url)}")
kronos-venv/lib/python3.12/site-packages/huggingface_hub/commands/repo_files.py
+128
View File
@@ -0,0 +1,128 @@
# 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 command to update or delete files in a repository using the CLI.
Usage:
# delete all
huggingface-cli repo-files <repo_id> delete "*"
# delete single file
huggingface-cli repo-files <repo_id> delete file.txt
# delete single folder
huggingface-cli repo-files <repo_id> delete folder/
# delete multiple
huggingface-cli repo-files <repo_id> delete file.txt folder/ file2.txt
# delete multiple patterns
huggingface-cli repo-files <repo_id> delete file.txt "*.json" "folder/*.parquet"
# delete from different revision / repo-type
huggingface-cli repo-files <repo_id> delete file.txt --revision=refs/pr/1 --repo-type=dataset
"""
from argparse import _SubParsersAction
from typing import List, Optional
from huggingface_hub import logging
from huggingface_hub.commands import BaseHuggingfaceCLICommand
from huggingface_hub.hf_api import HfApi
logger = logging.get_logger(__name__)
class DeleteFilesSubCommand:
def __init__(self, args) -> None:
self.args = args
self.repo_id: str = args.repo_id
self.repo_type: Optional[str] = args.repo_type
self.revision: Optional[str] = args.revision
self.api: HfApi = HfApi(token=args.token, library_name="huggingface-cli")
self.patterns: List[str] = args.patterns
self.commit_message: Optional[str] = args.commit_message
self.commit_description: Optional[str] = args.commit_description
self.create_pr: bool = args.create_pr
self.token: Optional[str] = args.token
def run(self) -> None:
logging.set_verbosity_info()
url = self.api.delete_files(
delete_patterns=self.patterns,
repo_id=self.repo_id,
repo_type=self.repo_type,
revision=self.revision,
commit_message=self.commit_message,
commit_description=self.commit_description,
create_pr=self.create_pr,
)
print(f"Files correctly deleted from repo. Commit: {url}.")
logging.set_verbosity_warning()
class RepoFilesCommand(BaseHuggingfaceCLICommand):
@staticmethod
def register_subcommand(parser: _SubParsersAction):
repo_files_parser = parser.add_parser("repo-files", help="Manage files in a repo on the Hub")
repo_files_parser.add_argument(
"repo_id", type=str, help="The ID of the repo to manage (e.g. `username/repo-name`)."
)
repo_files_subparsers = repo_files_parser.add_subparsers(
help="Action to execute against the files.",
required=True,
)
delete_subparser = repo_files_subparsers.add_parser(
"delete",
help="Delete files from a repo on the Hub",
)
delete_subparser.set_defaults(func=lambda args: DeleteFilesSubCommand(args))
delete_subparser.add_argument(
"patterns",
nargs="+",
type=str,
help="Glob patterns to match files to delete.",
)
delete_subparser.add_argument(
"--repo-type",
choices=["model", "dataset", "space"],
default="model",
help="Type of the repo to upload to (e.g. `dataset`).",
)
delete_subparser.add_argument(
"--revision",
type=str,
help=(
"An optional Git revision to push to. It can be a branch name "
"or a PR reference. If revision does not"
" exist and `--create-pr` is not set, a branch will be automatically created."
),
)
delete_subparser.add_argument(
"--commit-message", type=str, help="The summary / title / first line of the generated commit."
)
delete_subparser.add_argument(
"--commit-description", type=str, help="The description of the generated commit."
)
delete_subparser.add_argument(
"--create-pr", action="store_true", help="Whether to create a new Pull Request for these changes."
)
repo_files_parser.add_argument(
"--token",
type=str,
help="A User Access Token generated from https://huggingface.co/settings/tokens",
)
repo_files_parser.set_defaults(func=RepoFilesCommand)
kronos-venv/lib/python3.12/site-packages/huggingface_hub/commands/scan_cache.py
+181
View File
@@ -0,0 +1,181 @@
# coding=utf-8
# Copyright 2022-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 command to scan the HF cache directory.
Usage:
huggingface-cli scan-cache
huggingface-cli scan-cache -v
huggingface-cli scan-cache -vvv
huggingface-cli scan-cache --dir ~/.cache/huggingface/hub
"""
import time
from argparse import Namespace, _SubParsersAction
from typing import Optional
from ..utils import CacheNotFound, HFCacheInfo, scan_cache_dir
from . import BaseHuggingfaceCLICommand
from ._cli_utils import ANSI, tabulate
class ScanCacheCommand(BaseHuggingfaceCLICommand):
@staticmethod
def register_subcommand(parser: _SubParsersAction):
scan_cache_parser = parser.add_parser("scan-cache", help="Scan cache directory.")
scan_cache_parser.add_argument(
"--dir",
type=str,
default=None,
help="cache directory to scan (optional). Default to the default HuggingFace cache.",
)
scan_cache_parser.add_argument(
"-v",
"--verbose",
action="count",
default=0,
help="show a more verbose output",
)
scan_cache_parser.set_defaults(func=ScanCacheCommand)
def __init__(self, args: Namespace) -> None:
self.verbosity: int = args.verbose
self.cache_dir: Optional[str] = args.dir
def run(self):
try:
t0 = time.time()
hf_cache_info = scan_cache_dir(self.cache_dir)
t1 = time.time()
except CacheNotFound as exc:
cache_dir = exc.cache_dir
print(f"Cache directory not found: {cache_dir}")
return
self._print_hf_cache_info_as_table(hf_cache_info)
print(
f"\nDone in {round(t1 - t0, 1)}s. Scanned {len(hf_cache_info.repos)} repo(s)"
f" for a total of {ANSI.red(hf_cache_info.size_on_disk_str)}."
)
if len(hf_cache_info.warnings) > 0:
message = f"Got {len(hf_cache_info.warnings)} warning(s) while scanning."
if self.verbosity >= 3:
print(ANSI.gray(message))
for warning in hf_cache_info.warnings:
print(ANSI.gray(warning))
else:
print(ANSI.gray(message + " Use -vvv to print details."))
def _print_hf_cache_info_as_table(self, hf_cache_info: HFCacheInfo) -> None:
print(get_table(hf_cache_info, verbosity=self.verbosity))
def get_table(hf_cache_info: HFCacheInfo, *, verbosity: int = 0) -> str:
"""Generate a table from the [`HFCacheInfo`] object.
Pass `verbosity=0` to get a table with a single row per repo, with columns
"repo_id", "repo_type", "size_on_disk", "nb_files", "last_accessed", "last_modified", "refs", "local_path".
Pass `verbosity=1` to get a table with a row per repo and revision (thus multiple rows can appear for a single repo), with columns
"repo_id", "repo_type", "revision", "size_on_disk", "nb_files", "last_modified", "refs", "local_path".
Example:
```py
>>> from huggingface_hub.utils import scan_cache_dir
>>> from huggingface_hub.commands.scan_cache import get_table
>>> hf_cache_info = scan_cache_dir()
HFCacheInfo(...)
>>> print(get_table(hf_cache_info, verbosity=0))
REPO ID REPO TYPE SIZE ON DISK NB FILES LAST_ACCESSED LAST_MODIFIED REFS LOCAL PATH
--------------------------------------------------- --------- ------------ -------- ------------- ------------- ---- --------------------------------------------------------------------------------------------------
roberta-base model 2.7M 5 1 day ago 1 week ago main C:\\Users\\admin\\.cache\\huggingface\\hub\\models--roberta-base
suno/bark model 8.8K 1 1 week ago 1 week ago main C:\\Users\\admin\\.cache\\huggingface\\hub\\models--suno--bark
t5-base model 893.8M 4 4 days ago 7 months ago main C:\\Users\\admin\\.cache\\huggingface\\hub\\models--t5-base
t5-large model 3.0G 4 5 weeks ago 5 months ago main C:\\Users\\admin\\.cache\\huggingface\\hub\\models--t5-large
>>> print(get_table(hf_cache_info, verbosity=1))
REPO ID REPO TYPE REVISION SIZE ON DISK NB FILES LAST_MODIFIED REFS LOCAL PATH
--------------------------------------------------- --------- ---------------------------------------- ------------ -------- ------------- ---- -----------------------------------------------------------------------------------------------------------------------------------------------------
roberta-base model e2da8e2f811d1448a5b465c236feacd80ffbac7b 2.7M 5 1 week ago main C:\\Users\\admin\\.cache\\huggingface\\hub\\models--roberta-base\\snapshots\\e2da8e2f811d1448a5b465c236feacd80ffbac7b
suno/bark model 70a8a7d34168586dc5d028fa9666aceade177992 8.8K 1 1 week ago main C:\\Users\\admin\\.cache\\huggingface\\hub\\models--suno--bark\\snapshots\\70a8a7d34168586dc5d028fa9666aceade177992
t5-base model a9723ea7f1b39c1eae772870f3b547bf6ef7e6c1 893.8M 4 7 months ago main C:\\Users\\admin\\.cache\\huggingface\\hub\\models--t5-base\\snapshots\\a9723ea7f1b39c1eae772870f3b547bf6ef7e6c1
t5-large model 150ebc2c4b72291e770f58e6057481c8d2ed331a 3.0G 4 5 months ago main C:\\Users\\admin\\.cache\\huggingface\\hub\\models--t5-large\\snapshots\\150ebc2c4b72291e770f58e6057481c8d2ed331a ```
```
Args:
hf_cache_info ([`HFCacheInfo`]):
The HFCacheInfo object to print.
verbosity (`int`, *optional*):
The verbosity level. Defaults to 0.
Returns:
`str`: The table as a string.
"""
if verbosity == 0:
return tabulate(
rows=[
[
repo.repo_id,
repo.repo_type,
"{:>12}".format(repo.size_on_disk_str),
repo.nb_files,
repo.last_accessed_str,
repo.last_modified_str,
", ".join(sorted(repo.refs)),
str(repo.repo_path),
]
for repo in sorted(hf_cache_info.repos, key=lambda repo: repo.repo_path)
],
headers=[
"REPO ID",
"REPO TYPE",
"SIZE ON DISK",
"NB FILES",
"LAST_ACCESSED",
"LAST_MODIFIED",
"REFS",
"LOCAL PATH",
],
)
else:
return tabulate(
rows=[
[
repo.repo_id,
repo.repo_type,
revision.commit_hash,
"{:>12}".format(revision.size_on_disk_str),
revision.nb_files,
revision.last_modified_str,
", ".join(sorted(revision.refs)),
str(revision.snapshot_path),
]
for repo in sorted(hf_cache_info.repos, key=lambda repo: repo.repo_path)
for revision in sorted(repo.revisions, key=lambda revision: revision.commit_hash)
],
headers=[
"REPO ID",
"REPO TYPE",
"REVISION",
"SIZE ON DISK",
"NB FILES",
"LAST_MODIFIED",
"REFS",
"LOCAL PATH",
],
)
kronos-venv/lib/python3.12/site-packages/huggingface_hub/commands/tag.py
+159
View File
@@ -0,0 +1,159 @@
# coding=utf-8
# Copyright 2024-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 commands to perform tag management with the CLI.
Usage Examples:
- Create a tag:
$ huggingface-cli tag user/my-model 1.0 --message "First release"
$ huggingface-cli tag user/my-model 1.0 -m "First release" --revision develop
$ huggingface-cli tag user/my-dataset 1.0 -m "First release" --repo-type dataset
$ huggingface-cli tag user/my-space 1.0
- List all tags:
$ huggingface-cli tag -l user/my-model
$ huggingface-cli tag --list user/my-dataset --repo-type dataset
- Delete a tag:
$ huggingface-cli tag -d user/my-model 1.0
$ huggingface-cli tag --delete user/my-dataset 1.0 --repo-type dataset
$ huggingface-cli tag -d user/my-space 1.0 -y
"""
from argparse import Namespace, _SubParsersAction
from requests.exceptions import HTTPError
from huggingface_hub.commands import BaseHuggingfaceCLICommand
from huggingface_hub.constants import (
REPO_TYPES,
)
from huggingface_hub.hf_api import HfApi
from ..errors import HfHubHTTPError, RepositoryNotFoundError, RevisionNotFoundError
from ._cli_utils import ANSI
class TagCommands(BaseHuggingfaceCLICommand):
@staticmethod
def register_subcommand(parser: _SubParsersAction):
tag_parser = parser.add_parser("tag", help="(create, list, delete) tags for a repo in the hub")
tag_parser.add_argument("repo_id", type=str, help="The ID of the repo to tag (e.g. `username/repo-name`).")
tag_parser.add_argument("tag", nargs="?", type=str, help="The name of the tag for creation or deletion.")
tag_parser.add_argument("-m", "--message", type=str, help="The description of the tag to create.")
tag_parser.add_argument("--revision", type=str, help="The git revision to tag.")
tag_parser.add_argument(
"--token", type=str, help="A User Access Token generated from https://huggingface.co/settings/tokens."
)
tag_parser.add_argument(
"--repo-type",
choices=["model", "dataset", "space"],
default="model",
help="Set the type of repository (model, dataset, or space).",
)
tag_parser.add_argument("-y", "--yes", action="store_true", help="Answer Yes to prompts automatically.")
tag_parser.add_argument("-l", "--list", action="store_true", help="List tags for a repository.")
tag_parser.add_argument("-d", "--delete", action="store_true", help="Delete a tag for a repository.")
tag_parser.set_defaults(func=lambda args: handle_commands(args))
def handle_commands(args: Namespace):
if args.list:
return TagListCommand(args)
elif args.delete:
return TagDeleteCommand(args)
else:
return TagCreateCommand(args)
class TagCommand:
def __init__(self, args: Namespace):
self.args = args
self.api = HfApi(token=self.args.token)
self.repo_id = self.args.repo_id
self.repo_type = self.args.repo_type
if self.repo_type not in REPO_TYPES:
print("Invalid repo --repo-type")
exit(1)
class TagCreateCommand(TagCommand):
def run(self):
print(f"You are about to create tag {ANSI.bold(self.args.tag)} on {self.repo_type} {ANSI.bold(self.repo_id)}")
try:
self.api.create_tag(
repo_id=self.repo_id,
tag=self.args.tag,
tag_message=self.args.message,
revision=self.args.revision,
repo_type=self.repo_type,
)
except RepositoryNotFoundError:
print(f"{self.repo_type.capitalize()} {ANSI.bold(self.repo_id)} not found.")
exit(1)
except RevisionNotFoundError:
print(f"Revision {ANSI.bold(self.args.revision)} not found.")
exit(1)
except HfHubHTTPError as e:
if e.response.status_code == 409:
print(f"Tag {ANSI.bold(self.args.tag)} already exists on {ANSI.bold(self.repo_id)}")
exit(1)
raise e
print(f"Tag {ANSI.bold(self.args.tag)} created on {ANSI.bold(self.repo_id)}")
class TagListCommand(TagCommand):
def run(self):
try:
refs = self.api.list_repo_refs(
repo_id=self.repo_id,
repo_type=self.repo_type,
)
except RepositoryNotFoundError:
print(f"{self.repo_type.capitalize()} {ANSI.bold(self.repo_id)} not found.")
exit(1)
except HTTPError as e:
print(e)
print(ANSI.red(e.response.text))
exit(1)
if len(refs.tags) == 0:
print("No tags found")
exit(0)
print(f"Tags for {self.repo_type} {ANSI.bold(self.repo_id)}:")
for tag in refs.tags:
print(tag.name)
class TagDeleteCommand(TagCommand):
def run(self):
print(f"You are about to delete tag {ANSI.bold(self.args.tag)} on {self.repo_type} {ANSI.bold(self.repo_id)}")
if not self.args.yes:
choice = input("Proceed? [Y/n] ").lower()
if choice not in ("", "y", "yes"):
print("Abort")
exit()
try:
self.api.delete_tag(repo_id=self.repo_id, tag=self.args.tag, repo_type=self.repo_type)
except RepositoryNotFoundError:
print(f"{self.repo_type.capitalize()} {ANSI.bold(self.repo_id)} not found.")
exit(1)
except RevisionNotFoundError:
print(f"Tag {ANSI.bold(self.args.tag)} not found on {ANSI.bold(self.repo_id)}")
exit(1)
print(f"Tag {ANSI.bold(self.args.tag)} deleted on {ANSI.bold(self.repo_id)}")
kronos-venv/lib/python3.12/site-packages/huggingface_hub/commands/upload.py
+314
View File
@@ -0,0 +1,314 @@
# 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 command to upload a repo or file with the CLI.
Usage:
# Upload file (implicit)
huggingface-cli upload my-cool-model ./my-cool-model.safetensors
# Upload file (explicit)
huggingface-cli upload my-cool-model ./my-cool-model.safetensors model.safetensors
# Upload directory (implicit). If `my-cool-model/` is a directory it will be uploaded, otherwise an exception is raised.
huggingface-cli upload my-cool-model
# Upload directory (explicit)
huggingface-cli upload my-cool-model ./models/my-cool-model .
# Upload filtered directory (example: tensorboard logs except for the last run)
huggingface-cli upload my-cool-model ./model/training /logs --include "*.tfevents.*" --exclude "*20230905*"
# Upload with wildcard
huggingface-cli upload my-cool-model "./model/training/*.safetensors"
# Upload private dataset
huggingface-cli upload Wauplin/my-cool-dataset ./data . --repo-type=dataset --private
# Upload with token
huggingface-cli upload Wauplin/my-cool-model --token=hf_****
# Sync local Space with Hub (upload new files, delete removed files)
huggingface-cli upload Wauplin/space-example --repo-type=space --exclude="/logs/*" --delete="*" --commit-message="Sync local Space with Hub"
# Schedule commits every 30 minutes
huggingface-cli upload Wauplin/my-cool-model --every=30
"""
import os
import time
import warnings
from argparse import Namespace, _SubParsersAction
from typing import List, Optional
from huggingface_hub import logging
from huggingface_hub._commit_scheduler import CommitScheduler
from huggingface_hub.commands import BaseHuggingfaceCLICommand
from huggingface_hub.constants import HF_HUB_ENABLE_HF_TRANSFER
from huggingface_hub.errors import RevisionNotFoundError
from huggingface_hub.hf_api import HfApi
from huggingface_hub.utils import disable_progress_bars, enable_progress_bars
from huggingface_hub.utils._runtime import is_xet_available
logger = logging.get_logger(__name__)
class UploadCommand(BaseHuggingfaceCLICommand):
@staticmethod
def register_subcommand(parser: _SubParsersAction):
upload_parser = parser.add_parser("upload", help="Upload a file or a folder to a repo on the Hub")
upload_parser.add_argument(
"repo_id", type=str, help="The ID of the repo to upload to (e.g. `username/repo-name`)."
)
upload_parser.add_argument(
"local_path",
nargs="?",
help="Local path to the file or folder to upload. Wildcard patterns are supported. Defaults to current directory.",
)
upload_parser.add_argument(
"path_in_repo",
nargs="?",
help="Path of the file or folder in the repo. Defaults to the relative path of the file or folder.",
)
upload_parser.add_argument(
"--repo-type",
choices=["model", "dataset", "space"],
default="model",
help="Type of the repo to upload to (e.g. `dataset`).",
)
upload_parser.add_argument(
"--revision",
type=str,
help=(
"An optional Git revision to push to. It can be a branch name or a PR reference. If revision does not"
" exist and `--create-pr` is not set, a branch will be automatically created."
),
)
upload_parser.add_argument(
"--private",
action="store_true",
help=(
"Whether to create a private repo if repo doesn't exist on the Hub. Ignored if the repo already"
" exists."
),
)
upload_parser.add_argument("--include", nargs="*", type=str, help="Glob patterns to match files to upload.")
upload_parser.add_argument(
"--exclude", nargs="*", type=str, help="Glob patterns to exclude from files to upload."
)
upload_parser.add_argument(
"--delete",
nargs="*",
type=str,
help="Glob patterns for file to be deleted from the repo while committing.",
)
upload_parser.add_argument(
"--commit-message", type=str, help="The summary / title / first line of the generated commit."
)
upload_parser.add_argument("--commit-description", type=str, help="The description of the generated commit.")
upload_parser.add_argument(
"--create-pr", action="store_true", help="Whether to upload content as a new Pull Request."
)
upload_parser.add_argument(
"--every",
type=float,
help="If set, a background job is scheduled to create commits every `every` minutes.",
)
upload_parser.add_argument(
"--token", type=str, help="A User Access Token generated from https://huggingface.co/settings/tokens"
)
upload_parser.add_argument(
"--quiet",
action="store_true",
help="If True, progress bars are disabled and only the path to the uploaded files is printed.",
)
upload_parser.set_defaults(func=UploadCommand)
def __init__(self, args: Namespace) -> None:
self.repo_id: str = args.repo_id
self.repo_type: Optional[str] = args.repo_type
self.revision: Optional[str] = args.revision
self.private: bool = args.private
self.include: Optional[List[str]] = args.include
self.exclude: Optional[List[str]] = args.exclude
self.delete: Optional[List[str]] = args.delete
self.commit_message: Optional[str] = args.commit_message
self.commit_description: Optional[str] = args.commit_description
self.create_pr: bool = args.create_pr
self.api: HfApi = HfApi(token=args.token, library_name="huggingface-cli")
self.quiet: bool = args.quiet # disable warnings and progress bars
# Check `--every` is valid
if args.every is not None and args.every <= 0:
raise ValueError(f"`every` must be a positive value (got '{args.every}')")
self.every: Optional[float] = args.every
# Resolve `local_path` and `path_in_repo`
repo_name: str = args.repo_id.split("/")[-1] # e.g. "Wauplin/my-cool-model" => "my-cool-model"
self.local_path: str
self.path_in_repo: str
if args.local_path is not None and any(c in args.local_path for c in ["*", "?", "["]):
if args.include is not None:
raise ValueError("Cannot set `--include` when passing a `local_path` containing a wildcard.")
if args.path_in_repo is not None and args.path_in_repo != ".":
raise ValueError("Cannot set `path_in_repo` when passing a `local_path` containing a wildcard.")
self.local_path = "."
self.include = args.local_path
self.path_in_repo = "."
elif args.local_path is None and os.path.isfile(repo_name):
# Implicit case 1: user provided only a repo_id which happen to be a local file as well => upload it with same name
self.local_path = repo_name
self.path_in_repo = repo_name
elif args.local_path is None and os.path.isdir(repo_name):
# Implicit case 2: user provided only a repo_id which happen to be a local folder as well => upload it at root
self.local_path = repo_name
self.path_in_repo = "."
elif args.local_path is None:
# Implicit case 3: user provided only a repo_id that does not match a local file or folder
# => the user must explicitly provide a local_path => raise exception
raise ValueError(f"'{repo_name}' is not a local file or folder. Please set `local_path` explicitly.")
elif args.path_in_repo is None and os.path.isfile(args.local_path):
# Explicit local path to file, no path in repo => upload it at root with same name
self.local_path = args.local_path
self.path_in_repo = os.path.basename(args.local_path)
elif args.path_in_repo is None:
# Explicit local path to folder, no path in repo => upload at root
self.local_path = args.local_path
self.path_in_repo = "."
else:
# Finally, if both paths are explicit
self.local_path = args.local_path
self.path_in_repo = args.path_in_repo
def run(self) -> None:
if self.quiet:
disable_progress_bars()
with warnings.catch_warnings():
warnings.simplefilter("ignore")
print(self._upload())
enable_progress_bars()
else:
logging.set_verbosity_info()
print(self._upload())
logging.set_verbosity_warning()
def _upload(self) -> str:
if os.path.isfile(self.local_path):
if self.include is not None and len(self.include) > 0:
warnings.warn("Ignoring `--include` since a single file is uploaded.")
if self.exclude is not None and len(self.exclude) > 0:
warnings.warn("Ignoring `--exclude` since a single file is uploaded.")
if self.delete is not None and len(self.delete) > 0:
warnings.warn("Ignoring `--delete` since a single file is uploaded.")
if not is_xet_available() and not HF_HUB_ENABLE_HF_TRANSFER:
logger.info(
"Consider using `hf_transfer` for faster uploads. This solution comes with some limitations. See"
" https://huggingface.co/docs/huggingface_hub/hf_transfer for more details."
)
# Schedule commits if `every` is set
if self.every is not None:
if os.path.isfile(self.local_path):
# If file => watch entire folder + use allow_patterns
folder_path = os.path.dirname(self.local_path)
path_in_repo = (
self.path_in_repo[: -len(self.local_path)] # remove filename from path_in_repo
if self.path_in_repo.endswith(self.local_path)
else self.path_in_repo
)
allow_patterns = [self.local_path]
ignore_patterns = []
else:
folder_path = self.local_path
path_in_repo = self.path_in_repo
allow_patterns = self.include or []
ignore_patterns = self.exclude or []
if self.delete is not None and len(self.delete) > 0:
warnings.warn("Ignoring `--delete` when uploading with scheduled commits.")
scheduler = CommitScheduler(
folder_path=folder_path,
repo_id=self.repo_id,
repo_type=self.repo_type,
revision=self.revision,
allow_patterns=allow_patterns,
ignore_patterns=ignore_patterns,
path_in_repo=path_in_repo,
private=self.private,
every=self.every,
hf_api=self.api,
)
print(f"Scheduling commits every {self.every} minutes to {scheduler.repo_id}.")
try: # Block main thread until KeyboardInterrupt
while True:
time.sleep(100)
except KeyboardInterrupt:
scheduler.stop()
return "Stopped scheduled commits."
# Otherwise, create repo and proceed with the upload
if not os.path.isfile(self.local_path) and not os.path.isdir(self.local_path):
raise FileNotFoundError(f"No such file or directory: '{self.local_path}'.")
repo_id = self.api.create_repo(
repo_id=self.repo_id,
repo_type=self.repo_type,
exist_ok=True,
private=self.private,
space_sdk="gradio" if self.repo_type == "space" else None,
# ^ We don't want it to fail when uploading to a Space => let's set Gradio by default.
# ^ I'd rather not add CLI args to set it explicitly as we already have `huggingface-cli repo create` for that.
).repo_id
# Check if branch already exists and if not, create it
if self.revision is not None and not self.create_pr:
try:
self.api.repo_info(repo_id=repo_id, repo_type=self.repo_type, revision=self.revision)
except RevisionNotFoundError:
logger.info(f"Branch '{self.revision}' not found. Creating it...")
self.api.create_branch(repo_id=repo_id, repo_type=self.repo_type, branch=self.revision, exist_ok=True)
# ^ `exist_ok=True` to avoid race concurrency issues
# File-based upload
if os.path.isfile(self.local_path):
return self.api.upload_file(
path_or_fileobj=self.local_path,
path_in_repo=self.path_in_repo,
repo_id=repo_id,
repo_type=self.repo_type,
revision=self.revision,
commit_message=self.commit_message,
commit_description=self.commit_description,
create_pr=self.create_pr,
)
# Folder-based upload
else:
return self.api.upload_folder(
folder_path=self.local_path,
path_in_repo=self.path_in_repo,
repo_id=repo_id,
repo_type=self.repo_type,
revision=self.revision,
commit_message=self.commit_message,
commit_description=self.commit_description,
create_pr=self.create_pr,
allow_patterns=self.include,
ignore_patterns=self.exclude,
delete_patterns=self.delete,
)
kronos-venv/lib/python3.12/site-packages/huggingface_hub/commands/upload_large_folder.py
+129
View File
@@ -0,0 +1,129 @@
# 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 command to upload a large folder with the CLI."""
import os
from argparse import Namespace, _SubParsersAction
from typing import List, Optional
from huggingface_hub import logging
from huggingface_hub.commands import BaseHuggingfaceCLICommand
from huggingface_hub.hf_api import HfApi
from huggingface_hub.utils import disable_progress_bars
from ._cli_utils import ANSI
logger = logging.get_logger(__name__)
class UploadLargeFolderCommand(BaseHuggingfaceCLICommand):
@staticmethod
def register_subcommand(parser: _SubParsersAction):
subparser = parser.add_parser("upload-large-folder", help="Upload a large folder to a repo on the Hub")
subparser.add_argument(
"repo_id", type=str, help="The ID of the repo to upload to (e.g. `username/repo-name`)."
)
subparser.add_argument("local_path", type=str, help="Local path to the file or folder to upload.")
subparser.add_argument(
"--repo-type",
choices=["model", "dataset", "space"],
help="Type of the repo to upload to (e.g. `dataset`).",
)
subparser.add_argument(
"--revision",
type=str,
help=("An optional Git revision to push to. It can be a branch name or a PR reference."),
)
subparser.add_argument(
"--private",
action="store_true",
help=(
"Whether to create a private repo if repo doesn't exist on the Hub. Ignored if the repo already exists."
),
)
subparser.add_argument("--include", nargs="*", type=str, help="Glob patterns to match files to upload.")
subparser.add_argument("--exclude", nargs="*", type=str, help="Glob patterns to exclude from files to upload.")
subparser.add_argument(
"--token", type=str, help="A User Access Token generated from https://huggingface.co/settings/tokens"
)
subparser.add_argument(
"--num-workers", type=int, help="Number of workers to use to hash, upload and commit files."
)
subparser.add_argument("--no-report", action="store_true", help="Whether to disable regular status report.")
subparser.add_argument("--no-bars", action="store_true", help="Whether to disable progress bars.")
subparser.set_defaults(func=UploadLargeFolderCommand)
def __init__(self, args: Namespace) -> None:
self.repo_id: str = args.repo_id
self.local_path: str = args.local_path
self.repo_type: str = args.repo_type
self.revision: Optional[str] = args.revision
self.private: bool = args.private
self.include: Optional[List[str]] = args.include
self.exclude: Optional[List[str]] = args.exclude
self.api: HfApi = HfApi(token=args.token, library_name="huggingface-cli")
self.num_workers: Optional[int] = args.num_workers
self.no_report: bool = args.no_report
self.no_bars: bool = args.no_bars
if not os.path.isdir(self.local_path):
raise ValueError("Large upload is only supported for folders.")
def run(self) -> None:
logging.set_verbosity_info()
print(
ANSI.yellow(
"You are about to upload a large folder to the Hub using `huggingface-cli upload-large-folder`. "
"This is a new feature so feedback is very welcome!\n"
"\n"
"A few things to keep in mind:\n"
" - Repository limits still apply: https://huggingface.co/docs/hub/repositories-recommendations\n"
" - Do not start several processes in parallel.\n"
" - You can interrupt and resume the process at any time. "
"The script will pick up where it left off except for partially uploaded files that would have to be entirely reuploaded.\n"
" - Do not upload the same folder to several repositories. If you need to do so, you must delete the `./.cache/huggingface/` folder first.\n"
"\n"
f"Some temporary metadata will be stored under `{self.local_path}/.cache/huggingface`.\n"
" - You must not modify those files manually.\n"
" - You must not delete the `./.cache/huggingface/` folder while a process is running.\n"
" - You can delete the `./.cache/huggingface/` folder to reinitialize the upload state when process is not running. Files will have to be hashed and preuploaded again, except for already committed files.\n"
"\n"
"If the process output is too verbose, you can disable the progress bars with `--no-bars`. "
"You can also entirely disable the status report with `--no-report`.\n"
"\n"
"For more details, run `huggingface-cli upload-large-folder --help` or check the documentation at "
"https://huggingface.co/docs/huggingface_hub/guides/upload#upload-a-large-folder."
)
)
if self.no_bars:
disable_progress_bars()
self.api.upload_large_folder(
repo_id=self.repo_id,
folder_path=self.local_path,
repo_type=self.repo_type,
revision=self.revision,
private=self.private,
allow_patterns=self.include,
ignore_patterns=self.exclude,
num_workers=self.num_workers,
print_report=not self.no_report,
)
kronos-venv/lib/python3.12/site-packages/huggingface_hub/commands/user.py
+198
View File
@@ -0,0 +1,198 @@
# Copyright 2020 The HuggingFace Team. All rights reserved.
#
# 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 commands to authenticate to the Hugging Face Hub and interact with your repositories.
Usage:
# login and save token locally.
huggingface-cli login --token=hf_*** --add-to-git-credential
# switch between tokens
huggingface-cli auth switch
# list all tokens
huggingface-cli auth list
# logout from a specific token, if no token-name is provided, all tokens will be deleted from your machine.
huggingface-cli logout --token-name=your_token_name
# find out which huggingface.co account you are logged in as
huggingface-cli whoami
"""
from argparse import _SubParsersAction
from typing import List, Optional
from requests.exceptions import HTTPError
from huggingface_hub.commands import BaseHuggingfaceCLICommand
from huggingface_hub.constants import ENDPOINT
from huggingface_hub.hf_api import HfApi
from .._login import auth_list, auth_switch, login, logout
from ..utils import get_stored_tokens, get_token, logging
from ._cli_utils import ANSI
logger = logging.get_logger(__name__)
try:
from InquirerPy import inquirer
from InquirerPy.base.control import Choice
_inquirer_py_available = True
except ImportError:
_inquirer_py_available = False
class UserCommands(BaseHuggingfaceCLICommand):
@staticmethod
def register_subcommand(parser: _SubParsersAction):
login_parser = parser.add_parser("login", help="Log in using a token from huggingface.co/settings/tokens")
login_parser.add_argument(
"--token",
type=str,
help="Token generated from https://huggingface.co/settings/tokens",
)
login_parser.add_argument(
"--add-to-git-credential",
action="store_true",
help="Optional: Save token to git credential helper.",
)
login_parser.set_defaults(func=lambda args: LoginCommand(args))
whoami_parser = parser.add_parser("whoami", help="Find out which huggingface.co account you are logged in as.")
whoami_parser.set_defaults(func=lambda args: WhoamiCommand(args))
logout_parser = parser.add_parser("logout", help="Log out")
logout_parser.add_argument(
"--token-name",
type=str,
help="Optional: Name of the access token to log out from.",
)
logout_parser.set_defaults(func=lambda args: LogoutCommand(args))
auth_parser = parser.add_parser("auth", help="Other authentication related commands")
auth_subparsers = auth_parser.add_subparsers(help="Authentication subcommands")
auth_switch_parser = auth_subparsers.add_parser("switch", help="Switch between access tokens")
auth_switch_parser.add_argument(
"--token-name",
type=str,
help="Optional: Name of the access token to switch to.",
)
auth_switch_parser.add_argument(
"--add-to-git-credential",
action="store_true",
help="Optional: Save token to git credential helper.",
)
auth_switch_parser.set_defaults(func=lambda args: AuthSwitchCommand(args))
auth_list_parser = auth_subparsers.add_parser("list", help="List all stored access tokens")
auth_list_parser.set_defaults(func=lambda args: AuthListCommand(args))
class BaseUserCommand:
def __init__(self, args):
self.args = args
self._api = HfApi()
class LoginCommand(BaseUserCommand):
def run(self):
logging.set_verbosity_info()
login(
token=self.args.token,
add_to_git_credential=self.args.add_to_git_credential,
)
class LogoutCommand(BaseUserCommand):
def run(self):
logging.set_verbosity_info()
logout(token_name=self.args.token_name)
class AuthSwitchCommand(BaseUserCommand):
def run(self):
logging.set_verbosity_info()
token_name = self.args.token_name
if token_name is None:
token_name = self._select_token_name()
if token_name is None:
print("No token name provided. Aborting.")
exit()
auth_switch(token_name, add_to_git_credential=self.args.add_to_git_credential)
def _select_token_name(self) -> Optional[str]:
token_names = list(get_stored_tokens().keys())
if not token_names:
logger.error("No stored tokens found. Please login first.")
return None
if _inquirer_py_available:
return self._select_token_name_tui(token_names)
# if inquirer is not available, use a simpler terminal UI
print("Available stored tokens:")
for i, token_name in enumerate(token_names, 1):
print(f"{i}. {token_name}")
while True:
try:
choice = input("Enter the number of the token to switch to (or 'q' to quit): ")
if choice.lower() == "q":
return None
index = int(choice) - 1
if 0 <= index < len(token_names):
return token_names[index]
else:
print("Invalid selection. Please try again.")
except ValueError:
print("Invalid input. Please enter a number or 'q' to quit.")
def _select_token_name_tui(self, token_names: List[str]) -> Optional[str]:
choices = [Choice(token_name, name=token_name) for token_name in token_names]
try:
return inquirer.select(
message="Select a token to switch to:",
choices=choices,
default=None,
).execute()
except KeyboardInterrupt:
logger.info("Token selection cancelled.")
return None
class AuthListCommand(BaseUserCommand):
def run(self):
logging.set_verbosity_info()
auth_list()
class WhoamiCommand(BaseUserCommand):
def run(self):
token = get_token()
if token is None:
print("Not logged in")
exit()
try:
info = self._api.whoami(token)
print(info["name"])
orgs = [org["name"] for org in info["orgs"]]
if orgs:
print(ANSI.bold("orgs: "), ",".join(orgs))
if ENDPOINT != "https://huggingface.co":
print(f"Authenticated through private endpoint: {ENDPOINT}")
except HTTPError as e:
print(e)
print(ANSI.red(e.response.text))
exit(1)
kronos-venv/lib/python3.12/site-packages/huggingface_hub/commands/version.py
+37
View File
@@ -0,0 +1,37 @@
# Copyright 2022 The HuggingFace Team. All rights reserved.
#
# 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 command to print information about the version.
Usage:
huggingface-cli version
"""
from argparse import _SubParsersAction
from huggingface_hub import __version__
from . import BaseHuggingfaceCLICommand
class VersionCommand(BaseHuggingfaceCLICommand):
def __init__(self, args):
self.args = args
@staticmethod
def register_subcommand(parser: _SubParsersAction):
version_parser = parser.add_parser("version", help="Print information about the huggingface-cli version.")
version_parser.set_defaults(func=VersionCommand)
def run(self) -> None:
print(f"huggingface_hub version: {__version__}")
kronos-venv/lib/python3.12/site-packages/huggingface_hub/community.py
+355
View File
@@ -0,0 +1,355 @@
"""
Data structures to interact with Discussions and Pull Requests on the Hub.
See [the Discussions and Pull Requests guide](https://huggingface.co/docs/hub/repositories-pull-requests-discussions)
for more information on Pull Requests, Discussions, and the community tab.
"""
from dataclasses import dataclass
from datetime import datetime
from typing import List, Literal, Optional, Union
from . import constants
from .utils import parse_datetime
DiscussionStatus = Literal["open", "closed", "merged", "draft"]
@dataclass
class Discussion:
"""
A Discussion or Pull Request on the Hub.
This dataclass is not intended to be instantiated directly.
Attributes:
title (`str`):
The title of the Discussion / Pull Request
status (`str`):
The status of the Discussion / Pull Request.
It must be one of:
* `"open"`
* `"closed"`
* `"merged"` (only for Pull Requests )
* `"draft"` (only for Pull Requests )
num (`int`):
The number of the Discussion / Pull Request.
repo_id (`str`):
The id (`"{namespace}/{repo_name}"`) of the repo on which
the Discussion / Pull Request was open.
repo_type (`str`):
The type of the repo on which the Discussion / Pull Request was open.
Possible values are: `"model"`, `"dataset"`, `"space"`.
author (`str`):
The username of the Discussion / Pull Request author.
Can be `"deleted"` if the user has been deleted since.
is_pull_request (`bool`):
Whether or not this is a Pull Request.
created_at (`datetime`):
The `datetime` of creation of the Discussion / Pull Request.
endpoint (`str`):
Endpoint of the Hub. Default is https://huggingface.co.
git_reference (`str`, *optional*):
(property) Git reference to which changes can be pushed if this is a Pull Request, `None` otherwise.
url (`str`):
(property) URL of the discussion on the Hub.
"""
title: str
status: DiscussionStatus
num: int
repo_id: str
repo_type: str
author: str
is_pull_request: bool
created_at: datetime
endpoint: str
@property
def git_reference(self) -> Optional[str]:
"""
If this is a Pull Request , returns the git reference to which changes can be pushed.
Returns `None` otherwise.
"""
if self.is_pull_request:
return f"refs/pr/{self.num}"
return None
@property
def url(self) -> str:
"""Returns the URL of the discussion on the Hub."""
if self.repo_type is None or self.repo_type == constants.REPO_TYPE_MODEL:
return f"{self.endpoint}/{self.repo_id}/discussions/{self.num}"
return f"{self.endpoint}/{self.repo_type}s/{self.repo_id}/discussions/{self.num}"
@dataclass
class DiscussionWithDetails(Discussion):
"""
Subclass of [`Discussion`].
Attributes:
title (`str`):
The title of the Discussion / Pull Request
status (`str`):
The status of the Discussion / Pull Request.
It can be one of:
* `"open"`
* `"closed"`
* `"merged"` (only for Pull Requests )
* `"draft"` (only for Pull Requests )
num (`int`):
The number of the Discussion / Pull Request.
repo_id (`str`):
The id (`"{namespace}/{repo_name}"`) of the repo on which
the Discussion / Pull Request was open.
repo_type (`str`):
The type of the repo on which the Discussion / Pull Request was open.
Possible values are: `"model"`, `"dataset"`, `"space"`.
author (`str`):
The username of the Discussion / Pull Request author.
Can be `"deleted"` if the user has been deleted since.
is_pull_request (`bool`):
Whether or not this is a Pull Request.
created_at (`datetime`):
The `datetime` of creation of the Discussion / Pull Request.
events (`list` of [`DiscussionEvent`])
The list of [`DiscussionEvents`] in this Discussion or Pull Request.
conflicting_files (`Union[List[str], bool, None]`, *optional*):
A list of conflicting files if this is a Pull Request.
`None` if `self.is_pull_request` is `False`.
`True` if there are conflicting files but the list can't be retrieved.
target_branch (`str`, *optional*):
The branch into which changes are to be merged if this is a
Pull Request . `None` if `self.is_pull_request` is `False`.
merge_commit_oid (`str`, *optional*):
If this is a merged Pull Request , this is set to the OID / SHA of
the merge commit, `None` otherwise.
diff (`str`, *optional*):
The git diff if this is a Pull Request , `None` otherwise.
endpoint (`str`):
Endpoint of the Hub. Default is https://huggingface.co.
git_reference (`str`, *optional*):
(property) Git reference to which changes can be pushed if this is a Pull Request, `None` otherwise.
url (`str`):
(property) URL of the discussion on the Hub.
"""
events: List["DiscussionEvent"]
conflicting_files: Union[List[str], bool, None]
target_branch: Optional[str]
merge_commit_oid: Optional[str]
diff: Optional[str]
@dataclass
class DiscussionEvent:
"""
An event in a Discussion or Pull Request.
Use concrete classes:
* [`DiscussionComment`]
* [`DiscussionStatusChange`]
* [`DiscussionCommit`]
* [`DiscussionTitleChange`]
Attributes:
id (`str`):
The ID of the event. An hexadecimal string.
type (`str`):
The type of the event.
created_at (`datetime`):
A [`datetime`](https://docs.python.org/3/library/datetime.html?highlight=datetime#datetime.datetime)
object holding the creation timestamp for the event.
author (`str`):
The username of the Discussion / Pull Request author.
Can be `"deleted"` if the user has been deleted since.
"""
id: str
type: str
created_at: datetime
author: str
_event: dict
"""Stores the original event data, in case we need to access it later."""
@dataclass
class DiscussionComment(DiscussionEvent):
"""A comment in a Discussion / Pull Request.
Subclass of [`DiscussionEvent`].
Attributes:
id (`str`):
The ID of the event. An hexadecimal string.
type (`str`):
The type of the event.
created_at (`datetime`):
A [`datetime`](https://docs.python.org/3/library/datetime.html?highlight=datetime#datetime.datetime)
object holding the creation timestamp for the event.
author (`str`):
The username of the Discussion / Pull Request author.
Can be `"deleted"` if the user has been deleted since.
content (`str`):
The raw markdown content of the comment. Mentions, links and images are not rendered.
edited (`bool`):
Whether or not this comment has been edited.
hidden (`bool`):
Whether or not this comment has been hidden.
"""
content: str
edited: bool
hidden: bool
@property
def rendered(self) -> str:
"""The rendered comment, as a HTML string"""
return self._event["data"]["latest"]["html"]
@property
def last_edited_at(self) -> datetime:
"""The last edit time, as a `datetime` object."""
return parse_datetime(self._event["data"]["latest"]["updatedAt"])
@property
def last_edited_by(self) -> str:
"""The last edit time, as a `datetime` object."""
return self._event["data"]["latest"].get("author", {}).get("name", "deleted")
@property
def edit_history(self) -> List[dict]:
"""The edit history of the comment"""
return self._event["data"]["history"]
@property
def number_of_edits(self) -> int:
return len(self.edit_history)
@dataclass
class DiscussionStatusChange(DiscussionEvent):
"""A change of status in a Discussion / Pull Request.
Subclass of [`DiscussionEvent`].
Attributes:
id (`str`):
The ID of the event. An hexadecimal string.
type (`str`):
The type of the event.
created_at (`datetime`):
A [`datetime`](https://docs.python.org/3/library/datetime.html?highlight=datetime#datetime.datetime)
object holding the creation timestamp for the event.
author (`str`):
The username of the Discussion / Pull Request author.
Can be `"deleted"` if the user has been deleted since.
new_status (`str`):
The status of the Discussion / Pull Request after the change.
It can be one of:
* `"open"`
* `"closed"`
* `"merged"` (only for Pull Requests )
"""
new_status: str
@dataclass
class DiscussionCommit(DiscussionEvent):
"""A commit in a Pull Request.
Subclass of [`DiscussionEvent`].
Attributes:
id (`str`):
The ID of the event. An hexadecimal string.
type (`str`):
The type of the event.
created_at (`datetime`):
A [`datetime`](https://docs.python.org/3/library/datetime.html?highlight=datetime#datetime.datetime)
object holding the creation timestamp for the event.
author (`str`):
The username of the Discussion / Pull Request author.
Can be `"deleted"` if the user has been deleted since.
summary (`str`):
The summary of the commit.
oid (`str`):
The OID / SHA of the commit, as a hexadecimal string.
"""
summary: str
oid: str
@dataclass
class DiscussionTitleChange(DiscussionEvent):
"""A rename event in a Discussion / Pull Request.
Subclass of [`DiscussionEvent`].
Attributes:
id (`str`):
The ID of the event. An hexadecimal string.
type (`str`):
The type of the event.
created_at (`datetime`):
A [`datetime`](https://docs.python.org/3/library/datetime.html?highlight=datetime#datetime.datetime)
object holding the creation timestamp for the event.
author (`str`):
The username of the Discussion / Pull Request author.
Can be `"deleted"` if the user has been deleted since.
old_title (`str`):
The previous title for the Discussion / Pull Request.
new_title (`str`):
The new title.
"""
old_title: str
new_title: str
def deserialize_event(event: dict) -> DiscussionEvent:
"""Instantiates a [`DiscussionEvent`] from a dict"""
event_id: str = event["id"]
event_type: str = event["type"]
created_at = parse_datetime(event["createdAt"])
common_args = dict(
id=event_id,
type=event_type,
created_at=created_at,
author=event.get("author", {}).get("name", "deleted"),
_event=event,
)
if event_type == "comment":
return DiscussionComment(
**common_args,
edited=event["data"]["edited"],
hidden=event["data"]["hidden"],
content=event["data"]["latest"]["raw"],
)
if event_type == "status-change":
return DiscussionStatusChange(
**common_args,
new_status=event["data"]["status"],
)
if event_type == "commit":
return DiscussionCommit(
**common_args,
summary=event["data"]["subject"],
oid=event["data"]["oid"],
)
if event_type == "title-change":
return DiscussionTitleChange(
**common_args,
old_title=event["data"]["from"],
new_title=event["data"]["to"],
)
return DiscussionEvent(**common_args)
kronos-venv/lib/python3.12/site-packages/huggingface_hub/constants.py
+293
View File
@@ -0,0 +1,293 @@
import os
import re
import typing
from typing import Literal, Optional, Tuple
# Possible values for env variables
ENV_VARS_TRUE_VALUES = {"1", "ON", "YES", "TRUE"}
ENV_VARS_TRUE_AND_AUTO_VALUES = ENV_VARS_TRUE_VALUES.union({"AUTO"})
def _is_true(value: Optional[str]) -> bool:
if value is None:
return False
return value.upper() in ENV_VARS_TRUE_VALUES
def _as_int(value: Optional[str]) -> Optional[int]:
if value is None:
return None
return int(value)
# Constants for file downloads
PYTORCH_WEIGHTS_NAME = "pytorch_model.bin"
TF2_WEIGHTS_NAME = "tf_model.h5"
TF_WEIGHTS_NAME = "model.ckpt"
FLAX_WEIGHTS_NAME = "flax_model.msgpack"
CONFIG_NAME = "config.json"
REPOCARD_NAME = "README.md"
DEFAULT_ETAG_TIMEOUT = 10
DEFAULT_DOWNLOAD_TIMEOUT = 10
DEFAULT_REQUEST_TIMEOUT = 10
DOWNLOAD_CHUNK_SIZE = 10 * 1024 * 1024
HF_TRANSFER_CONCURRENCY = 100
MAX_HTTP_DOWNLOAD_SIZE = 50 * 1000 * 1000 * 1000 # 50 GB
# Constants for serialization
PYTORCH_WEIGHTS_FILE_PATTERN = "pytorch_model{suffix}.bin" # Unsafe pickle: use safetensors instead
SAFETENSORS_WEIGHTS_FILE_PATTERN = "model{suffix}.safetensors"
TF2_WEIGHTS_FILE_PATTERN = "tf_model{suffix}.h5"
# Constants for safetensors repos
SAFETENSORS_SINGLE_FILE = "model.safetensors"
SAFETENSORS_INDEX_FILE = "model.safetensors.index.json"
SAFETENSORS_MAX_HEADER_LENGTH = 25_000_000
# Timeout of aquiring file lock and logging the attempt
FILELOCK_LOG_EVERY_SECONDS = 10
# Git-related constants
DEFAULT_REVISION = "main"
REGEX_COMMIT_OID = re.compile(r"[A-Fa-f0-9]{5,40}")
HUGGINGFACE_CO_URL_HOME = "https://huggingface.co/"
_staging_mode = _is_true(os.environ.get("HUGGINGFACE_CO_STAGING"))
_HF_DEFAULT_ENDPOINT = "https://huggingface.co"
_HF_DEFAULT_STAGING_ENDPOINT = "https://hub-ci.huggingface.co"
ENDPOINT = os.getenv("HF_ENDPOINT", _HF_DEFAULT_ENDPOINT).rstrip("/")
HUGGINGFACE_CO_URL_TEMPLATE = ENDPOINT + "/{repo_id}/resolve/{revision}/{filename}"
if _staging_mode:
ENDPOINT = _HF_DEFAULT_STAGING_ENDPOINT
HUGGINGFACE_CO_URL_TEMPLATE = _HF_DEFAULT_STAGING_ENDPOINT + "/{repo_id}/resolve/{revision}/{filename}"
HUGGINGFACE_HEADER_X_REPO_COMMIT = "X-Repo-Commit"
HUGGINGFACE_HEADER_X_LINKED_ETAG = "X-Linked-Etag"
HUGGINGFACE_HEADER_X_LINKED_SIZE = "X-Linked-Size"
HUGGINGFACE_HEADER_X_BILL_TO = "X-HF-Bill-To"
INFERENCE_ENDPOINT = os.environ.get("HF_INFERENCE_ENDPOINT", "https://api-inference.huggingface.co")
# See https://huggingface.co/docs/inference-endpoints/index
INFERENCE_ENDPOINTS_ENDPOINT = "https://api.endpoints.huggingface.cloud/v2"
INFERENCE_CATALOG_ENDPOINT = "https://endpoints.huggingface.co/api/catalog"
# See https://api.endpoints.huggingface.cloud/#post-/v2/endpoint/-namespace-
INFERENCE_ENDPOINT_IMAGE_KEYS = [
"custom",
"huggingface",
"huggingfaceNeuron",
"llamacpp",
"tei",
"tgi",
"tgiNeuron",
]
# Proxy for third-party providers
INFERENCE_PROXY_TEMPLATE = "https://router.huggingface.co/{provider}"
REPO_ID_SEPARATOR = "--"
# ^ this substring is not allowed in repo_ids on hf.co
# and is the canonical one we use for serialization of repo ids elsewhere.
REPO_TYPE_DATASET = "dataset"
REPO_TYPE_SPACE = "space"
REPO_TYPE_MODEL = "model"
REPO_TYPES = [None, REPO_TYPE_MODEL, REPO_TYPE_DATASET, REPO_TYPE_SPACE]
SPACES_SDK_TYPES = ["gradio", "streamlit", "docker", "static"]
REPO_TYPES_URL_PREFIXES = {
REPO_TYPE_DATASET: "datasets/",
REPO_TYPE_SPACE: "spaces/",
}
REPO_TYPES_MAPPING = {
"datasets": REPO_TYPE_DATASET,
"spaces": REPO_TYPE_SPACE,
"models": REPO_TYPE_MODEL,
}
DiscussionTypeFilter = Literal["all", "discussion", "pull_request"]
DISCUSSION_TYPES: Tuple[DiscussionTypeFilter, ...] = typing.get_args(DiscussionTypeFilter)
DiscussionStatusFilter = Literal["all", "open", "closed"]
DISCUSSION_STATUS: Tuple[DiscussionTypeFilter, ...] = typing.get_args(DiscussionStatusFilter)
# Webhook subscription types
WEBHOOK_DOMAIN_T = Literal["repo", "discussions"]
# default cache
default_home = os.path.join(os.path.expanduser("~"), ".cache")
HF_HOME = os.path.expandvars(
os.path.expanduser(
os.getenv(
"HF_HOME",
os.path.join(os.getenv("XDG_CACHE_HOME", default_home), "huggingface"),
)
)
)
hf_cache_home = HF_HOME # for backward compatibility. TODO: remove this in 1.0.0
default_cache_path = os.path.join(HF_HOME, "hub")
default_assets_cache_path = os.path.join(HF_HOME, "assets")
# Legacy env variables
HUGGINGFACE_HUB_CACHE = os.getenv("HUGGINGFACE_HUB_CACHE", default_cache_path)
HUGGINGFACE_ASSETS_CACHE = os.getenv("HUGGINGFACE_ASSETS_CACHE", default_assets_cache_path)
# New env variables
HF_HUB_CACHE = os.path.expandvars(
os.path.expanduser(
os.getenv(
"HF_HUB_CACHE",
HUGGINGFACE_HUB_CACHE,
)
)
)
HF_ASSETS_CACHE = os.path.expandvars(
os.path.expanduser(
os.getenv(
"HF_ASSETS_CACHE",
HUGGINGFACE_ASSETS_CACHE,
)
)
)
HF_HUB_OFFLINE = _is_true(os.environ.get("HF_HUB_OFFLINE") or os.environ.get("TRANSFORMERS_OFFLINE"))
# If set, log level will be set to DEBUG and all requests made to the Hub will be logged
# as curl commands for reproducibility.
HF_DEBUG = _is_true(os.environ.get("HF_DEBUG"))
# Opt-out from telemetry requests
HF_HUB_DISABLE_TELEMETRY = (
_is_true(os.environ.get("HF_HUB_DISABLE_TELEMETRY")) # HF-specific env variable
or _is_true(os.environ.get("DISABLE_TELEMETRY"))
or _is_true(os.environ.get("DO_NOT_TRACK")) # https://consoledonottrack.com/
)
HF_TOKEN_PATH = os.path.expandvars(
os.path.expanduser(
os.getenv(
"HF_TOKEN_PATH",
os.path.join(HF_HOME, "token"),
)
)
)
HF_STORED_TOKENS_PATH = os.path.join(os.path.dirname(HF_TOKEN_PATH), "stored_tokens")
if _staging_mode:
# In staging mode, we use a different cache to ensure we don't mix up production and staging data or tokens
# In practice in `huggingface_hub` tests, we monkeypatch these values with temporary directories. The following
# lines are only used in third-party libraries tests (e.g. `transformers`, `diffusers`, etc.).
_staging_home = os.path.join(os.path.expanduser("~"), ".cache", "huggingface_staging")
HUGGINGFACE_HUB_CACHE = os.path.join(_staging_home, "hub")
HF_TOKEN_PATH = os.path.join(_staging_home, "token")
# Here, `True` will disable progress bars globally without possibility of enabling it
# programmatically. `False` will enable them without possibility of disabling them.
# If environment variable is not set (None), then the user is free to enable/disable
# them programmatically.
# TL;DR: env variable has priority over code
__HF_HUB_DISABLE_PROGRESS_BARS = os.environ.get("HF_HUB_DISABLE_PROGRESS_BARS")
HF_HUB_DISABLE_PROGRESS_BARS: Optional[bool] = (
_is_true(__HF_HUB_DISABLE_PROGRESS_BARS) if __HF_HUB_DISABLE_PROGRESS_BARS is not None else None
)
# Disable warning on machines that do not support symlinks (e.g. Windows non-developer)
HF_HUB_DISABLE_SYMLINKS_WARNING: bool = _is_true(os.environ.get("HF_HUB_DISABLE_SYMLINKS_WARNING"))
# Disable warning when using experimental features
HF_HUB_DISABLE_EXPERIMENTAL_WARNING: bool = _is_true(os.environ.get("HF_HUB_DISABLE_EXPERIMENTAL_WARNING"))
# Disable sending the cached token by default is all HTTP requests to the Hub
HF_HUB_DISABLE_IMPLICIT_TOKEN: bool = _is_true(os.environ.get("HF_HUB_DISABLE_IMPLICIT_TOKEN"))
# Enable fast-download using external dependency "hf_transfer"
# See:
# - https://pypi.org/project/hf-transfer/
# - https://github.com/huggingface/hf_transfer (private)
HF_HUB_ENABLE_HF_TRANSFER: bool = _is_true(os.environ.get("HF_HUB_ENABLE_HF_TRANSFER"))
# UNUSED
# We don't use symlinks in local dir anymore.
HF_HUB_LOCAL_DIR_AUTO_SYMLINK_THRESHOLD: int = (
_as_int(os.environ.get("HF_HUB_LOCAL_DIR_AUTO_SYMLINK_THRESHOLD")) or 5 * 1024 * 1024
)
# Used to override the etag timeout on a system level
HF_HUB_ETAG_TIMEOUT: int = _as_int(os.environ.get("HF_HUB_ETAG_TIMEOUT")) or DEFAULT_ETAG_TIMEOUT
# Used to override the get request timeout on a system level
HF_HUB_DOWNLOAD_TIMEOUT: int = _as_int(os.environ.get("HF_HUB_DOWNLOAD_TIMEOUT")) or DEFAULT_DOWNLOAD_TIMEOUT
# Allows to add information about the requester in the user-agent (eg. partner name)
HF_HUB_USER_AGENT_ORIGIN: Optional[str] = os.environ.get("HF_HUB_USER_AGENT_ORIGIN")
# List frameworks that are handled by the InferenceAPI service. Useful to scan endpoints and check which models are
# deployed and running. Since 95% of the models are using the top 4 frameworks listed below, we scan only those by
# default. We still keep the full list of supported frameworks in case we want to scan all of them.
MAIN_INFERENCE_API_FRAMEWORKS = [
"diffusers",
"sentence-transformers",
"text-generation-inference",
"transformers",
]
ALL_INFERENCE_API_FRAMEWORKS = MAIN_INFERENCE_API_FRAMEWORKS + [
"adapter-transformers",
"allennlp",
"asteroid",
"bertopic",
"doctr",
"espnet",
"fairseq",
"fastai",
"fasttext",
"flair",
"k2",
"keras",
"mindspore",
"nemo",
"open_clip",
"paddlenlp",
"peft",
"pyannote-audio",
"sklearn",
"spacy",
"span-marker",
"speechbrain",
"stanza",
"timm",
]
# If OAuth didn't work after 2 redirects, there's likely a third-party cookie issue in the Space iframe view.
# In this case, we redirect the user to the non-iframe view.
OAUTH_MAX_REDIRECTS = 2
# OAuth-related environment variables injected by the Space
OAUTH_CLIENT_ID = os.environ.get("OAUTH_CLIENT_ID")
OAUTH_CLIENT_SECRET = os.environ.get("OAUTH_CLIENT_SECRET")
OAUTH_SCOPES = os.environ.get("OAUTH_SCOPES")
OPENID_PROVIDER_URL = os.environ.get("OPENID_PROVIDER_URL")
# Xet constants
HUGGINGFACE_HEADER_X_XET_ENDPOINT = "X-Xet-Cas-Url"
HUGGINGFACE_HEADER_X_XET_ACCESS_TOKEN = "X-Xet-Access-Token"
HUGGINGFACE_HEADER_X_XET_EXPIRATION = "X-Xet-Token-Expiration"
HUGGINGFACE_HEADER_X_XET_HASH = "X-Xet-Hash"
HUGGINGFACE_HEADER_X_XET_REFRESH_ROUTE = "X-Xet-Refresh-Route"
HUGGINGFACE_HEADER_LINK_XET_AUTH_KEY = "xet-auth"
default_xet_cache_path = os.path.join(HF_HOME, "xet")
HF_XET_CACHE = os.getenv("HF_XET_CACHE", default_xet_cache_path)
kronos-venv/lib/python3.12/site-packages/huggingface_hub/dataclasses.py
+481
View File
@@ -0,0 +1,481 @@
import inspect
from dataclasses import _MISSING_TYPE, MISSING, Field, field, fields
from functools import wraps
from typing import (
Any,
Callable,
Dict,
List,
Literal,
Optional,
Tuple,
Type,
TypeVar,
Union,
get_args,
get_origin,
overload,
)
from .errors import (
StrictDataclassClassValidationError,
StrictDataclassDefinitionError,
StrictDataclassFieldValidationError,
)
Validator_T = Callable[[Any], None]
T = TypeVar("T")
# The overload decorator helps type checkers understand the different return types
@overload
def strict(cls: Type[T]) -> Type[T]: ...
@overload
def strict(*, accept_kwargs: bool = False) -> Callable[[Type[T]], Type[T]]: ...
def strict(
cls: Optional[Type[T]] = None, *, accept_kwargs: bool = False
) -> Union[Type[T], Callable[[Type[T]], Type[T]]]:
"""
Decorator to add strict validation to a dataclass.
This decorator must be used on top of `@dataclass` to ensure IDEs and static typing tools
recognize the class as a dataclass.
Can be used with or without arguments:
- `@strict`
- `@strict(accept_kwargs=True)`
Args:
cls:
The class to convert to a strict dataclass.
accept_kwargs (`bool`, *optional*):
If True, allows arbitrary keyword arguments in `__init__`. Defaults to False.
Returns:
The enhanced dataclass with strict validation on field assignment.
Example:
```py
>>> from dataclasses import dataclass
>>> from huggingface_hub.dataclasses import as_validated_field, strict, validated_field
>>> @as_validated_field
>>> def positive_int(value: int):
... if not value >= 0:
... raise ValueError(f"Value must be positive, got {value}")
>>> @strict(accept_kwargs=True)
... @dataclass
... class User:
... name: str
... age: int = positive_int(default=10)
# Initialize
>>> User(name="John")
User(name='John', age=10)
# Extra kwargs are accepted
>>> User(name="John", age=30, lastname="Doe")
User(name='John', age=30, *lastname='Doe')
# Invalid type => raises
>>> User(name="John", age="30")
huggingface_hub.errors.StrictDataclassFieldValidationError: Validation error for field 'age':
TypeError: Field 'age' expected int, got str (value: '30')
# Invalid value => raises
>>> User(name="John", age=-1)
huggingface_hub.errors.StrictDataclassFieldValidationError: Validation error for field 'age':
ValueError: Value must be positive, got -1
```
"""
def wrap(cls: Type[T]) -> Type[T]:
if not hasattr(cls, "__dataclass_fields__"):
raise StrictDataclassDefinitionError(
f"Class '{cls.__name__}' must be a dataclass before applying @strict."
)
# List and store validators
field_validators: Dict[str, List[Validator_T]] = {}
for f in fields(cls): # type: ignore [arg-type]
validators = []
validators.append(_create_type_validator(f))
custom_validator = f.metadata.get("validator")
if custom_validator is not None:
if not isinstance(custom_validator, list):
custom_validator = [custom_validator]
for validator in custom_validator:
if not _is_validator(validator):
raise StrictDataclassDefinitionError(
f"Invalid validator for field '{f.name}': {validator}. Must be a callable taking a single argument."
)
validators.extend(custom_validator)
field_validators[f.name] = validators
cls.__validators__ = field_validators # type: ignore
# Override __setattr__ to validate fields on assignment
original_setattr = cls.__setattr__
def __strict_setattr__(self: Any, name: str, value: Any) -> None:
"""Custom __setattr__ method for strict dataclasses."""
# Run all validators
for validator in self.__validators__.get(name, []):
try:
validator(value)
except (ValueError, TypeError) as e:
raise StrictDataclassFieldValidationError(field=name, cause=e) from e
# If validation passed, set the attribute
original_setattr(self, name, value)
cls.__setattr__ = __strict_setattr__ # type: ignore[method-assign]
if accept_kwargs:
# (optional) Override __init__ to accept arbitrary keyword arguments
original_init = cls.__init__
@wraps(original_init)
def __init__(self, **kwargs: Any) -> None:
# Extract only the fields that are part of the dataclass
dataclass_fields = {f.name for f in fields(cls)} # type: ignore [arg-type]
standard_kwargs = {k: v for k, v in kwargs.items() if k in dataclass_fields}
# Call the original __init__ with standard fields
original_init(self, **standard_kwargs)
# Add any additional kwargs as attributes
for name, value in kwargs.items():
if name not in dataclass_fields:
self.__setattr__(name, value)
cls.__init__ = __init__ # type: ignore[method-assign]
# (optional) Override __repr__ to include additional kwargs
original_repr = cls.__repr__
@wraps(original_repr)
def __repr__(self) -> str:
# Call the original __repr__ to get the standard fields
standard_repr = original_repr(self)
# Get additional kwargs
additional_kwargs = [
# add a '*' in front of additional kwargs to let the user know they are not part of the dataclass
f"*{k}={v!r}"
for k, v in self.__dict__.items()
if k not in cls.__dataclass_fields__ # type: ignore [attr-defined]
]
additional_repr = ", ".join(additional_kwargs)
# Combine both representations
return f"{standard_repr[:-1]}, {additional_repr})" if additional_kwargs else standard_repr
cls.__repr__ = __repr__ # type: ignore [method-assign]
# List all public methods starting with `validate_` => class validators.
class_validators = []
for name in dir(cls):
if not name.startswith("validate_"):
continue
method = getattr(cls, name)
if not callable(method):
continue
if len(inspect.signature(method).parameters) != 1:
raise StrictDataclassDefinitionError(
f"Class '{cls.__name__}' has a class validator '{name}' that takes more than one argument."
" Class validators must take only 'self' as an argument. Methods starting with 'validate_'"
" are considered to be class validators."
)
class_validators.append(method)
cls.__class_validators__ = class_validators # type: ignore [attr-defined]
# Add `validate` method to the class, but first check if it already exists
def validate(self: T) -> None:
"""Run class validators on the instance."""
for validator in cls.__class_validators__: # type: ignore [attr-defined]
try:
validator(self)
except (ValueError, TypeError) as e:
raise StrictDataclassClassValidationError(validator=validator.__name__, cause=e) from e
# Hack to be able to raise if `.validate()` already exists except if it was created by this decorator on a parent class
# (in which case we just override it)
validate.__is_defined_by_strict_decorator__ = True # type: ignore [attr-defined]
if hasattr(cls, "validate"):
if not getattr(cls.validate, "__is_defined_by_strict_decorator__", False): # type: ignore [attr-defined]
raise StrictDataclassDefinitionError(
f"Class '{cls.__name__}' already implements a method called 'validate'."
" This method name is reserved when using the @strict decorator on a dataclass."
" If you want to keep your own method, please rename it."
)
cls.validate = validate # type: ignore
# Run class validators after initialization
initial_init = cls.__init__
@wraps(initial_init)
def init_with_validate(self, *args, **kwargs) -> None:
"""Run class validators after initialization."""
initial_init(self, *args, **kwargs) # type: ignore [call-arg]
cls.validate(self) # type: ignore [attr-defined]
setattr(cls, "__init__", init_with_validate)
return cls
# Return wrapped class or the decorator itself
return wrap(cls) if cls is not None else wrap
def validated_field(
validator: Union[List[Validator_T], Validator_T],
default: Union[Any, _MISSING_TYPE] = MISSING,
default_factory: Union[Callable[[], Any], _MISSING_TYPE] = MISSING,
init: bool = True,
repr: bool = True,
hash: Optional[bool] = None,
compare: bool = True,
metadata: Optional[Dict] = None,
**kwargs: Any,
) -> Any:
"""
Create a dataclass field with a custom validator.
Useful to apply several checks to a field. If only applying one rule, check out the [`as_validated_field`] decorator.
Args:
validator (`Callable` or `List[Callable]`):
A method that takes a value as input and raises ValueError/TypeError if the value is invalid.
Can be a list of validators to apply multiple checks.
**kwargs:
Additional arguments to pass to `dataclasses.field()`.
Returns:
A field with the validator attached in metadata
"""
if not isinstance(validator, list):
validator = [validator]
if metadata is None:
metadata = {}
metadata["validator"] = validator
return field( # type: ignore
default=default, # type: ignore [arg-type]
default_factory=default_factory, # type: ignore [arg-type]
init=init,
repr=repr,
hash=hash,
compare=compare,
metadata=metadata,
**kwargs,
)
def as_validated_field(validator: Validator_T):
"""
Decorates a validator function as a [`validated_field`] (i.e. a dataclass field with a custom validator).
Args:
validator (`Callable`):
A method that takes a value as input and raises ValueError/TypeError if the value is invalid.
"""
def _inner(
default: Union[Any, _MISSING_TYPE] = MISSING,
default_factory: Union[Callable[[], Any], _MISSING_TYPE] = MISSING,
init: bool = True,
repr: bool = True,
hash: Optional[bool] = None,
compare: bool = True,
metadata: Optional[Dict] = None,
**kwargs: Any,
):
return validated_field(
validator,
default=default,
default_factory=default_factory,
init=init,
repr=repr,
hash=hash,
compare=compare,
metadata=metadata,
**kwargs,
)
return _inner
def type_validator(name: str, value: Any, expected_type: Any) -> None:
"""Validate that 'value' matches 'expected_type'."""
origin = get_origin(expected_type)
args = get_args(expected_type)
if expected_type is Any:
return
elif validator := _BASIC_TYPE_VALIDATORS.get(origin):
validator(name, value, args)
elif isinstance(expected_type, type): # simple types
_validate_simple_type(name, value, expected_type)
else:
raise TypeError(f"Unsupported type for field '{name}': {expected_type}")
def _validate_union(name: str, value: Any, args: Tuple[Any, ...]) -> None:
"""Validate that value matches one of the types in a Union."""
errors = []
for t in args:
try:
type_validator(name, value, t)
return # Valid if any type matches
except TypeError as e:
errors.append(str(e))
raise TypeError(
f"Field '{name}' with value {repr(value)} doesn't match any type in {args}. Errors: {'; '.join(errors)}"
)
def _validate_literal(name: str, value: Any, args: Tuple[Any, ...]) -> None:
"""Validate Literal type."""
if value not in args:
raise TypeError(f"Field '{name}' expected one of {args}, got {value}")
def _validate_list(name: str, value: Any, args: Tuple[Any, ...]) -> None:
"""Validate List[T] type."""
if not isinstance(value, list):
raise TypeError(f"Field '{name}' expected a list, got {type(value).__name__}")
# Validate each item in the list
item_type = args[0]
for i, item in enumerate(value):
try:
type_validator(f"{name}[{i}]", item, item_type)
except TypeError as e:
raise TypeError(f"Invalid item at index {i} in list '{name}'") from e
def _validate_dict(name: str, value: Any, args: Tuple[Any, ...]) -> None:
"""Validate Dict[K, V] type."""
if not isinstance(value, dict):
raise TypeError(f"Field '{name}' expected a dict, got {type(value).__name__}")
# Validate keys and values
key_type, value_type = args
for k, v in value.items():
try:
type_validator(f"{name}.key", k, key_type)
type_validator(f"{name}[{k!r}]", v, value_type)
except TypeError as e:
raise TypeError(f"Invalid key or value in dict '{name}'") from e
def _validate_tuple(name: str, value: Any, args: Tuple[Any, ...]) -> None:
"""Validate Tuple type."""
if not isinstance(value, tuple):
raise TypeError(f"Field '{name}' expected a tuple, got {type(value).__name__}")
# Handle variable-length tuples: Tuple[T, ...]
if len(args) == 2 and args[1] is Ellipsis:
for i, item in enumerate(value):
try:
type_validator(f"{name}[{i}]", item, args[0])
except TypeError as e:
raise TypeError(f"Invalid item at index {i} in tuple '{name}'") from e
# Handle fixed-length tuples: Tuple[T1, T2, ...]
elif len(args) != len(value):
raise TypeError(f"Field '{name}' expected a tuple of length {len(args)}, got {len(value)}")
else:
for i, (item, expected) in enumerate(zip(value, args)):
try:
type_validator(f"{name}[{i}]", item, expected)
except TypeError as e:
raise TypeError(f"Invalid item at index {i} in tuple '{name}'") from e
def _validate_set(name: str, value: Any, args: Tuple[Any, ...]) -> None:
"""Validate Set[T] type."""
if not isinstance(value, set):
raise TypeError(f"Field '{name}' expected a set, got {type(value).__name__}")
# Validate each item in the set
item_type = args[0]
for i, item in enumerate(value):
try:
type_validator(f"{name} item", item, item_type)
except TypeError as e:
raise TypeError(f"Invalid item in set '{name}'") from e
def _validate_simple_type(name: str, value: Any, expected_type: type) -> None:
"""Validate simple type (int, str, etc.)."""
if not isinstance(value, expected_type):
raise TypeError(
f"Field '{name}' expected {expected_type.__name__}, got {type(value).__name__} (value: {repr(value)})"
)
def _create_type_validator(field: Field) -> Validator_T:
"""Create a type validator function for a field."""
# Hacky: we cannot use a lambda here because of reference issues
def validator(value: Any) -> None:
type_validator(field.name, value, field.type)
return validator
def _is_validator(validator: Any) -> bool:
"""Check if a function is a validator.
A validator is a Callable that can be called with a single positional argument.
The validator can have more arguments with default values.
Basically, returns True if `validator(value)` is possible.
"""
if not callable(validator):
return False
signature = inspect.signature(validator)
parameters = list(signature.parameters.values())
if len(parameters) == 0:
return False
if parameters[0].kind not in (
inspect.Parameter.POSITIONAL_OR_KEYWORD,
inspect.Parameter.POSITIONAL_ONLY,
inspect.Parameter.VAR_POSITIONAL,
):
return False
for parameter in parameters[1:]:
if parameter.default == inspect.Parameter.empty:
return False
return True
_BASIC_TYPE_VALIDATORS = {
Union: _validate_union,
Literal: _validate_literal,
list: _validate_list,
dict: _validate_dict,
tuple: _validate_tuple,
set: _validate_set,
}
__all__ = [
"strict",
"validated_field",
"Validator_T",
"StrictDataclassClassValidationError",
"StrictDataclassDefinitionError",
"StrictDataclassFieldValidationError",
]
kronos-venv/lib/python3.12/site-packages/huggingface_hub/errors.py
+377
View File
@@ -0,0 +1,377 @@
"""Contains all custom errors."""
from pathlib import Path
from typing import Optional, Union
from requests import HTTPError, Response
# CACHE ERRORS
class CacheNotFound(Exception):
"""Exception thrown when the Huggingface cache is not found."""
cache_dir: Union[str, Path]
def __init__(self, msg: str, cache_dir: Union[str, Path], *args, **kwargs):
super().__init__(msg, *args, **kwargs)
self.cache_dir = cache_dir
class CorruptedCacheException(Exception):
"""Exception for any unexpected structure in the Huggingface cache-system."""
# HEADERS ERRORS
class LocalTokenNotFoundError(EnvironmentError):
"""Raised if local token is required but not found."""
# HTTP ERRORS
class OfflineModeIsEnabled(ConnectionError):
"""Raised when a request is made but `HF_HUB_OFFLINE=1` is set as environment variable."""
class HfHubHTTPError(HTTPError):
"""
HTTPError to inherit from for any custom HTTP Error raised in HF Hub.
Any HTTPError is converted at least into a `HfHubHTTPError`. If some information is
sent back by the server, it will be added to the error message.
Added details:
- Request id from "X-Request-Id" header if exists. If not, fallback to "X-Amzn-Trace-Id" header if exists.
- Server error message from the header "X-Error-Message".
- Server error message if we can found one in the response body.
Example:
```py
import requests
from huggingface_hub.utils import get_session, hf_raise_for_status, HfHubHTTPError
response = get_session().post(...)
try:
hf_raise_for_status(response)
except HfHubHTTPError as e:
print(str(e)) # formatted message
e.request_id, e.server_message # details returned by server
# Complete the error message with additional information once it's raised
e.append_to_message("\n`create_commit` expects the repository to exist.")
raise
```
"""
def __init__(self, message: str, response: Optional[Response] = None, *, server_message: Optional[str] = None):
self.request_id = (
response.headers.get("x-request-id") or response.headers.get("X-Amzn-Trace-Id")
if response is not None
else None
)
self.server_message = server_message
super().__init__(
message,
response=response, # type: ignore [arg-type]
request=response.request if response is not None else None, # type: ignore [arg-type]
)
def append_to_message(self, additional_message: str) -> None:
"""Append additional information to the `HfHubHTTPError` initial message."""
self.args = (self.args[0] + additional_message,) + self.args[1:]
# INFERENCE CLIENT ERRORS
class InferenceTimeoutError(HTTPError, TimeoutError):
"""Error raised when a model is unavailable or the request times out."""
# INFERENCE ENDPOINT ERRORS
class InferenceEndpointError(Exception):
"""Generic exception when dealing with Inference Endpoints."""
class InferenceEndpointTimeoutError(InferenceEndpointError, TimeoutError):
"""Exception for timeouts while waiting for Inference Endpoint."""
# SAFETENSORS ERRORS
class SafetensorsParsingError(Exception):
"""Raised when failing to parse a safetensors file metadata.
This can be the case if the file is not a safetensors file or does not respect the specification.
"""
class NotASafetensorsRepoError(Exception):
"""Raised when a repo is not a Safetensors repo i.e. doesn't have either a `model.safetensors` or a
`model.safetensors.index.json` file.
"""
# TEXT GENERATION ERRORS
class TextGenerationError(HTTPError):
"""Generic error raised if text-generation went wrong."""
# Text Generation Inference Errors
class ValidationError(TextGenerationError):
"""Server-side validation error."""
class GenerationError(TextGenerationError):
pass
class OverloadedError(TextGenerationError):
pass
class IncompleteGenerationError(TextGenerationError):
pass
class UnknownError(TextGenerationError):
pass
# VALIDATION ERRORS
class HFValidationError(ValueError):
"""Generic exception thrown by `huggingface_hub` validators.
Inherits from [`ValueError`](https://docs.python.org/3/library/exceptions.html#ValueError).
"""
# FILE METADATA ERRORS
class FileMetadataError(OSError):
"""Error triggered when the metadata of a file on the Hub cannot be retrieved (missing ETag or commit_hash).
Inherits from `OSError` for backward compatibility.
"""
# REPOSITORY ERRORS
class RepositoryNotFoundError(HfHubHTTPError):
"""
Raised when trying to access a hf.co URL with an invalid repository name, or
with a private repo name the user does not have access to.
Example:
```py
>>> from huggingface_hub import model_info
>>> model_info("<non_existent_repository>")
(...)
huggingface_hub.utils._errors.RepositoryNotFoundError: 401 Client Error. (Request ID: PvMw_VjBMjVdMz53WKIzP)
Repository Not Found for url: https://huggingface.co/api/models/%3Cnon_existent_repository%3E.
Please make sure you specified the correct `repo_id` and `repo_type`.
If the repo is private, make sure you are authenticated.
Invalid username or password.
```
"""
class GatedRepoError(RepositoryNotFoundError):
"""
Raised when trying to access a gated repository for which the user is not on the
authorized list.
Note: derives from `RepositoryNotFoundError` to ensure backward compatibility.
Example:
```py
>>> from huggingface_hub import model_info
>>> model_info("<gated_repository>")
(...)
huggingface_hub.utils._errors.GatedRepoError: 403 Client Error. (Request ID: ViT1Bf7O_026LGSQuVqfa)
Cannot access gated repo for url https://huggingface.co/api/models/ardent-figment/gated-model.
Access to model ardent-figment/gated-model is restricted and you are not in the authorized list.
Visit https://huggingface.co/ardent-figment/gated-model to ask for access.
```
"""
class DisabledRepoError(HfHubHTTPError):
"""
Raised when trying to access a repository that has been disabled by its author.
Example:
```py
>>> from huggingface_hub import dataset_info
>>> dataset_info("laion/laion-art")
(...)
huggingface_hub.utils._errors.DisabledRepoError: 403 Client Error. (Request ID: Root=1-659fc3fa-3031673e0f92c71a2260dbe2;bc6f4dfb-b30a-4862-af0a-5cfe827610d8)
Cannot access repository for url https://huggingface.co/api/datasets/laion/laion-art.
Access to this resource is disabled.
```
"""
# REVISION ERROR
class RevisionNotFoundError(HfHubHTTPError):
"""
Raised when trying to access a hf.co URL with a valid repository but an invalid
revision.
Example:
```py
>>> from huggingface_hub import hf_hub_download
>>> hf_hub_download('bert-base-cased', 'config.json', revision='<non-existent-revision>')
(...)
huggingface_hub.utils._errors.RevisionNotFoundError: 404 Client Error. (Request ID: Mwhe_c3Kt650GcdKEFomX)
Revision Not Found for url: https://huggingface.co/bert-base-cased/resolve/%3Cnon-existent-revision%3E/config.json.
```
"""
# ENTRY ERRORS
class EntryNotFoundError(HfHubHTTPError):
"""
Raised when trying to access a hf.co URL with a valid repository and revision
but an invalid filename.
Example:
```py
>>> from huggingface_hub import hf_hub_download
>>> hf_hub_download('bert-base-cased', '<non-existent-file>')
(...)
huggingface_hub.utils._errors.EntryNotFoundError: 404 Client Error. (Request ID: 53pNl6M0MxsnG5Sw8JA6x)
Entry Not Found for url: https://huggingface.co/bert-base-cased/resolve/main/%3Cnon-existent-file%3E.
```
"""
class LocalEntryNotFoundError(EntryNotFoundError, FileNotFoundError, ValueError):
"""
Raised when trying to access a file or snapshot that is not on the disk when network is
disabled or unavailable (connection issue). The entry may exist on the Hub.
Note: `ValueError` type is to ensure backward compatibility.
Note: `LocalEntryNotFoundError` derives from `HTTPError` because of `EntryNotFoundError`
even when it is not a network issue.
Example:
```py
>>> from huggingface_hub import hf_hub_download
>>> hf_hub_download('bert-base-cased', '<non-cached-file>', local_files_only=True)
(...)
huggingface_hub.utils._errors.LocalEntryNotFoundError: Cannot find the requested files in the disk cache and outgoing traffic has been disabled. To enable hf.co look-ups and downloads online, set 'local_files_only' to False.
```
"""
def __init__(self, message: str):
super().__init__(message, response=None)
# REQUEST ERROR
class BadRequestError(HfHubHTTPError, ValueError):
"""
Raised by `hf_raise_for_status` when the server returns a HTTP 400 error.
Example:
```py
>>> resp = requests.post("hf.co/api/check", ...)
>>> hf_raise_for_status(resp, endpoint_name="check")
huggingface_hub.utils._errors.BadRequestError: Bad request for check endpoint: {details} (Request ID: XXX)
```
"""
# DDUF file format ERROR
class DDUFError(Exception):
"""Base exception for errors related to the DDUF format."""
class DDUFCorruptedFileError(DDUFError):
"""Exception thrown when the DDUF file is corrupted."""
class DDUFExportError(DDUFError):
"""Base exception for errors during DDUF export."""
class DDUFInvalidEntryNameError(DDUFExportError):
"""Exception thrown when the entry name is invalid."""
# STRICT DATACLASSES ERRORS
class StrictDataclassError(Exception):
"""Base exception for strict dataclasses."""
class StrictDataclassDefinitionError(StrictDataclassError):
"""Exception thrown when a strict dataclass is defined incorrectly."""
class StrictDataclassFieldValidationError(StrictDataclassError):
"""Exception thrown when a strict dataclass fails validation for a given field."""
def __init__(self, field: str, cause: Exception):
error_message = f"Validation error for field '{field}':"
error_message += f"\n {cause.__class__.__name__}: {cause}"
super().__init__(error_message)
class StrictDataclassClassValidationError(StrictDataclassError):
"""Exception thrown when a strict dataclass fails validation on a class validator."""
def __init__(self, validator: str, cause: Exception):
error_message = f"Class validation error for validator '{validator}':"
error_message += f"\n {cause.__class__.__name__}: {cause}"
super().__init__(error_message)
# XET ERRORS
class XetError(Exception):
"""Base exception for errors related to Xet Storage."""
class XetAuthorizationError(XetError):
"""Exception thrown when the user does not have the right authorization to use Xet Storage."""
class XetRefreshTokenError(XetError):
"""Exception thrown when the refresh token is invalid."""
class XetDownloadError(Exception):
"""Exception thrown when the download from Xet Storage fails."""
kronos-venv/lib/python3.12/site-packages/huggingface_hub/fastai_utils.py
+425
View File
@@ -0,0 +1,425 @@
import json
import os
from pathlib import Path
from pickle import DEFAULT_PROTOCOL, PicklingError
from typing import Any, Dict, List, Optional, Union
from packaging import version
from huggingface_hub import constants, snapshot_download
from huggingface_hub.hf_api import HfApi
from huggingface_hub.utils import (
SoftTemporaryDirectory,
get_fastai_version,
get_fastcore_version,
get_python_version,
)
from .utils import logging, validate_hf_hub_args
from .utils._runtime import _PY_VERSION # noqa: F401 # for backward compatibility...
logger = logging.get_logger(__name__)
def _check_fastai_fastcore_versions(
fastai_min_version: str = "2.4",
fastcore_min_version: str = "1.3.27",
):
"""
Checks that the installed fastai and fastcore versions are compatible for pickle serialization.
Args:
fastai_min_version (`str`, *optional*):
The minimum fastai version supported.
fastcore_min_version (`str`, *optional*):
The minimum fastcore version supported.
<Tip>
Raises the following error:
- [`ImportError`](https://docs.python.org/3/library/exceptions.html#ImportError)
if the fastai or fastcore libraries are not available or are of an invalid version.
</Tip>
"""
if (get_fastcore_version() or get_fastai_version()) == "N/A":
raise ImportError(
f"fastai>={fastai_min_version} and fastcore>={fastcore_min_version} are"
f" required. Currently using fastai=={get_fastai_version()} and"
f" fastcore=={get_fastcore_version()}."
)
current_fastai_version = version.Version(get_fastai_version())
current_fastcore_version = version.Version(get_fastcore_version())
if current_fastai_version < version.Version(fastai_min_version):
raise ImportError(
"`push_to_hub_fastai` and `from_pretrained_fastai` require a"
f" fastai>={fastai_min_version} version, but you are using fastai version"
f" {get_fastai_version()} which is incompatible. Upgrade with `pip install"
" fastai==2.5.6`."
)
if current_fastcore_version < version.Version(fastcore_min_version):
raise ImportError(
"`push_to_hub_fastai` and `from_pretrained_fastai` require a"
f" fastcore>={fastcore_min_version} version, but you are using fastcore"
f" version {get_fastcore_version()} which is incompatible. Upgrade with"
" `pip install fastcore==1.3.27`."
)
def _check_fastai_fastcore_pyproject_versions(
storage_folder: str,
fastai_min_version: str = "2.4",
fastcore_min_version: str = "1.3.27",
):
"""
Checks that the `pyproject.toml` file in the directory `storage_folder` has fastai and fastcore versions
that are compatible with `from_pretrained_fastai` and `push_to_hub_fastai`. If `pyproject.toml` does not exist
or does not contain versions for fastai and fastcore, then it logs a warning.
Args:
storage_folder (`str`):
Folder to look for the `pyproject.toml` file.
fastai_min_version (`str`, *optional*):
The minimum fastai version supported.
fastcore_min_version (`str`, *optional*):
The minimum fastcore version supported.
<Tip>
Raises the following errors:
- [`ImportError`](https://docs.python.org/3/library/exceptions.html#ImportError)
if the `toml` module is not installed.
- [`ImportError`](https://docs.python.org/3/library/exceptions.html#ImportError)
if the `pyproject.toml` indicates a lower than minimum supported version of fastai or fastcore.
</Tip>
"""
try:
import toml
except ModuleNotFoundError:
raise ImportError(
"`push_to_hub_fastai` and `from_pretrained_fastai` require the toml module."
" Install it with `pip install toml`."
)
# Checks that a `pyproject.toml`, with `build-system` and `requires` sections, exists in the repository. If so, get a list of required packages.
if not os.path.isfile(f"{storage_folder}/pyproject.toml"):
logger.warning(
"There is no `pyproject.toml` in the repository that contains the fastai"
" `Learner`. The `pyproject.toml` would allow us to verify that your fastai"
" and fastcore versions are compatible with those of the model you want to"
" load."
)
return
pyproject_toml = toml.load(f"{storage_folder}/pyproject.toml")
if "build-system" not in pyproject_toml.keys():
logger.warning(
"There is no `build-system` section in the pyproject.toml of the repository"
" that contains the fastai `Learner`. The `build-system` would allow us to"
" verify that your fastai and fastcore versions are compatible with those"
" of the model you want to load."
)
return
build_system_toml = pyproject_toml["build-system"]
if "requires" not in build_system_toml.keys():
logger.warning(
"There is no `requires` section in the pyproject.toml of the repository"
" that contains the fastai `Learner`. The `requires` would allow us to"
" verify that your fastai and fastcore versions are compatible with those"
" of the model you want to load."
)
return
package_versions = build_system_toml["requires"]
# Extracts contains fastai and fastcore versions from `pyproject.toml` if available.
# If the package is specified but not the version (e.g. "fastai" instead of "fastai=2.4"), the default versions are the highest.
fastai_packages = [pck for pck in package_versions if pck.startswith("fastai")]
if len(fastai_packages) == 0:
logger.warning("The repository does not have a fastai version specified in the `pyproject.toml`.")
# fastai_version is an empty string if not specified
else:
fastai_version = str(fastai_packages[0]).partition("=")[2]
if fastai_version != "" and version.Version(fastai_version) < version.Version(fastai_min_version):
raise ImportError(
"`from_pretrained_fastai` requires"
f" fastai>={fastai_min_version} version but the model to load uses"
f" {fastai_version} which is incompatible."
)
fastcore_packages = [pck for pck in package_versions if pck.startswith("fastcore")]
if len(fastcore_packages) == 0:
logger.warning("The repository does not have a fastcore version specified in the `pyproject.toml`.")
# fastcore_version is an empty string if not specified
else:
fastcore_version = str(fastcore_packages[0]).partition("=")[2]
if fastcore_version != "" and version.Version(fastcore_version) < version.Version(fastcore_min_version):
raise ImportError(
"`from_pretrained_fastai` requires"
f" fastcore>={fastcore_min_version} version, but you are using fastcore"
f" version {fastcore_version} which is incompatible."
)
README_TEMPLATE = """---
tags:
- fastai
---
# Amazing!
🥳 Congratulations on hosting your fastai model on the Hugging Face Hub!
# Some next steps
1. Fill out this model card with more information (see the template below and the [documentation here](https://huggingface.co/docs/hub/model-repos))!
2. Create a demo in Gradio or Streamlit using 🤗 Spaces ([documentation here](https://huggingface.co/docs/hub/spaces)).
3. Join the fastai community on the [Fastai Discord](https://discord.com/invite/YKrxeNn)!
Greetings fellow fastlearner 🤝! Don't forget to delete this content from your model card.
---
# Model card
## Model description
More information needed
## Intended uses & limitations
More information needed
## Training and evaluation data
More information needed
"""
PYPROJECT_TEMPLATE = f"""[build-system]
requires = ["setuptools>=40.8.0", "wheel", "python={get_python_version()}", "fastai={get_fastai_version()}", "fastcore={get_fastcore_version()}"]
build-backend = "setuptools.build_meta:__legacy__"
"""
def _create_model_card(repo_dir: Path):
"""
Creates a model card for the repository.
Args:
repo_dir (`Path`):
Directory where model card is created.
"""
readme_path = repo_dir / "README.md"
if not readme_path.exists():
with readme_path.open("w", encoding="utf-8") as f:
f.write(README_TEMPLATE)
def _create_model_pyproject(repo_dir: Path):
"""
Creates a `pyproject.toml` for the repository.
Args:
repo_dir (`Path`):
Directory where `pyproject.toml` is created.
"""
pyproject_path = repo_dir / "pyproject.toml"
if not pyproject_path.exists():
with pyproject_path.open("w", encoding="utf-8") as f:
f.write(PYPROJECT_TEMPLATE)
def _save_pretrained_fastai(
learner,
save_directory: Union[str, Path],
config: Optional[Dict[str, Any]] = None,
):
"""
Saves a fastai learner to `save_directory` in pickle format using the default pickle protocol for the version of python used.
Args:
learner (`Learner`):
The `fastai.Learner` you'd like to save.
save_directory (`str` or `Path`):
Specific directory in which you want to save the fastai learner.
config (`dict`, *optional*):
Configuration object. Will be uploaded as a .json file. Example: 'https://huggingface.co/espejelomar/fastai-pet-breeds-classification/blob/main/config.json'.
<Tip>
Raises the following error:
- [`RuntimeError`](https://docs.python.org/3/library/exceptions.html#RuntimeError)
if the config file provided is not a dictionary.
</Tip>
"""
_check_fastai_fastcore_versions()
os.makedirs(save_directory, exist_ok=True)
# if the user provides config then we update it with the fastai and fastcore versions in CONFIG_TEMPLATE.
if config is not None:
if not isinstance(config, dict):
raise RuntimeError(f"Provided config should be a dict. Got: '{type(config)}'")
path = os.path.join(save_directory, constants.CONFIG_NAME)
with open(path, "w") as f:
json.dump(config, f)
_create_model_card(Path(save_directory))
_create_model_pyproject(Path(save_directory))
# learner.export saves the model in `self.path`.
learner.path = Path(save_directory)
os.makedirs(save_directory, exist_ok=True)
try:
learner.export(
fname="model.pkl",
pickle_protocol=DEFAULT_PROTOCOL,
)
except PicklingError:
raise PicklingError(
"You are using a lambda function, i.e., an anonymous function. `pickle`"
" cannot pickle function objects and requires that all functions have"
" names. One possible solution is to name the function."
)
@validate_hf_hub_args
def from_pretrained_fastai(
repo_id: str,
revision: Optional[str] = None,
):
"""
Load pretrained fastai model from the Hub or from a local directory.
Args:
repo_id (`str`):
The location where the pickled fastai.Learner is. It can be either of the two:
- Hosted on the Hugging Face Hub. E.g.: 'espejelomar/fatai-pet-breeds-classification' or 'distilgpt2'.
You can add a `revision` by appending `@` at the end of `repo_id`. E.g.: `dbmdz/bert-base-german-cased@main`.
Revision is the specific model version to use. Since we use a git-based system for storing models and other
artifacts on the Hugging Face Hub, it can be a branch name, a tag name, or a commit id.
- Hosted locally. `repo_id` would be a directory containing the pickle and a pyproject.toml
indicating the fastai and fastcore versions used to build the `fastai.Learner`. E.g.: `./my_model_directory/`.
revision (`str`, *optional*):
Revision at which the repo's files are downloaded. See documentation of `snapshot_download`.
Returns:
The `fastai.Learner` model in the `repo_id` repo.
"""
_check_fastai_fastcore_versions()
# Load the `repo_id` repo.
# `snapshot_download` returns the folder where the model was stored.
# `cache_dir` will be the default '/root/.cache/huggingface/hub'
if not os.path.isdir(repo_id):
storage_folder = snapshot_download(
repo_id=repo_id,
revision=revision,
library_name="fastai",
library_version=get_fastai_version(),
)
else:
storage_folder = repo_id
_check_fastai_fastcore_pyproject_versions(storage_folder)
from fastai.learner import load_learner # type: ignore
return load_learner(os.path.join(storage_folder, "model.pkl"))
@validate_hf_hub_args
def push_to_hub_fastai(
learner,
*,
repo_id: str,
commit_message: str = "Push FastAI model using huggingface_hub.",
private: Optional[bool] = None,
token: Optional[str] = None,
config: Optional[dict] = None,
branch: Optional[str] = None,
create_pr: Optional[bool] = None,
allow_patterns: Optional[Union[List[str], str]] = None,
ignore_patterns: Optional[Union[List[str], str]] = None,
delete_patterns: Optional[Union[List[str], str]] = None,
api_endpoint: Optional[str] = None,
):
"""
Upload learner checkpoint files to the Hub.
Use `allow_patterns` and `ignore_patterns` to precisely filter which files should be pushed to the hub. Use
`delete_patterns` to delete existing remote files in the same commit. See [`upload_folder`] reference for more
details.
Args:
learner (`Learner`):
The `fastai.Learner' you'd like to push to the Hub.
repo_id (`str`):
The repository id for your model in Hub in the format of "namespace/repo_name". The namespace can be your individual account or an organization to which you have write access (for example, 'stanfordnlp/stanza-de').
commit_message (`str`, *optional*):
Message to commit while pushing. Will default to :obj:`"add model"`.
private (`bool`, *optional*):
Whether or not the repository created should be private.
If `None` (default), will default to been public except if the organization's default is private.
token (`str`, *optional*):
The Hugging Face account token to use as HTTP bearer authorization for remote files. If :obj:`None`, the token will be asked by a prompt.
config (`dict`, *optional*):
Configuration object to be saved alongside the model weights.
branch (`str`, *optional*):
The git branch on which to push the model. This defaults to
the default branch as specified in your repository, which
defaults to `"main"`.
create_pr (`boolean`, *optional*):
Whether or not to create a Pull Request from `branch` with that commit.
Defaults to `False`.
api_endpoint (`str`, *optional*):
The API endpoint to use when pushing the model to the hub.
allow_patterns (`List[str]` or `str`, *optional*):
If provided, only files matching at least one pattern are pushed.
ignore_patterns (`List[str]` or `str`, *optional*):
If provided, files matching any of the patterns are not pushed.
delete_patterns (`List[str]` or `str`, *optional*):
If provided, remote files matching any of the patterns will be deleted from the repo.
Returns:
The url of the commit of your model in the given repository.
<Tip>
Raises the following error:
- [`ValueError`](https://docs.python.org/3/library/exceptions.html#ValueError)
if the user is not log on to the Hugging Face Hub.
</Tip>
"""
_check_fastai_fastcore_versions()
api = HfApi(endpoint=api_endpoint)
repo_id = api.create_repo(repo_id=repo_id, token=token, private=private, exist_ok=True).repo_id
# Push the files to the repo in a single commit
with SoftTemporaryDirectory() as tmp:
saved_path = Path(tmp) / repo_id
_save_pretrained_fastai(learner, saved_path, config=config)
return api.upload_folder(
repo_id=repo_id,
token=token,
folder_path=saved_path,
commit_message=commit_message,
revision=branch,
create_pr=create_pr,
allow_patterns=allow_patterns,
ignore_patterns=ignore_patterns,
delete_patterns=delete_patterns,
)
kronos-venv/lib/python3.12/site-packages/huggingface_hub/file_download.py
+1803
View File
File diff suppressed because it is too large Load Diff
kronos-venv/lib/python3.12/site-packages/huggingface_hub/hf_api.py
+10097
View File
File diff suppressed because it is too large Load Diff
kronos-venv/lib/python3.12/site-packages/huggingface_hub/hf_file_system.py
+1142
View File
File diff suppressed because it is too large Load Diff
kronos-venv/lib/python3.12/site-packages/huggingface_hub/hub_mixin.py
+851
View File
@@ -0,0 +1,851 @@
import inspect
import json
import os
from dataclasses import Field, asdict, dataclass, is_dataclass
from pathlib import Path
from typing import Any, Callable, ClassVar, Dict, List, Optional, Protocol, Tuple, Type, TypeVar, Union
import packaging.version
from . import constants
from .errors import EntryNotFoundError, HfHubHTTPError
from .file_download import hf_hub_download
from .hf_api import HfApi
from .repocard import ModelCard, ModelCardData
from .utils import (
SoftTemporaryDirectory,
is_jsonable,
is_safetensors_available,
is_simple_optional_type,
is_torch_available,
logging,
unwrap_simple_optional_type,
validate_hf_hub_args,
)
if is_torch_available():
import torch # type: ignore
if is_safetensors_available():
import safetensors
from safetensors.torch import load_model as load_model_as_safetensor
from safetensors.torch import save_model as save_model_as_safetensor
logger = logging.get_logger(__name__)
# Type alias for dataclass instances, copied from https://github.com/python/typeshed/blob/9f28171658b9ca6c32a7cb93fbb99fc92b17858b/stdlib/_typeshed/__init__.pyi#L349
class DataclassInstance(Protocol):
__dataclass_fields__: ClassVar[Dict[str, Field]]
# Generic variable that is either ModelHubMixin or a subclass thereof
T = TypeVar("T", bound="ModelHubMixin")
# Generic variable to represent an args type
ARGS_T = TypeVar("ARGS_T")
ENCODER_T = Callable[[ARGS_T], Any]
DECODER_T = Callable[[Any], ARGS_T]
CODER_T = Tuple[ENCODER_T, DECODER_T]
DEFAULT_MODEL_CARD = """
---
# For reference on model card metadata, see the spec: https://github.com/huggingface/hub-docs/blob/main/modelcard.md?plain=1
# Doc / guide: https://huggingface.co/docs/hub/model-cards
{{ card_data }}
---
This model has been pushed to the Hub using the [PytorchModelHubMixin](https://huggingface.co/docs/huggingface_hub/package_reference/mixins#huggingface_hub.PyTorchModelHubMixin) integration:
- Code: {{ repo_url | default("[More Information Needed]", true) }}
- Paper: {{ paper_url | default("[More Information Needed]", true) }}
- Docs: {{ docs_url | default("[More Information Needed]", true) }}
"""
@dataclass
class MixinInfo:
model_card_template: str
model_card_data: ModelCardData
docs_url: Optional[str] = None
paper_url: Optional[str] = None
repo_url: Optional[str] = None
class ModelHubMixin:
"""
A generic mixin to integrate ANY machine learning framework with the Hub.
To integrate your framework, your model class must inherit from this class. Custom logic for saving/loading models
have to be overwritten in [`_from_pretrained`] and [`_save_pretrained`]. [`PyTorchModelHubMixin`] is a good example
of mixin integration with the Hub. Check out our [integration guide](../guides/integrations) for more instructions.
When inheriting from [`ModelHubMixin`], you can define class-level attributes. These attributes are not passed to
`__init__` but to the class definition itself. This is useful to define metadata about the library integrating
[`ModelHubMixin`].
For more details on how to integrate the mixin with your library, checkout the [integration guide](../guides/integrations).
Args:
repo_url (`str`, *optional*):
URL of the library repository. Used to generate model card.
paper_url (`str`, *optional*):
URL of the library paper. Used to generate model card.
docs_url (`str`, *optional*):
URL of the library documentation. Used to generate model card.
model_card_template (`str`, *optional*):
Template of the model card. Used to generate model card. Defaults to a generic template.
language (`str` or `List[str]`, *optional*):
Language supported by the library. Used to generate model card.
library_name (`str`, *optional*):
Name of the library integrating ModelHubMixin. Used to generate model card.
license (`str`, *optional*):
License of the library integrating ModelHubMixin. Used to generate model card.
E.g: "apache-2.0"
license_name (`str`, *optional*):
Name of the library integrating ModelHubMixin. Used to generate model card.
Only used if `license` is set to `other`.
E.g: "coqui-public-model-license".
license_link (`str`, *optional*):
URL to the license of the library integrating ModelHubMixin. Used to generate model card.
Only used if `license` is set to `other` and `license_name` is set.
E.g: "https://coqui.ai/cpml".
pipeline_tag (`str`, *optional*):
Tag of the pipeline. Used to generate model card. E.g. "text-classification".
tags (`List[str]`, *optional*):
Tags to be added to the model card. Used to generate model card. E.g. ["computer-vision"]
coders (`Dict[Type, Tuple[Callable, Callable]]`, *optional*):
Dictionary of custom types and their encoders/decoders. Used to encode/decode arguments that are not
jsonable by default. E.g dataclasses, argparse.Namespace, OmegaConf, etc.
Example:
```python
>>> from huggingface_hub import ModelHubMixin
# Inherit from ModelHubMixin
>>> class MyCustomModel(
... ModelHubMixin,
... library_name="my-library",
... tags=["computer-vision"],
... repo_url="https://github.com/huggingface/my-cool-library",
... paper_url="https://arxiv.org/abs/2304.12244",
... docs_url="https://huggingface.co/docs/my-cool-library",
... # ^ optional metadata to generate model card
... ):
... def __init__(self, size: int = 512, device: str = "cpu"):
... # define how to initialize your model
... super().__init__()
... ...
...
... def _save_pretrained(self, save_directory: Path) -> None:
... # define how to serialize your model
... ...
...
... @classmethod
... def from_pretrained(
... cls: Type[T],
... pretrained_model_name_or_path: Union[str, Path],
... *,
... force_download: bool = False,
... resume_download: Optional[bool] = None,
... proxies: Optional[Dict] = None,
... token: Optional[Union[str, bool]] = None,
... cache_dir: Optional[Union[str, Path]] = None,
... local_files_only: bool = False,
... revision: Optional[str] = None,
... **model_kwargs,
... ) -> T:
... # define how to deserialize your model
... ...
>>> model = MyCustomModel(size=256, device="gpu")
# Save model weights to local directory
>>> model.save_pretrained("my-awesome-model")
# Push model weights to the Hub
>>> model.push_to_hub("my-awesome-model")
# Download and initialize weights from the Hub
>>> reloaded_model = MyCustomModel.from_pretrained("username/my-awesome-model")
>>> reloaded_model.size
256
# Model card has been correctly populated
>>> from huggingface_hub import ModelCard
>>> card = ModelCard.load("username/my-awesome-model")
>>> card.data.tags
["x-custom-tag", "pytorch_model_hub_mixin", "model_hub_mixin"]
>>> card.data.library_name
"my-library"
```
"""
_hub_mixin_config: Optional[Union[dict, DataclassInstance]] = None
# ^ optional config attribute automatically set in `from_pretrained`
_hub_mixin_info: MixinInfo
# ^ information about the library integrating ModelHubMixin (used to generate model card)
_hub_mixin_inject_config: bool # whether `_from_pretrained` expects `config` or not
_hub_mixin_init_parameters: Dict[str, inspect.Parameter] # __init__ parameters
_hub_mixin_jsonable_default_values: Dict[str, Any] # default values for __init__ parameters
_hub_mixin_jsonable_custom_types: Tuple[Type, ...] # custom types that can be encoded/decoded
_hub_mixin_coders: Dict[Type, CODER_T] # encoders/decoders for custom types
# ^ internal values to handle config
def __init_subclass__(
cls,
*,
# Generic info for model card
repo_url: Optional[str] = None,
paper_url: Optional[str] = None,
docs_url: Optional[str] = None,
# Model card template
model_card_template: str = DEFAULT_MODEL_CARD,
# Model card metadata
language: Optional[List[str]] = None,
library_name: Optional[str] = None,
license: Optional[str] = None,
license_name: Optional[str] = None,
license_link: Optional[str] = None,
pipeline_tag: Optional[str] = None,
tags: Optional[List[str]] = None,
# How to encode/decode arguments with custom type into a JSON config?
coders: Optional[
Dict[Type, CODER_T]
# Key is a type.
# Value is a tuple (encoder, decoder).
# Example: {MyCustomType: (lambda x: x.value, lambda data: MyCustomType(data))}
] = None,
) -> None:
"""Inspect __init__ signature only once when subclassing + handle modelcard."""
super().__init_subclass__()
# Will be reused when creating modelcard
tags = tags or []
tags.append("model_hub_mixin")
# Initialize MixinInfo if not existent
info = MixinInfo(model_card_template=model_card_template, model_card_data=ModelCardData())
# If parent class has a MixinInfo, inherit from it as a copy
if hasattr(cls, "_hub_mixin_info"):
# Inherit model card template from parent class if not explicitly set
if model_card_template == DEFAULT_MODEL_CARD:
info.model_card_template = cls._hub_mixin_info.model_card_template
# Inherit from parent model card data
info.model_card_data = ModelCardData(**cls._hub_mixin_info.model_card_data.to_dict())
# Inherit other info
info.docs_url = cls._hub_mixin_info.docs_url
info.paper_url = cls._hub_mixin_info.paper_url
info.repo_url = cls._hub_mixin_info.repo_url
cls._hub_mixin_info = info
# Update MixinInfo with metadata
if model_card_template is not None and model_card_template != DEFAULT_MODEL_CARD:
info.model_card_template = model_card_template
if repo_url is not None:
info.repo_url = repo_url
if paper_url is not None:
info.paper_url = paper_url
if docs_url is not None:
info.docs_url = docs_url
if language is not None:
info.model_card_data.language = language
if library_name is not None:
info.model_card_data.library_name = library_name
if license is not None:
info.model_card_data.license = license
if license_name is not None:
info.model_card_data.license_name = license_name
if license_link is not None:
info.model_card_data.license_link = license_link
if pipeline_tag is not None:
info.model_card_data.pipeline_tag = pipeline_tag
if tags is not None:
if info.model_card_data.tags is not None:
info.model_card_data.tags.extend(tags)
else:
info.model_card_data.tags = tags
info.model_card_data.tags = sorted(set(info.model_card_data.tags))
# Handle encoders/decoders for args
cls._hub_mixin_coders = coders or {}
cls._hub_mixin_jsonable_custom_types = tuple(cls._hub_mixin_coders.keys())
# Inspect __init__ signature to handle config
cls._hub_mixin_init_parameters = dict(inspect.signature(cls.__init__).parameters)
cls._hub_mixin_jsonable_default_values = {
param.name: cls._encode_arg(param.default)
for param in cls._hub_mixin_init_parameters.values()
if param.default is not inspect.Parameter.empty and cls._is_jsonable(param.default)
}
cls._hub_mixin_inject_config = "config" in inspect.signature(cls._from_pretrained).parameters
def __new__(cls: Type[T], *args, **kwargs) -> T:
"""Create a new instance of the class and handle config.
3 cases:
- If `self._hub_mixin_config` is already set, do nothing.
- If `config` is passed as a dataclass, set it as `self._hub_mixin_config`.
- Otherwise, build `self._hub_mixin_config` from default values and passed values.
"""
instance = super().__new__(cls)
# If `config` is already set, return early
if instance._hub_mixin_config is not None:
return instance
# Infer passed values
passed_values = {
**{
key: value
for key, value in zip(
# [1:] to skip `self` parameter
list(cls._hub_mixin_init_parameters)[1:],
args,
)
},
**kwargs,
}
# If config passed as dataclass => set it and return early
if is_dataclass(passed_values.get("config")):
instance._hub_mixin_config = passed_values["config"]
return instance
# Otherwise, build config from default + passed values
init_config = {
# default values
**cls._hub_mixin_jsonable_default_values,
# passed values
**{
key: cls._encode_arg(value) # Encode custom types as jsonable value
for key, value in passed_values.items()
if instance._is_jsonable(value) # Only if jsonable or we have a custom encoder
},
}
passed_config = init_config.pop("config", {})
# Populate `init_config` with provided config
if isinstance(passed_config, dict):
init_config.update(passed_config)
# Set `config` attribute and return
if init_config != {}:
instance._hub_mixin_config = init_config
return instance
@classmethod
def _is_jsonable(cls, value: Any) -> bool:
"""Check if a value is JSON serializable."""
if is_dataclass(value):
return True
if isinstance(value, cls._hub_mixin_jsonable_custom_types):
return True
return is_jsonable(value)
@classmethod
def _encode_arg(cls, arg: Any) -> Any:
"""Encode an argument into a JSON serializable format."""
if is_dataclass(arg):
return asdict(arg) # type: ignore[arg-type]
for type_, (encoder, _) in cls._hub_mixin_coders.items():
if isinstance(arg, type_):
if arg is None:
return None
return encoder(arg)
return arg
@classmethod
def _decode_arg(cls, expected_type: Type[ARGS_T], value: Any) -> Optional[ARGS_T]:
"""Decode a JSON serializable value into an argument."""
if is_simple_optional_type(expected_type):
if value is None:
return None
expected_type = unwrap_simple_optional_type(expected_type)
# Dataclass => handle it
if is_dataclass(expected_type):
return _load_dataclass(expected_type, value) # type: ignore[return-value]
# Otherwise => check custom decoders
for type_, (_, decoder) in cls._hub_mixin_coders.items():
if inspect.isclass(expected_type) and issubclass(expected_type, type_):
return decoder(value)
# Otherwise => don't decode
return value
def save_pretrained(
self,
save_directory: Union[str, Path],
*,
config: Optional[Union[dict, DataclassInstance]] = None,
repo_id: Optional[str] = None,
push_to_hub: bool = False,
model_card_kwargs: Optional[Dict[str, Any]] = None,
**push_to_hub_kwargs,
) -> Optional[str]:
"""
Save weights in local directory.
Args:
save_directory (`str` or `Path`):
Path to directory in which the model weights and configuration will be saved.
config (`dict` or `DataclassInstance`, *optional*):
Model configuration specified as a key/value dictionary or a dataclass instance.
push_to_hub (`bool`, *optional*, defaults to `False`):
Whether or not to push your model to the Huggingface Hub after saving it.
repo_id (`str`, *optional*):
ID of your repository on the Hub. Used only if `push_to_hub=True`. Will default to the folder name if
not provided.
model_card_kwargs (`Dict[str, Any]`, *optional*):
Additional arguments passed to the model card template to customize the model card.
push_to_hub_kwargs:
Additional key word arguments passed along to the [`~ModelHubMixin.push_to_hub`] method.
Returns:
`str` or `None`: url of the commit on the Hub if `push_to_hub=True`, `None` otherwise.
"""
save_directory = Path(save_directory)
save_directory.mkdir(parents=True, exist_ok=True)
# Remove config.json if already exists. After `_save_pretrained` we don't want to overwrite config.json
# as it might have been saved by the custom `_save_pretrained` already. However we do want to overwrite
# an existing config.json if it was not saved by `_save_pretrained`.
config_path = save_directory / constants.CONFIG_NAME
config_path.unlink(missing_ok=True)
# save model weights/files (framework-specific)
self._save_pretrained(save_directory)
# save config (if provided and if not serialized yet in `_save_pretrained`)
if config is None:
config = self._hub_mixin_config
if config is not None:
if is_dataclass(config):
config = asdict(config) # type: ignore[arg-type]
if not config_path.exists():
config_str = json.dumps(config, sort_keys=True, indent=2)
config_path.write_text(config_str)
# save model card
model_card_path = save_directory / "README.md"
model_card_kwargs = model_card_kwargs if model_card_kwargs is not None else {}
if not model_card_path.exists(): # do not overwrite if already exists
self.generate_model_card(**model_card_kwargs).save(save_directory / "README.md")
# push to the Hub if required
if push_to_hub:
kwargs = push_to_hub_kwargs.copy() # soft-copy to avoid mutating input
if config is not None: # kwarg for `push_to_hub`
kwargs["config"] = config
if repo_id is None:
repo_id = save_directory.name # Defaults to `save_directory` name
return self.push_to_hub(repo_id=repo_id, model_card_kwargs=model_card_kwargs, **kwargs)
return None
def _save_pretrained(self, save_directory: Path) -> None:
"""
Overwrite this method in subclass to define how to save your model.
Check out our [integration guide](../guides/integrations) for instructions.
Args:
save_directory (`str` or `Path`):
Path to directory in which the model weights and configuration will be saved.
"""
raise NotImplementedError
@classmethod
@validate_hf_hub_args
def from_pretrained(
cls: Type[T],
pretrained_model_name_or_path: Union[str, Path],
*,
force_download: bool = False,
resume_download: Optional[bool] = None,
proxies: Optional[Dict] = None,
token: Optional[Union[str, bool]] = None,
cache_dir: Optional[Union[str, Path]] = None,
local_files_only: bool = False,
revision: Optional[str] = None,
**model_kwargs,
) -> T:
"""
Download a model from the Huggingface Hub and instantiate it.
Args:
pretrained_model_name_or_path (`str`, `Path`):
- Either the `model_id` (string) of a model hosted on the Hub, e.g. `bigscience/bloom`.
- Or a path to a `directory` containing model weights saved using
[`~transformers.PreTrainedModel.save_pretrained`], e.g., `../path/to/my_model_directory/`.
revision (`str`, *optional*):
Revision of the model on the Hub. Can be a branch name, a git tag or any commit id.
Defaults to the latest commit on `main` branch.
force_download (`bool`, *optional*, defaults to `False`):
Whether to force (re-)downloading the model weights and configuration files from the Hub, overriding
the existing cache.
proxies (`Dict[str, str]`, *optional*):
A dictionary of proxy servers to use by protocol or endpoint, e.g., `{'http': 'foo.bar:3128',
'http://hostname': 'foo.bar:4012'}`. The proxies are used on every request.
token (`str` or `bool`, *optional*):
The token to use as HTTP bearer authorization for remote files. By default, it will use the token
cached when running `huggingface-cli login`.
cache_dir (`str`, `Path`, *optional*):
Path to the folder where cached files are stored.
local_files_only (`bool`, *optional*, defaults to `False`):
If `True`, avoid downloading the file and return the path to the local cached file if it exists.
model_kwargs (`Dict`, *optional*):
Additional kwargs to pass to the model during initialization.
"""
model_id = str(pretrained_model_name_or_path)
config_file: Optional[str] = None
if os.path.isdir(model_id):
if constants.CONFIG_NAME in os.listdir(model_id):
config_file = os.path.join(model_id, constants.CONFIG_NAME)
else:
logger.warning(f"{constants.CONFIG_NAME} not found in {Path(model_id).resolve()}")
else:
try:
config_file = hf_hub_download(
repo_id=model_id,
filename=constants.CONFIG_NAME,
revision=revision,
cache_dir=cache_dir,
force_download=force_download,
proxies=proxies,
resume_download=resume_download,
token=token,
local_files_only=local_files_only,
)
except HfHubHTTPError as e:
logger.info(f"{constants.CONFIG_NAME} not found on the HuggingFace Hub: {str(e)}")
# Read config
config = None
if config_file is not None:
with open(config_file, "r", encoding="utf-8") as f:
config = json.load(f)
# Decode custom types in config
for key, value in config.items():
if key in cls._hub_mixin_init_parameters:
expected_type = cls._hub_mixin_init_parameters[key].annotation
if expected_type is not inspect.Parameter.empty:
config[key] = cls._decode_arg(expected_type, value)
# Populate model_kwargs from config
for param in cls._hub_mixin_init_parameters.values():
if param.name not in model_kwargs and param.name in config:
model_kwargs[param.name] = config[param.name]
# Check if `config` argument was passed at init
if "config" in cls._hub_mixin_init_parameters and "config" not in model_kwargs:
# Decode `config` argument if it was passed
config_annotation = cls._hub_mixin_init_parameters["config"].annotation
config = cls._decode_arg(config_annotation, config)
# Forward config to model initialization
model_kwargs["config"] = config
# Inject config if `**kwargs` are expected
if is_dataclass(cls):
for key in cls.__dataclass_fields__:
if key not in model_kwargs and key in config:
model_kwargs[key] = config[key]
elif any(param.kind == inspect.Parameter.VAR_KEYWORD for param in cls._hub_mixin_init_parameters.values()):
for key, value in config.items():
if key not in model_kwargs:
model_kwargs[key] = value
# Finally, also inject if `_from_pretrained` expects it
if cls._hub_mixin_inject_config and "config" not in model_kwargs:
model_kwargs["config"] = config
instance = cls._from_pretrained(
model_id=str(model_id),
revision=revision,
cache_dir=cache_dir,
force_download=force_download,
proxies=proxies,
resume_download=resume_download,
local_files_only=local_files_only,
token=token,
**model_kwargs,
)
# Implicitly set the config as instance attribute if not already set by the class
# This way `config` will be available when calling `save_pretrained` or `push_to_hub`.
if config is not None and (getattr(instance, "_hub_mixin_config", None) in (None, {})):
instance._hub_mixin_config = config
return instance
@classmethod
def _from_pretrained(
cls: Type[T],
*,
model_id: str,
revision: Optional[str],
cache_dir: Optional[Union[str, Path]],
force_download: bool,
proxies: Optional[Dict],
resume_download: Optional[bool],
local_files_only: bool,
token: Optional[Union[str, bool]],
**model_kwargs,
) -> T:
"""Overwrite this method in subclass to define how to load your model from pretrained.
Use [`hf_hub_download`] or [`snapshot_download`] to download files from the Hub before loading them. Most
args taken as input can be directly passed to those 2 methods. If needed, you can add more arguments to this
method using "model_kwargs". For example [`PyTorchModelHubMixin._from_pretrained`] takes as input a `map_location`
parameter to set on which device the model should be loaded.
Check out our [integration guide](../guides/integrations) for more instructions.
Args:
model_id (`str`):
ID of the model to load from the Huggingface Hub (e.g. `bigscience/bloom`).
revision (`str`, *optional*):
Revision of the model on the Hub. Can be a branch name, a git tag or any commit id. Defaults to the
latest commit on `main` branch.
force_download (`bool`, *optional*, defaults to `False`):
Whether to force (re-)downloading the model weights and configuration files from the Hub, overriding
the existing cache.
proxies (`Dict[str, str]`, *optional*):
A dictionary of proxy servers to use by protocol or endpoint (e.g., `{'http': 'foo.bar:3128',
'http://hostname': 'foo.bar:4012'}`).
token (`str` or `bool`, *optional*):
The token to use as HTTP bearer authorization for remote files. By default, it will use the token
cached when running `huggingface-cli login`.
cache_dir (`str`, `Path`, *optional*):
Path to the folder where cached files are stored.
local_files_only (`bool`, *optional*, defaults to `False`):
If `True`, avoid downloading the file and return the path to the local cached file if it exists.
model_kwargs:
Additional keyword arguments passed along to the [`~ModelHubMixin._from_pretrained`] method.
"""
raise NotImplementedError
@validate_hf_hub_args
def push_to_hub(
self,
repo_id: str,
*,
config: Optional[Union[dict, DataclassInstance]] = None,
commit_message: str = "Push model using huggingface_hub.",
private: Optional[bool] = None,
token: Optional[str] = None,
branch: Optional[str] = None,
create_pr: Optional[bool] = None,
allow_patterns: Optional[Union[List[str], str]] = None,
ignore_patterns: Optional[Union[List[str], str]] = None,
delete_patterns: Optional[Union[List[str], str]] = None,
model_card_kwargs: Optional[Dict[str, Any]] = None,
) -> str:
"""
Upload model checkpoint to the Hub.
Use `allow_patterns` and `ignore_patterns` to precisely filter which files should be pushed to the hub. Use
`delete_patterns` to delete existing remote files in the same commit. See [`upload_folder`] reference for more
details.
Args:
repo_id (`str`):
ID of the repository to push to (example: `"username/my-model"`).
config (`dict` or `DataclassInstance`, *optional*):
Model configuration specified as a key/value dictionary or a dataclass instance.
commit_message (`str`, *optional*):
Message to commit while pushing.
private (`bool`, *optional*):
Whether the repository created should be private.
If `None` (default), the repo will be public unless the organization's default is private.
token (`str`, *optional*):
The token to use as HTTP bearer authorization for remote files. By default, it will use the token
cached when running `huggingface-cli login`.
branch (`str`, *optional*):
The git branch on which to push the model. This defaults to `"main"`.
create_pr (`boolean`, *optional*):
Whether or not to create a Pull Request from `branch` with that commit. Defaults to `False`.
allow_patterns (`List[str]` or `str`, *optional*):
If provided, only files matching at least one pattern are pushed.
ignore_patterns (`List[str]` or `str`, *optional*):
If provided, files matching any of the patterns are not pushed.
delete_patterns (`List[str]` or `str`, *optional*):
If provided, remote files matching any of the patterns will be deleted from the repo.
model_card_kwargs (`Dict[str, Any]`, *optional*):
Additional arguments passed to the model card template to customize the model card.
Returns:
The url of the commit of your model in the given repository.
"""
api = HfApi(token=token)
repo_id = api.create_repo(repo_id=repo_id, private=private, exist_ok=True).repo_id
# Push the files to the repo in a single commit
with SoftTemporaryDirectory() as tmp:
saved_path = Path(tmp) / repo_id
self.save_pretrained(saved_path, config=config, model_card_kwargs=model_card_kwargs)
return api.upload_folder(
repo_id=repo_id,
repo_type="model",
folder_path=saved_path,
commit_message=commit_message,
revision=branch,
create_pr=create_pr,
allow_patterns=allow_patterns,
ignore_patterns=ignore_patterns,
delete_patterns=delete_patterns,
)
def generate_model_card(self, *args, **kwargs) -> ModelCard:
card = ModelCard.from_template(
card_data=self._hub_mixin_info.model_card_data,
template_str=self._hub_mixin_info.model_card_template,
repo_url=self._hub_mixin_info.repo_url,
paper_url=self._hub_mixin_info.paper_url,
docs_url=self._hub_mixin_info.docs_url,
**kwargs,
)
return card
class PyTorchModelHubMixin(ModelHubMixin):
"""
Implementation of [`ModelHubMixin`] to provide model Hub upload/download capabilities to PyTorch models. The model
is set in evaluation mode by default using `model.eval()` (dropout modules are deactivated). To train the model,
you should first set it back in training mode with `model.train()`.
See [`ModelHubMixin`] for more details on how to use the mixin.
Example:
```python
>>> import torch
>>> import torch.nn as nn
>>> from huggingface_hub import PyTorchModelHubMixin
>>> class MyModel(
... nn.Module,
... PyTorchModelHubMixin,
... library_name="keras-nlp",
... repo_url="https://github.com/keras-team/keras-nlp",
... paper_url="https://arxiv.org/abs/2304.12244",
... docs_url="https://keras.io/keras_nlp/",
... # ^ optional metadata to generate model card
... ):
... def __init__(self, hidden_size: int = 512, vocab_size: int = 30000, output_size: int = 4):
... super().__init__()
... self.param = nn.Parameter(torch.rand(hidden_size, vocab_size))
... self.linear = nn.Linear(output_size, vocab_size)
... def forward(self, x):
... return self.linear(x + self.param)
>>> model = MyModel(hidden_size=256)
# Save model weights to local directory
>>> model.save_pretrained("my-awesome-model")
# Push model weights to the Hub
>>> model.push_to_hub("my-awesome-model")
# Download and initialize weights from the Hub
>>> model = MyModel.from_pretrained("username/my-awesome-model")
>>> model.hidden_size
256
```
"""
def __init_subclass__(cls, *args, tags: Optional[List[str]] = None, **kwargs) -> None:
tags = tags or []
tags.append("pytorch_model_hub_mixin")
kwargs["tags"] = tags
return super().__init_subclass__(*args, **kwargs)
def _save_pretrained(self, save_directory: Path) -> None:
"""Save weights from a Pytorch model to a local directory."""
model_to_save = self.module if hasattr(self, "module") else self # type: ignore
save_model_as_safetensor(model_to_save, str(save_directory / constants.SAFETENSORS_SINGLE_FILE)) # type: ignore [arg-type]
@classmethod
def _from_pretrained(
cls,
*,
model_id: str,
revision: Optional[str],
cache_dir: Optional[Union[str, Path]],
force_download: bool,
proxies: Optional[Dict],
resume_download: Optional[bool],
local_files_only: bool,
token: Union[str, bool, None],
map_location: str = "cpu",
strict: bool = False,
**model_kwargs,
):
"""Load Pytorch pretrained weights and return the loaded model."""
model = cls(**model_kwargs)
if os.path.isdir(model_id):
print("Loading weights from local directory")
model_file = os.path.join(model_id, constants.SAFETENSORS_SINGLE_FILE)
return cls._load_as_safetensor(model, model_file, map_location, strict)
else:
try:
model_file = hf_hub_download(
repo_id=model_id,
filename=constants.SAFETENSORS_SINGLE_FILE,
revision=revision,
cache_dir=cache_dir,
force_download=force_download,
proxies=proxies,
resume_download=resume_download,
token=token,
local_files_only=local_files_only,
)
return cls._load_as_safetensor(model, model_file, map_location, strict)
except EntryNotFoundError:
model_file = hf_hub_download(
repo_id=model_id,
filename=constants.PYTORCH_WEIGHTS_NAME,
revision=revision,
cache_dir=cache_dir,
force_download=force_download,
proxies=proxies,
resume_download=resume_download,
token=token,
local_files_only=local_files_only,
)
return cls._load_as_pickle(model, model_file, map_location, strict)
@classmethod
def _load_as_pickle(cls, model: T, model_file: str, map_location: str, strict: bool) -> T:
state_dict = torch.load(model_file, map_location=torch.device(map_location), weights_only=True)
model.load_state_dict(state_dict, strict=strict) # type: ignore
model.eval() # type: ignore
return model
@classmethod
def _load_as_safetensor(cls, model: T, model_file: str, map_location: str, strict: bool) -> T:
if packaging.version.parse(safetensors.__version__) < packaging.version.parse("0.4.3"): # type: ignore [attr-defined]
load_model_as_safetensor(model, model_file, strict=strict) # type: ignore [arg-type]
if map_location != "cpu":
logger.warning(
"Loading model weights on other devices than 'cpu' is not supported natively in your version of safetensors."
" This means that the model is loaded on 'cpu' first and then copied to the device."
" This leads to a slower loading time."
" Please update safetensors to version 0.4.3 or above for improved performance."
)
model.to(map_location) # type: ignore [attr-defined]
else:
safetensors.torch.load_model(model, model_file, strict=strict, device=map_location) # type: ignore [arg-type]
return model
def _load_dataclass(datacls: Type[DataclassInstance], data: dict) -> DataclassInstance:
"""Load a dataclass instance from a dictionary.
Fields not expected by the dataclass are ignored.
"""
return datacls(**{k: v for k, v in data.items() if k in datacls.__dataclass_fields__})
kronos-venv/lib/python3.12/site-packages/huggingface_hub/inference/__init__.py
View File
kronos-venv/lib/python3.12/site-packages/huggingface_hub/inference/_client.py
+3473
View File
File diff suppressed because it is too large Load Diff
kronos-venv/lib/python3.12/site-packages/huggingface_hub/inference/_common.py
+422
View File
@@ -0,0 +1,422 @@
# 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
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,
Union,
overload,
)
from requests import HTTPError
from huggingface_hub.errors import (
GenerationError,
IncompleteGenerationError,
OverloadedError,
TextGenerationError,
UnknownError,
ValidationError,
)
from ..utils import get_session, is_aiohttp_available, is_numpy_available, is_pillow_available
from ._generated.types import ChatCompletionStreamOutput, TextGenerationStreamOutput
if TYPE_CHECKING:
from aiohttp import ClientResponse, ClientSession
from PIL.Image 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__)
@dataclass
class RequestParameters:
url: str
task: str
model: Optional[str]
json: Optional[Union[str, Dict, List]]
data: Optional[ContentT]
headers: Dict[str, Any]
# Add dataclass for ModelStatus. We use this dataclass in get_model_status function.
@dataclass
class ModelStatus:
"""
This Dataclass represents the model status in the HF Inference API.
Args:
loaded (`bool`):
If the model is currently loaded into HF's Inference API. 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
## 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 bytes, 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))
def _as_dict(response: Union[bytes, Dict]) -> Dict:
return json.loads(response) if isinstance(response, bytes) else response
## PAYLOAD UTILS
## 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:
try:
output = _format_text_generation_stream_output(byte_payload, details)
except StopIteration:
break
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:
try:
output = _format_text_generation_stream_output(byte_payload, details)
except StopIteration:
break
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
if byte_payload.strip() == b"data: [DONE]":
raise StopIteration("[DONE] signal received.")
# 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(
bytes_lines: Iterable[bytes],
) -> Iterable[ChatCompletionStreamOutput]:
"""Used in `InferenceClient.chat_completion` if model is served with TGI."""
for item in bytes_lines:
try:
output = _format_chat_completion_stream_output(item)
except StopIteration:
break
if output is not None:
yield output
async def _async_stream_chat_completion_response(
bytes_lines: AsyncIterable[bytes],
) -> AsyncIterable[ChatCompletionStreamOutput]:
"""Used in `AsyncInferenceClient.chat_completion`."""
async for item in bytes_lines:
try:
output = _format_chat_completion_stream_output(item)
except StopIteration:
break
if output is not None:
yield output
def _format_chat_completion_stream_output(
byte_payload: bytes,
) -> Optional[ChatCompletionStreamOutput]:
if not byte_payload.startswith(b"data:"):
return None # empty line
if byte_payload.strip() == b"data: [DONE]":
raise StopIteration("[DONE] signal received.")
# 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
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.strip()
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. When that's the case, We remember the unsupported
# attributes for this model in the `_UNSUPPORTED_TEXT_GENERATION_KWARGS` 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.
_UNSUPPORTED_TEXT_GENERATION_KWARGS: Dict[Optional[str], List[str]] = {}
def _set_unsupported_text_generation_kwargs(model: Optional[str], unsupported_kwargs: List[str]) -> None:
_UNSUPPORTED_TEXT_GENERATION_KWARGS.setdefault(model, []).extend(unsupported_kwargs)
def _get_unsupported_text_generation_kwargs(model: Optional[str]) -> List[str]:
return _UNSUPPORTED_TEXT_GENERATION_KWARGS.get(model, [])
# 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
kronos-venv/lib/python3.12/site-packages/huggingface_hub/inference/_generated/__init__.py
View File
kronos-venv/lib/python3.12/site-packages/huggingface_hub/inference/_generated/_async_client.py
+3585
View File
File diff suppressed because it is too large Load Diff
kronos-venv/lib/python3.12/site-packages/huggingface_hub/inference/_generated/types/__init__.py
+191
View File
@@ -0,0 +1,191 @@
# This file is auto-generated by `utils/generate_inference_types.py`.
# Do not modify it manually.
#
# ruff: noqa: F401
from .audio_classification import (
AudioClassificationInput,
AudioClassificationOutputElement,
AudioClassificationOutputTransform,
AudioClassificationParameters,
)
from .audio_to_audio import AudioToAudioInput, AudioToAudioOutputElement
from .automatic_speech_recognition import (
AutomaticSpeechRecognitionEarlyStoppingEnum,
AutomaticSpeechRecognitionGenerationParameters,
AutomaticSpeechRecognitionInput,
AutomaticSpeechRecognitionOutput,
AutomaticSpeechRecognitionOutputChunk,
AutomaticSpeechRecognitionParameters,
)
from .base import BaseInferenceType
from .chat_completion import (
ChatCompletionInput,
ChatCompletionInputFunctionDefinition,
ChatCompletionInputFunctionName,
ChatCompletionInputGrammarType,
ChatCompletionInputJSONSchema,
ChatCompletionInputMessage,
ChatCompletionInputMessageChunk,
ChatCompletionInputMessageChunkType,
ChatCompletionInputResponseFormatJSONObject,
ChatCompletionInputResponseFormatJSONSchema,
ChatCompletionInputResponseFormatText,
ChatCompletionInputStreamOptions,
ChatCompletionInputTool,
ChatCompletionInputToolCall,
ChatCompletionInputToolChoiceClass,
ChatCompletionInputToolChoiceEnum,
ChatCompletionInputURL,
ChatCompletionOutput,
ChatCompletionOutputComplete,
ChatCompletionOutputFunctionDefinition,
ChatCompletionOutputLogprob,
ChatCompletionOutputLogprobs,
ChatCompletionOutputMessage,
ChatCompletionOutputToolCall,
ChatCompletionOutputTopLogprob,
ChatCompletionOutputUsage,
ChatCompletionStreamOutput,
ChatCompletionStreamOutputChoice,
ChatCompletionStreamOutputDelta,
ChatCompletionStreamOutputDeltaToolCall,
ChatCompletionStreamOutputFunction,
ChatCompletionStreamOutputLogprob,
ChatCompletionStreamOutputLogprobs,
ChatCompletionStreamOutputTopLogprob,
ChatCompletionStreamOutputUsage,
)
from .depth_estimation import DepthEstimationInput, DepthEstimationOutput
from .document_question_answering import (
DocumentQuestionAnsweringInput,
DocumentQuestionAnsweringInputData,
DocumentQuestionAnsweringOutputElement,
DocumentQuestionAnsweringParameters,
)
from .feature_extraction import FeatureExtractionInput, FeatureExtractionInputTruncationDirection
from .fill_mask import FillMaskInput, FillMaskOutputElement, FillMaskParameters
from .image_classification import (
ImageClassificationInput,
ImageClassificationOutputElement,
ImageClassificationOutputTransform,
ImageClassificationParameters,
)
from .image_segmentation import (
ImageSegmentationInput,
ImageSegmentationOutputElement,
ImageSegmentationParameters,
ImageSegmentationSubtask,
)
from .image_to_image import ImageToImageInput, ImageToImageOutput, ImageToImageParameters, ImageToImageTargetSize
from .image_to_text import (
ImageToTextEarlyStoppingEnum,
ImageToTextGenerationParameters,
ImageToTextInput,
ImageToTextOutput,
ImageToTextParameters,
)
from .object_detection import (
ObjectDetectionBoundingBox,
ObjectDetectionInput,
ObjectDetectionOutputElement,
ObjectDetectionParameters,
)
from .question_answering import (
QuestionAnsweringInput,
QuestionAnsweringInputData,
QuestionAnsweringOutputElement,
QuestionAnsweringParameters,
)
from .sentence_similarity import SentenceSimilarityInput, SentenceSimilarityInputData
from .summarization import (
SummarizationInput,
SummarizationOutput,
SummarizationParameters,
SummarizationTruncationStrategy,
)
from .table_question_answering import (
Padding,
TableQuestionAnsweringInput,
TableQuestionAnsweringInputData,
TableQuestionAnsweringOutputElement,
TableQuestionAnsweringParameters,
)
from .text2text_generation import (
Text2TextGenerationInput,
Text2TextGenerationOutput,
Text2TextGenerationParameters,
Text2TextGenerationTruncationStrategy,
)
from .text_classification import (
TextClassificationInput,
TextClassificationOutputElement,
TextClassificationOutputTransform,
TextClassificationParameters,
)
from .text_generation import (
TextGenerationInput,
TextGenerationInputGenerateParameters,
TextGenerationInputGrammarType,
TextGenerationOutput,
TextGenerationOutputBestOfSequence,
TextGenerationOutputDetails,
TextGenerationOutputFinishReason,
TextGenerationOutputPrefillToken,
TextGenerationOutputToken,
TextGenerationStreamOutput,
TextGenerationStreamOutputStreamDetails,
TextGenerationStreamOutputToken,
TypeEnum,
)
from .text_to_audio import (
TextToAudioEarlyStoppingEnum,
TextToAudioGenerationParameters,
TextToAudioInput,
TextToAudioOutput,
TextToAudioParameters,
)
from .text_to_image import TextToImageInput, TextToImageOutput, TextToImageParameters
from .text_to_speech import (
TextToSpeechEarlyStoppingEnum,
TextToSpeechGenerationParameters,
TextToSpeechInput,
TextToSpeechOutput,
TextToSpeechParameters,
)
from .text_to_video import TextToVideoInput, TextToVideoOutput, TextToVideoParameters
from .token_classification import (
TokenClassificationAggregationStrategy,
TokenClassificationInput,
TokenClassificationOutputElement,
TokenClassificationParameters,
)
from .translation import TranslationInput, TranslationOutput, TranslationParameters, TranslationTruncationStrategy
from .video_classification import (
VideoClassificationInput,
VideoClassificationOutputElement,
VideoClassificationOutputTransform,
VideoClassificationParameters,
)
from .visual_question_answering import (
VisualQuestionAnsweringInput,
VisualQuestionAnsweringInputData,
VisualQuestionAnsweringOutputElement,
VisualQuestionAnsweringParameters,
)
from .zero_shot_classification import (
ZeroShotClassificationInput,
ZeroShotClassificationOutputElement,
ZeroShotClassificationParameters,
)
from .zero_shot_image_classification import (
ZeroShotImageClassificationInput,
ZeroShotImageClassificationOutputElement,
ZeroShotImageClassificationParameters,
)
from .zero_shot_object_detection import (
ZeroShotObjectDetectionBoundingBox,
ZeroShotObjectDetectionInput,
ZeroShotObjectDetectionOutputElement,
ZeroShotObjectDetectionParameters,
)
kronos-venv/lib/python3.12/site-packages/huggingface_hub/inference/_generated/types/audio_classification.py
+43
View File
@@ -0,0 +1,43 @@
# Inference code generated from the JSON schema spec in @huggingface/tasks.
#
# See:
# - script: https://github.com/huggingface/huggingface.js/blob/main/packages/tasks/scripts/inference-codegen.ts
# - specs: https://github.com/huggingface/huggingface.js/tree/main/packages/tasks/src/tasks.
from typing import Literal, Optional
from .base import BaseInferenceType, dataclass_with_extra
AudioClassificationOutputTransform = Literal["sigmoid", "softmax", "none"]
@dataclass_with_extra
class AudioClassificationParameters(BaseInferenceType):
"""Additional inference parameters for Audio Classification"""
function_to_apply: Optional["AudioClassificationOutputTransform"] = None
"""The function to apply to the model outputs in order to retrieve the scores."""
top_k: Optional[int] = None
"""When specified, limits the output to the top K most probable classes."""
@dataclass_with_extra
class AudioClassificationInput(BaseInferenceType):
"""Inputs for Audio Classification inference"""
inputs: str
"""The input audio data as a base64-encoded string. If no `parameters` are provided, you can
also provide the audio data as a raw bytes payload.
"""
parameters: Optional[AudioClassificationParameters] = None
"""Additional inference parameters for Audio Classification"""
@dataclass_with_extra
class AudioClassificationOutputElement(BaseInferenceType):
"""Outputs for Audio Classification inference"""
label: str
"""The predicted class label."""
score: float
"""The corresponding probability."""
kronos-venv/lib/python3.12/site-packages/huggingface_hub/inference/_generated/types/audio_to_audio.py
+30
View File
@@ -0,0 +1,30 @@
# Inference code generated from the JSON schema spec in @huggingface/tasks.
#
# See:
# - script: https://github.com/huggingface/huggingface.js/blob/main/packages/tasks/scripts/inference-codegen.ts
# - specs: https://github.com/huggingface/huggingface.js/tree/main/packages/tasks/src/tasks.
from typing import Any
from .base import BaseInferenceType, dataclass_with_extra
@dataclass_with_extra
class AudioToAudioInput(BaseInferenceType):
"""Inputs for Audio to Audio inference"""
inputs: Any
"""The input audio data"""
@dataclass_with_extra
class AudioToAudioOutputElement(BaseInferenceType):
"""Outputs of inference for the Audio To Audio task
A generated audio file with its label.
"""
blob: Any
"""The generated audio file."""
content_type: str
"""The content type of audio file."""
label: str
"""The label of the audio file."""
kronos-venv/lib/python3.12/site-packages/huggingface_hub/inference/_generated/types/automatic_speech_recognition.py
+113
View File
@@ -0,0 +1,113 @@
# Inference code generated from the JSON schema spec in @huggingface/tasks.
#
# See:
# - script: https://github.com/huggingface/huggingface.js/blob/main/packages/tasks/scripts/inference-codegen.ts
# - specs: https://github.com/huggingface/huggingface.js/tree/main/packages/tasks/src/tasks.
from typing import List, Literal, Optional, Union
from .base import BaseInferenceType, dataclass_with_extra
AutomaticSpeechRecognitionEarlyStoppingEnum = Literal["never"]
@dataclass_with_extra
class AutomaticSpeechRecognitionGenerationParameters(BaseInferenceType):
"""Parametrization of the text generation process"""
do_sample: Optional[bool] = None
"""Whether to use sampling instead of greedy decoding when generating new tokens."""
early_stopping: Optional[Union[bool, "AutomaticSpeechRecognitionEarlyStoppingEnum"]] = None
"""Controls the stopping condition for beam-based methods."""
epsilon_cutoff: Optional[float] = None
"""If set to float strictly between 0 and 1, only tokens with a conditional probability
greater than epsilon_cutoff will be sampled. In the paper, suggested values range from
3e-4 to 9e-4, depending on the size of the model. See [Truncation Sampling as Language
Model Desmoothing](https://hf.co/papers/2210.15191) for more details.
"""
eta_cutoff: Optional[float] = None
"""Eta sampling is a hybrid of locally typical sampling and epsilon sampling. If set to
float strictly between 0 and 1, a token is only considered if it is greater than either
eta_cutoff or sqrt(eta_cutoff) * exp(-entropy(softmax(next_token_logits))). The latter
term is intuitively the expected next token probability, scaled by sqrt(eta_cutoff). In
the paper, suggested values range from 3e-4 to 2e-3, depending on the size of the model.
See [Truncation Sampling as Language Model Desmoothing](https://hf.co/papers/2210.15191)
for more details.
"""
max_length: Optional[int] = None
"""The maximum length (in tokens) of the generated text, including the input."""
max_new_tokens: Optional[int] = None
"""The maximum number of tokens to generate. Takes precedence over max_length."""
min_length: Optional[int] = None
"""The minimum length (in tokens) of the generated text, including the input."""
min_new_tokens: Optional[int] = None
"""The minimum number of tokens to generate. Takes precedence over min_length."""
num_beam_groups: Optional[int] = None
"""Number of groups to divide num_beams into in order to ensure diversity among different
groups of beams. See [this paper](https://hf.co/papers/1610.02424) for more details.
"""
num_beams: Optional[int] = None
"""Number of beams to use for beam search."""
penalty_alpha: Optional[float] = None
"""The value balances the model confidence and the degeneration penalty in contrastive
search decoding.
"""
temperature: Optional[float] = None
"""The value used to modulate the next token probabilities."""
top_k: Optional[int] = None
"""The number of highest probability vocabulary tokens to keep for top-k-filtering."""
top_p: Optional[float] = None
"""If set to float < 1, only the smallest set of most probable tokens with probabilities
that add up to top_p or higher are kept for generation.
"""
typical_p: Optional[float] = None
"""Local typicality measures how similar the conditional probability of predicting a target
token next is to the expected conditional probability of predicting a random token next,
given the partial text already generated. If set to float < 1, the smallest set of the
most locally typical tokens with probabilities that add up to typical_p or higher are
kept for generation. See [this paper](https://hf.co/papers/2202.00666) for more details.
"""
use_cache: Optional[bool] = None
"""Whether the model should use the past last key/values attentions to speed up decoding"""
@dataclass_with_extra
class AutomaticSpeechRecognitionParameters(BaseInferenceType):
"""Additional inference parameters for Automatic Speech Recognition"""
generation_parameters: Optional[AutomaticSpeechRecognitionGenerationParameters] = None
"""Parametrization of the text generation process"""
return_timestamps: Optional[bool] = None
"""Whether to output corresponding timestamps with the generated text"""
@dataclass_with_extra
class AutomaticSpeechRecognitionInput(BaseInferenceType):
"""Inputs for Automatic Speech Recognition inference"""
inputs: str
"""The input audio data as a base64-encoded string. If no `parameters` are provided, you can
also provide the audio data as a raw bytes payload.
"""
parameters: Optional[AutomaticSpeechRecognitionParameters] = None
"""Additional inference parameters for Automatic Speech Recognition"""
@dataclass_with_extra
class AutomaticSpeechRecognitionOutputChunk(BaseInferenceType):
text: str
"""A chunk of text identified by the model"""
timestamp: List[float]
"""The start and end timestamps corresponding with the text"""
@dataclass_with_extra
class AutomaticSpeechRecognitionOutput(BaseInferenceType):
"""Outputs of inference for the Automatic Speech Recognition task"""
text: str
"""The recognized text."""
chunks: Optional[List[AutomaticSpeechRecognitionOutputChunk]] = None
"""When returnTimestamps is enabled, chunks contains a list of audio chunks identified by
the model.
"""
kronos-venv/lib/python3.12/site-packages/huggingface_hub/inference/_generated/types/base.py
+161
View File
@@ -0,0 +1,161 @@
# Copyright 2024 The HuggingFace Team. All rights reserved.
#
# 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 a base class for all inference types."""
import inspect
import json
from dataclasses import asdict, dataclass
from typing import Any, Dict, List, Type, TypeVar, Union, get_args
T = TypeVar("T", bound="BaseInferenceType")
def _repr_with_extra(self):
fields = list(self.__dataclass_fields__.keys())
other_fields = list(k for k in self.__dict__ if k not in fields)
return f"{self.__class__.__name__}({', '.join(f'{k}={self.__dict__[k]!r}' for k in fields + other_fields)})"
def dataclass_with_extra(cls: Type[T]) -> Type[T]:
"""Decorator to add a custom __repr__ method to a dataclass, showing all fields, including extra ones.
This decorator only works with dataclasses that inherit from `BaseInferenceType`.
"""
cls = dataclass(cls)
cls.__repr__ = _repr_with_extra # type: ignore[method-assign]
return cls
@dataclass
class BaseInferenceType(dict):
"""Base class for all inference types.
Object is a dataclass and a dict for backward compatibility but plan is to remove the dict part in the future.
Handle parsing from dict, list and json strings in a permissive way to ensure future-compatibility (e.g. all fields
are made optional, and non-expected fields are added as dict attributes).
"""
@classmethod
def parse_obj_as_list(cls: Type[T], data: Union[bytes, str, List, Dict]) -> List[T]:
"""Alias to parse server response and return a single instance.
See `parse_obj` for more details.
"""
output = cls.parse_obj(data)
if not isinstance(output, list):
raise ValueError(f"Invalid input data for {cls}. Expected a list, but got {type(output)}.")
return output
@classmethod
def parse_obj_as_instance(cls: Type[T], data: Union[bytes, str, List, Dict]) -> T:
"""Alias to parse server response and return a single instance.
See `parse_obj` for more details.
"""
output = cls.parse_obj(data)
if isinstance(output, list):
raise ValueError(f"Invalid input data for {cls}. Expected a single instance, but got a list.")
return output
@classmethod
def parse_obj(cls: Type[T], data: Union[bytes, str, List, Dict]) -> Union[List[T], T]:
"""Parse server response as a dataclass or list of dataclasses.
To enable future-compatibility, we want to handle cases where the server return more fields than expected.
In such cases, we don't want to raise an error but still create the dataclass object. Remaining fields are
added as dict attributes.
"""
# Parse server response (from bytes)
if isinstance(data, bytes):
data = data.decode()
if isinstance(data, str):
data = json.loads(data)
# If a list, parse each item individually
if isinstance(data, List):
return [cls.parse_obj(d) for d in data] # type: ignore [misc]
# At this point, we expect a dict
if not isinstance(data, dict):
raise ValueError(f"Invalid data type: {type(data)}")
init_values = {}
other_values = {}
for key, value in data.items():
key = normalize_key(key)
if key in cls.__dataclass_fields__ and cls.__dataclass_fields__[key].init:
if isinstance(value, dict) or isinstance(value, list):
field_type = cls.__dataclass_fields__[key].type
# if `field_type` is a `BaseInferenceType`, parse it
if inspect.isclass(field_type) and issubclass(field_type, BaseInferenceType):
value = field_type.parse_obj(value)
# otherwise, recursively parse nested dataclasses (if possible)
# `get_args` returns handle Union and Optional for us
else:
expected_types = get_args(field_type)
for expected_type in expected_types:
if getattr(expected_type, "_name", None) == "List":
expected_type = get_args(expected_type)[
0
] # assume same type for all items in the list
if inspect.isclass(expected_type) and issubclass(expected_type, BaseInferenceType):
value = expected_type.parse_obj(value)
break
init_values[key] = value
else:
other_values[key] = value
# Make all missing fields default to None
# => ensure that dataclass initialization will never fail even if the server does not return all fields.
for key in cls.__dataclass_fields__:
if key not in init_values:
init_values[key] = None
# Initialize dataclass with expected values
item = cls(**init_values)
# Add remaining fields as dict attributes
item.update(other_values)
# Add remaining fields as extra dataclass fields.
# They won't be part of the dataclass fields but will be accessible as attributes.
# Use @dataclass_with_extra to show them in __repr__.
item.__dict__.update(other_values)
return item
def __post_init__(self):
self.update(asdict(self))
def __setitem__(self, __key: Any, __value: Any) -> None:
# Hacky way to keep dataclass values in sync when dict is updated
super().__setitem__(__key, __value)
if __key in self.__dataclass_fields__ and getattr(self, __key, None) != __value:
self.__setattr__(__key, __value)
return
def __setattr__(self, __name: str, __value: Any) -> None:
# Hacky way to keep dict values is sync when dataclass is updated
super().__setattr__(__name, __value)
if self.get(__name) != __value:
self[__name] = __value
return
def normalize_key(key: str) -> str:
# e.g "content-type" -> "content_type", "Accept" -> "accept"
return key.replace("-", "_").replace(" ", "_").lower()
kronos-venv/lib/python3.12/site-packages/huggingface_hub/inference/_generated/types/chat_completion.py
+345
View File
@@ -0,0 +1,345 @@
# Inference code generated from the JSON schema spec in @huggingface/tasks.
#
# See:
# - script: https://github.com/huggingface/huggingface.js/blob/main/packages/tasks/scripts/inference-codegen.ts
# - specs: https://github.com/huggingface/huggingface.js/tree/main/packages/tasks/src/tasks.
from typing import Any, Dict, List, Literal, Optional, Union
from .base import BaseInferenceType, dataclass_with_extra
@dataclass_with_extra
class ChatCompletionInputURL(BaseInferenceType):
url: str
ChatCompletionInputMessageChunkType = Literal["text", "image_url"]
@dataclass_with_extra
class ChatCompletionInputMessageChunk(BaseInferenceType):
type: "ChatCompletionInputMessageChunkType"
image_url: Optional[ChatCompletionInputURL] = None
text: Optional[str] = None
@dataclass_with_extra
class ChatCompletionInputFunctionDefinition(BaseInferenceType):
name: str
parameters: Any
description: Optional[str] = None
@dataclass_with_extra
class ChatCompletionInputToolCall(BaseInferenceType):
function: ChatCompletionInputFunctionDefinition
id: str
type: str
@dataclass_with_extra
class ChatCompletionInputMessage(BaseInferenceType):
role: str
content: Optional[Union[List[ChatCompletionInputMessageChunk], str]] = None
name: Optional[str] = None
tool_calls: Optional[List[ChatCompletionInputToolCall]] = None
@dataclass_with_extra
class ChatCompletionInputJSONSchema(BaseInferenceType):
name: str
"""
The name of the response format.
"""
description: Optional[str] = None
"""
A description of what the response format is for, used by the model to determine
how to respond in the format.
"""
schema: Optional[Dict[str, object]] = None
"""
The schema for the response format, described as a JSON Schema object. Learn how
to build JSON schemas [here](https://json-schema.org/).
"""
strict: Optional[bool] = None
"""
Whether to enable strict schema adherence when generating the output. If set to
true, the model will always follow the exact schema defined in the `schema`
field.
"""
@dataclass_with_extra
class ChatCompletionInputResponseFormatText(BaseInferenceType):
type: Literal["text"]
@dataclass_with_extra
class ChatCompletionInputResponseFormatJSONSchema(BaseInferenceType):
type: Literal["json_schema"]
json_schema: ChatCompletionInputJSONSchema
@dataclass_with_extra
class ChatCompletionInputResponseFormatJSONObject(BaseInferenceType):
type: Literal["json_object"]
ChatCompletionInputGrammarType = Union[
ChatCompletionInputResponseFormatText,
ChatCompletionInputResponseFormatJSONSchema,
ChatCompletionInputResponseFormatJSONObject,
]
@dataclass_with_extra
class ChatCompletionInputStreamOptions(BaseInferenceType):
include_usage: Optional[bool] = None
"""If set, an additional chunk will be streamed before the data: [DONE] message. The usage
field on this chunk shows the token usage statistics for the entire request, and the
choices field will always be an empty array. All other chunks will also include a usage
field, but with a null value.
"""
@dataclass_with_extra
class ChatCompletionInputFunctionName(BaseInferenceType):
name: str
@dataclass_with_extra
class ChatCompletionInputToolChoiceClass(BaseInferenceType):
function: ChatCompletionInputFunctionName
ChatCompletionInputToolChoiceEnum = Literal["auto", "none", "required"]
@dataclass_with_extra
class ChatCompletionInputTool(BaseInferenceType):
function: ChatCompletionInputFunctionDefinition
type: str
@dataclass_with_extra
class ChatCompletionInput(BaseInferenceType):
"""Chat Completion Input.
Auto-generated from TGI specs.
For more details, check out
https://github.com/huggingface/huggingface.js/blob/main/packages/tasks/scripts/inference-tgi-import.ts.
"""
messages: List[ChatCompletionInputMessage]
"""A list of messages comprising the conversation so far."""
frequency_penalty: Optional[float] = None
"""Number between -2.0 and 2.0. Positive values penalize new tokens based on their existing
frequency in the text so far,
decreasing the model's likelihood to repeat the same line verbatim.
"""
logit_bias: Optional[List[float]] = None
"""UNUSED
Modify the likelihood of specified tokens appearing in the completion. Accepts a JSON
object that maps tokens
(specified by their token ID in the tokenizer) to an associated bias value from -100 to
100. Mathematically,
the bias is added to the logits generated by the model prior to sampling. The exact
effect will vary per model,
but values between -1 and 1 should decrease or increase likelihood of selection; values
like -100 or 100 should
result in a ban or exclusive selection of the relevant token.
"""
logprobs: Optional[bool] = None
"""Whether to return log probabilities of the output tokens or not. If true, returns the log
probabilities of each
output token returned in the content of message.
"""
max_tokens: Optional[int] = None
"""The maximum number of tokens that can be generated in the chat completion."""
model: Optional[str] = None
"""[UNUSED] ID of the model to use. See the model endpoint compatibility table for details
on which models work with the Chat API.
"""
n: Optional[int] = None
"""UNUSED
How many chat completion choices to generate for each input message. Note that you will
be charged based on the
number of generated tokens across all of the choices. Keep n as 1 to minimize costs.
"""
presence_penalty: Optional[float] = None
"""Number between -2.0 and 2.0. Positive values penalize new tokens based on whether they
appear in the text so far,
increasing the model's likelihood to talk about new topics
"""
response_format: Optional[ChatCompletionInputGrammarType] = None
seed: Optional[int] = None
stop: Optional[List[str]] = None
"""Up to 4 sequences where the API will stop generating further tokens."""
stream: Optional[bool] = None
stream_options: Optional[ChatCompletionInputStreamOptions] = None
temperature: Optional[float] = None
"""What sampling temperature to use, between 0 and 2. Higher values like 0.8 will make the
output more random, while
lower values like 0.2 will make it more focused and deterministic.
We generally recommend altering this or `top_p` but not both.
"""
tool_choice: Optional[Union[ChatCompletionInputToolChoiceClass, "ChatCompletionInputToolChoiceEnum"]] = None
tool_prompt: Optional[str] = None
"""A prompt to be appended before the tools"""
tools: Optional[List[ChatCompletionInputTool]] = None
"""A list of tools the model may call. Currently, only functions are supported as a tool.
Use this to provide a list of
functions the model may generate JSON inputs for.
"""
top_logprobs: Optional[int] = None
"""An integer between 0 and 5 specifying the number of most likely tokens to return at each
token position, each with
an associated log probability. logprobs must be set to true if this parameter is used.
"""
top_p: Optional[float] = None
"""An alternative to sampling with temperature, called nucleus sampling, where the model
considers the results of the
tokens with top_p probability mass. So 0.1 means only the tokens comprising the top 10%
probability mass are considered.
"""
@dataclass_with_extra
class ChatCompletionOutputTopLogprob(BaseInferenceType):
logprob: float
token: str
@dataclass_with_extra
class ChatCompletionOutputLogprob(BaseInferenceType):
logprob: float
token: str
top_logprobs: List[ChatCompletionOutputTopLogprob]
@dataclass_with_extra
class ChatCompletionOutputLogprobs(BaseInferenceType):
content: List[ChatCompletionOutputLogprob]
@dataclass_with_extra
class ChatCompletionOutputFunctionDefinition(BaseInferenceType):
arguments: str
name: str
description: Optional[str] = None
@dataclass_with_extra
class ChatCompletionOutputToolCall(BaseInferenceType):
function: ChatCompletionOutputFunctionDefinition
id: str
type: str
@dataclass_with_extra
class ChatCompletionOutputMessage(BaseInferenceType):
role: str
content: Optional[str] = None
tool_call_id: Optional[str] = None
tool_calls: Optional[List[ChatCompletionOutputToolCall]] = None
@dataclass_with_extra
class ChatCompletionOutputComplete(BaseInferenceType):
finish_reason: str
index: int
message: ChatCompletionOutputMessage
logprobs: Optional[ChatCompletionOutputLogprobs] = None
@dataclass_with_extra
class ChatCompletionOutputUsage(BaseInferenceType):
completion_tokens: int
prompt_tokens: int
total_tokens: int
@dataclass_with_extra
class ChatCompletionOutput(BaseInferenceType):
"""Chat Completion Output.
Auto-generated from TGI specs.
For more details, check out
https://github.com/huggingface/huggingface.js/blob/main/packages/tasks/scripts/inference-tgi-import.ts.
"""
choices: List[ChatCompletionOutputComplete]
created: int
id: str
model: str
system_fingerprint: str
usage: ChatCompletionOutputUsage
@dataclass_with_extra
class ChatCompletionStreamOutputFunction(BaseInferenceType):
arguments: str
name: Optional[str] = None
@dataclass_with_extra
class ChatCompletionStreamOutputDeltaToolCall(BaseInferenceType):
function: ChatCompletionStreamOutputFunction
id: str
index: int
type: str
@dataclass_with_extra
class ChatCompletionStreamOutputDelta(BaseInferenceType):
role: str
content: Optional[str] = None
tool_call_id: Optional[str] = None
tool_calls: Optional[List[ChatCompletionStreamOutputDeltaToolCall]] = None
@dataclass_with_extra
class ChatCompletionStreamOutputTopLogprob(BaseInferenceType):
logprob: float
token: str
@dataclass_with_extra
class ChatCompletionStreamOutputLogprob(BaseInferenceType):
logprob: float
token: str
top_logprobs: List[ChatCompletionStreamOutputTopLogprob]
@dataclass_with_extra
class ChatCompletionStreamOutputLogprobs(BaseInferenceType):
content: List[ChatCompletionStreamOutputLogprob]
@dataclass_with_extra
class ChatCompletionStreamOutputChoice(BaseInferenceType):
delta: ChatCompletionStreamOutputDelta
index: int
finish_reason: Optional[str] = None
logprobs: Optional[ChatCompletionStreamOutputLogprobs] = None
@dataclass_with_extra
class ChatCompletionStreamOutputUsage(BaseInferenceType):
completion_tokens: int
prompt_tokens: int
total_tokens: int
@dataclass_with_extra
class ChatCompletionStreamOutput(BaseInferenceType):
"""Chat Completion Stream Output.
Auto-generated from TGI specs.
For more details, check out
https://github.com/huggingface/huggingface.js/blob/main/packages/tasks/scripts/inference-tgi-import.ts.
"""
choices: List[ChatCompletionStreamOutputChoice]
created: int
id: str
model: str
system_fingerprint: str
usage: Optional[ChatCompletionStreamOutputUsage] = None
kronos-venv/lib/python3.12/site-packages/huggingface_hub/inference/_generated/types/depth_estimation.py
+28
View File
@@ -0,0 +1,28 @@
# Inference code generated from the JSON schema spec in @huggingface/tasks.
#
# See:
# - script: https://github.com/huggingface/huggingface.js/blob/main/packages/tasks/scripts/inference-codegen.ts
# - specs: https://github.com/huggingface/huggingface.js/tree/main/packages/tasks/src/tasks.
from typing import Any, Dict, Optional
from .base import BaseInferenceType, dataclass_with_extra
@dataclass_with_extra
class DepthEstimationInput(BaseInferenceType):
"""Inputs for Depth Estimation inference"""
inputs: Any
"""The input image data"""
parameters: Optional[Dict[str, Any]] = None
"""Additional inference parameters for Depth Estimation"""
@dataclass_with_extra
class DepthEstimationOutput(BaseInferenceType):
"""Outputs of inference for the Depth Estimation task"""
depth: Any
"""The predicted depth as an image"""
predicted_depth: Any
"""The predicted depth as a tensor"""
kronos-venv/lib/python3.12/site-packages/huggingface_hub/inference/_generated/types/document_question_answering.py
+80
View File
@@ -0,0 +1,80 @@
# Inference code generated from the JSON schema spec in @huggingface/tasks.
#
# See:
# - script: https://github.com/huggingface/huggingface.js/blob/main/packages/tasks/scripts/inference-codegen.ts
# - specs: https://github.com/huggingface/huggingface.js/tree/main/packages/tasks/src/tasks.
from typing import Any, List, Optional, Union
from .base import BaseInferenceType, dataclass_with_extra
@dataclass_with_extra
class DocumentQuestionAnsweringInputData(BaseInferenceType):
"""One (document, question) pair to answer"""
image: Any
"""The image on which the question is asked"""
question: str
"""A question to ask of the document"""
@dataclass_with_extra
class DocumentQuestionAnsweringParameters(BaseInferenceType):
"""Additional inference parameters for Document Question Answering"""
doc_stride: Optional[int] = None
"""If the words in the document are too long to fit with the question for the model, it will
be split in several chunks with some overlap. This argument controls the size of that
overlap.
"""
handle_impossible_answer: Optional[bool] = None
"""Whether to accept impossible as an answer"""
lang: Optional[str] = None
"""Language to use while running OCR. Defaults to english."""
max_answer_len: Optional[int] = None
"""The maximum length of predicted answers (e.g., only answers with a shorter length are
considered).
"""
max_question_len: Optional[int] = None
"""The maximum length of the question after tokenization. It will be truncated if needed."""
max_seq_len: Optional[int] = None
"""The maximum length of the total sentence (context + question) in tokens of each chunk
passed to the model. The context will be split in several chunks (using doc_stride as
overlap) if needed.
"""
top_k: Optional[int] = None
"""The number of answers to return (will be chosen by order of likelihood). Can return less
than top_k answers if there are not enough options available within the context.
"""
word_boxes: Optional[List[Union[List[float], str]]] = None
"""A list of words and bounding boxes (normalized 0->1000). If provided, the inference will
skip the OCR step and use the provided bounding boxes instead.
"""
@dataclass_with_extra
class DocumentQuestionAnsweringInput(BaseInferenceType):
"""Inputs for Document Question Answering inference"""
inputs: DocumentQuestionAnsweringInputData
"""One (document, question) pair to answer"""
parameters: Optional[DocumentQuestionAnsweringParameters] = None
"""Additional inference parameters for Document Question Answering"""
@dataclass_with_extra
class DocumentQuestionAnsweringOutputElement(BaseInferenceType):
"""Outputs of inference for the Document Question Answering task"""
answer: str
"""The answer to the question."""
end: int
"""The end word index of the answer (in the OCR’d version of the input or provided word
boxes).
"""
score: float
"""The probability associated to the answer."""
start: int
"""The start word index of the answer (in the OCR’d version of the input or provided word
boxes).
"""
kronos-venv/lib/python3.12/site-packages/huggingface_hub/inference/_generated/types/feature_extraction.py
+36
View File
@@ -0,0 +1,36 @@
# Inference code generated from the JSON schema spec in @huggingface/tasks.
#
# See:
# - script: https://github.com/huggingface/huggingface.js/blob/main/packages/tasks/scripts/inference-codegen.ts
# - specs: https://github.com/huggingface/huggingface.js/tree/main/packages/tasks/src/tasks.
from typing import List, Literal, Optional, Union
from .base import BaseInferenceType, dataclass_with_extra
FeatureExtractionInputTruncationDirection = Literal["Left", "Right"]
@dataclass_with_extra
class FeatureExtractionInput(BaseInferenceType):
"""Feature Extraction Input.
Auto-generated from TEI specs.
For more details, check out
https://github.com/huggingface/huggingface.js/blob/main/packages/tasks/scripts/inference-tei-import.ts.
"""
inputs: Union[List[str], str]
"""The text or list of texts to embed."""
normalize: Optional[bool] = None
prompt_name: Optional[str] = None
"""The name of the prompt that should be used by for encoding. If not set, no prompt
will be applied.
Must be a key in the `sentence-transformers` configuration `prompts` dictionary.
For example if ``prompt_name`` is "query" and the ``prompts`` is {"query": "query: ",
...},
then the sentence "What is the capital of France?" will be encoded as
"query: What is the capital of France?" because the prompt text will be prepended before
any text to encode.
"""
truncate: Optional[bool] = None
truncation_direction: Optional["FeatureExtractionInputTruncationDirection"] = None
kronos-venv/lib/python3.12/site-packages/huggingface_hub/inference/_generated/types/fill_mask.py
+47
View File
@@ -0,0 +1,47 @@
# Inference code generated from the JSON schema spec in @huggingface/tasks.
#
# See:
# - script: https://github.com/huggingface/huggingface.js/blob/main/packages/tasks/scripts/inference-codegen.ts
# - specs: https://github.com/huggingface/huggingface.js/tree/main/packages/tasks/src/tasks.
from typing import Any, List, Optional
from .base import BaseInferenceType, dataclass_with_extra
@dataclass_with_extra
class FillMaskParameters(BaseInferenceType):
"""Additional inference parameters for Fill Mask"""
targets: Optional[List[str]] = None
"""When passed, the model will limit the scores to the passed targets instead of looking up
in the whole vocabulary. If the provided targets are not in the model vocab, they will be
tokenized and the first resulting token will be used (with a warning, and that might be
slower).
"""
top_k: Optional[int] = None
"""When passed, overrides the number of predictions to return."""
@dataclass_with_extra
class FillMaskInput(BaseInferenceType):
"""Inputs for Fill Mask inference"""
inputs: str
"""The text with masked tokens"""
parameters: Optional[FillMaskParameters] = None
"""Additional inference parameters for Fill Mask"""
@dataclass_with_extra
class FillMaskOutputElement(BaseInferenceType):
"""Outputs of inference for the Fill Mask task"""
score: float
"""The corresponding probability"""
sequence: str
"""The corresponding input with the mask token prediction."""
token: int
"""The predicted token id (to replace the masked one)."""
token_str: Any
fill_mask_output_token_str: Optional[str] = None
"""The predicted token (to replace the masked one)."""
kronos-venv/lib/python3.12/site-packages/huggingface_hub/inference/_generated/types/image_classification.py
+43
View File
@@ -0,0 +1,43 @@
# Inference code generated from the JSON schema spec in @huggingface/tasks.
#
# See:
# - script: https://github.com/huggingface/huggingface.js/blob/main/packages/tasks/scripts/inference-codegen.ts
# - specs: https://github.com/huggingface/huggingface.js/tree/main/packages/tasks/src/tasks.
from typing import Literal, Optional
from .base import BaseInferenceType, dataclass_with_extra
ImageClassificationOutputTransform = Literal["sigmoid", "softmax", "none"]
@dataclass_with_extra
class ImageClassificationParameters(BaseInferenceType):
"""Additional inference parameters for Image Classification"""
function_to_apply: Optional["ImageClassificationOutputTransform"] = None
"""The function to apply to the model outputs in order to retrieve the scores."""
top_k: Optional[int] = None
"""When specified, limits the output to the top K most probable classes."""
@dataclass_with_extra
class ImageClassificationInput(BaseInferenceType):
"""Inputs for Image Classification inference"""
inputs: str
"""The input image data as a base64-encoded string. If no `parameters` are provided, you can
also provide the image data as a raw bytes payload.
"""
parameters: Optional[ImageClassificationParameters] = None
"""Additional inference parameters for Image Classification"""
@dataclass_with_extra
class ImageClassificationOutputElement(BaseInferenceType):
"""Outputs of inference for the Image Classification task"""
label: str
"""The predicted class label."""
score: float
"""The corresponding probability."""
kronos-venv/lib/python3.12/site-packages/huggingface_hub/inference/_generated/types/image_segmentation.py
+51
View File
@@ -0,0 +1,51 @@
# Inference code generated from the JSON schema spec in @huggingface/tasks.
#
# See:
# - script: https://github.com/huggingface/huggingface.js/blob/main/packages/tasks/scripts/inference-codegen.ts
# - specs: https://github.com/huggingface/huggingface.js/tree/main/packages/tasks/src/tasks.
from typing import Literal, Optional
from .base import BaseInferenceType, dataclass_with_extra
ImageSegmentationSubtask = Literal["instance", "panoptic", "semantic"]
@dataclass_with_extra
class ImageSegmentationParameters(BaseInferenceType):
"""Additional inference parameters for Image Segmentation"""
mask_threshold: Optional[float] = None
"""Threshold to use when turning the predicted masks into binary values."""
overlap_mask_area_threshold: Optional[float] = None
"""Mask overlap threshold to eliminate small, disconnected segments."""
subtask: Optional["ImageSegmentationSubtask"] = None
"""Segmentation task to be performed, depending on model capabilities."""
threshold: Optional[float] = None
"""Probability threshold to filter out predicted masks."""
@dataclass_with_extra
class ImageSegmentationInput(BaseInferenceType):
"""Inputs for Image Segmentation inference"""
inputs: str
"""The input image data as a base64-encoded string. If no `parameters` are provided, you can
also provide the image data as a raw bytes payload.
"""
parameters: Optional[ImageSegmentationParameters] = None
"""Additional inference parameters for Image Segmentation"""
@dataclass_with_extra
class ImageSegmentationOutputElement(BaseInferenceType):
"""Outputs of inference for the Image Segmentation task
A predicted mask / segment
"""
label: str
"""The label of the predicted segment."""
mask: str
"""The corresponding mask as a black-and-white image (base64-encoded)."""
score: Optional[float] = None
"""The score or confidence degree the model has."""
kronos-venv/lib/python3.12/site-packages/huggingface_hub/inference/_generated/types/image_to_image.py
+56
View File
@@ -0,0 +1,56 @@
# Inference code generated from the JSON schema spec in @huggingface/tasks.
#
# See:
# - script: https://github.com/huggingface/huggingface.js/blob/main/packages/tasks/scripts/inference-codegen.ts
# - specs: https://github.com/huggingface/huggingface.js/tree/main/packages/tasks/src/tasks.
from typing import Any, Optional
from .base import BaseInferenceType, dataclass_with_extra
@dataclass_with_extra
class ImageToImageTargetSize(BaseInferenceType):
"""The size in pixel of the output image."""
height: int
width: int
@dataclass_with_extra
class ImageToImageParameters(BaseInferenceType):
"""Additional inference parameters for Image To Image"""
guidance_scale: Optional[float] = None
"""For diffusion models. A higher guidance scale value encourages the model to generate
images closely linked to the text prompt at the expense of lower image quality.
"""
negative_prompt: Optional[str] = None
"""One prompt to guide what NOT to include in image generation."""
num_inference_steps: Optional[int] = None
"""For diffusion models. The number of denoising steps. More denoising steps usually lead to
a higher quality image at the expense of slower inference.
"""
prompt: Optional[str] = None
"""The text prompt to guide the image generation."""
target_size: Optional[ImageToImageTargetSize] = None
"""The size in pixel of the output image."""
@dataclass_with_extra
class ImageToImageInput(BaseInferenceType):
"""Inputs for Image To Image inference"""
inputs: str
"""The input image data as a base64-encoded string. If no `parameters` are provided, you can
also provide the image data as a raw bytes payload.
"""
parameters: Optional[ImageToImageParameters] = None
"""Additional inference parameters for Image To Image"""
@dataclass_with_extra
class ImageToImageOutput(BaseInferenceType):
"""Outputs of inference for the Image To Image task"""
image: Any
"""The output image returned as raw bytes in the payload."""
kronos-venv/lib/python3.12/site-packages/huggingface_hub/inference/_generated/types/image_to_text.py
+100
View File
@@ -0,0 +1,100 @@
# Inference code generated from the JSON schema spec in @huggingface/tasks.
#
# See:
# - script: https://github.com/huggingface/huggingface.js/blob/main/packages/tasks/scripts/inference-codegen.ts
# - specs: https://github.com/huggingface/huggingface.js/tree/main/packages/tasks/src/tasks.
from typing import Any, Literal, Optional, Union
from .base import BaseInferenceType, dataclass_with_extra
ImageToTextEarlyStoppingEnum = Literal["never"]
@dataclass_with_extra
class ImageToTextGenerationParameters(BaseInferenceType):
"""Parametrization of the text generation process"""
do_sample: Optional[bool] = None
"""Whether to use sampling instead of greedy decoding when generating new tokens."""
early_stopping: Optional[Union[bool, "ImageToTextEarlyStoppingEnum"]] = None
"""Controls the stopping condition for beam-based methods."""
epsilon_cutoff: Optional[float] = None
"""If set to float strictly between 0 and 1, only tokens with a conditional probability
greater than epsilon_cutoff will be sampled. In the paper, suggested values range from
3e-4 to 9e-4, depending on the size of the model. See [Truncation Sampling as Language
Model Desmoothing](https://hf.co/papers/2210.15191) for more details.
"""
eta_cutoff: Optional[float] = None
"""Eta sampling is a hybrid of locally typical sampling and epsilon sampling. If set to
float strictly between 0 and 1, a token is only considered if it is greater than either
eta_cutoff or sqrt(eta_cutoff) * exp(-entropy(softmax(next_token_logits))). The latter
term is intuitively the expected next token probability, scaled by sqrt(eta_cutoff). In
the paper, suggested values range from 3e-4 to 2e-3, depending on the size of the model.
See [Truncation Sampling as Language Model Desmoothing](https://hf.co/papers/2210.15191)
for more details.
"""
max_length: Optional[int] = None
"""The maximum length (in tokens) of the generated text, including the input."""
max_new_tokens: Optional[int] = None
"""The maximum number of tokens to generate. Takes precedence over max_length."""
min_length: Optional[int] = None
"""The minimum length (in tokens) of the generated text, including the input."""
min_new_tokens: Optional[int] = None
"""The minimum number of tokens to generate. Takes precedence over min_length."""
num_beam_groups: Optional[int] = None
"""Number of groups to divide num_beams into in order to ensure diversity among different
groups of beams. See [this paper](https://hf.co/papers/1610.02424) for more details.
"""
num_beams: Optional[int] = None
"""Number of beams to use for beam search."""
penalty_alpha: Optional[float] = None
"""The value balances the model confidence and the degeneration penalty in contrastive
search decoding.
"""
temperature: Optional[float] = None
"""The value used to modulate the next token probabilities."""
top_k: Optional[int] = None
"""The number of highest probability vocabulary tokens to keep for top-k-filtering."""
top_p: Optional[float] = None
"""If set to float < 1, only the smallest set of most probable tokens with probabilities
that add up to top_p or higher are kept for generation.
"""
typical_p: Optional[float] = None
"""Local typicality measures how similar the conditional probability of predicting a target
token next is to the expected conditional probability of predicting a random token next,
given the partial text already generated. If set to float < 1, the smallest set of the
most locally typical tokens with probabilities that add up to typical_p or higher are
kept for generation. See [this paper](https://hf.co/papers/2202.00666) for more details.
"""
use_cache: Optional[bool] = None
"""Whether the model should use the past last key/values attentions to speed up decoding"""
@dataclass_with_extra
class ImageToTextParameters(BaseInferenceType):
"""Additional inference parameters for Image To Text"""
generation_parameters: Optional[ImageToTextGenerationParameters] = None
"""Parametrization of the text generation process"""
max_new_tokens: Optional[int] = None
"""The amount of maximum tokens to generate."""
@dataclass_with_extra
class ImageToTextInput(BaseInferenceType):
"""Inputs for Image To Text inference"""
inputs: Any
"""The input image data"""
parameters: Optional[ImageToTextParameters] = None
"""Additional inference parameters for Image To Text"""
@dataclass_with_extra
class ImageToTextOutput(BaseInferenceType):
"""Outputs of inference for the Image To Text task"""
generated_text: Any
image_to_text_output_generated_text: Optional[str] = None
"""The generated text."""
kronos-venv/lib/python3.12/site-packages/huggingface_hub/inference/_generated/types/object_detection.py
+58
View File
@@ -0,0 +1,58 @@
# Inference code generated from the JSON schema spec in @huggingface/tasks.
#
# See:
# - script: https://github.com/huggingface/huggingface.js/blob/main/packages/tasks/scripts/inference-codegen.ts
# - specs: https://github.com/huggingface/huggingface.js/tree/main/packages/tasks/src/tasks.
from typing import Optional
from .base import BaseInferenceType, dataclass_with_extra
@dataclass_with_extra
class ObjectDetectionParameters(BaseInferenceType):
"""Additional inference parameters for Object Detection"""
threshold: Optional[float] = None
"""The probability necessary to make a prediction."""
@dataclass_with_extra
class ObjectDetectionInput(BaseInferenceType):
"""Inputs for Object Detection inference"""
inputs: str
"""The input image data as a base64-encoded string. If no `parameters` are provided, you can
also provide the image data as a raw bytes payload.
"""
parameters: Optional[ObjectDetectionParameters] = None
"""Additional inference parameters for Object Detection"""
@dataclass_with_extra
class ObjectDetectionBoundingBox(BaseInferenceType):
"""The predicted bounding box. Coordinates are relative to the top left corner of the input
image.
"""
xmax: int
"""The x-coordinate of the bottom-right corner of the bounding box."""
xmin: int
"""The x-coordinate of the top-left corner of the bounding box."""
ymax: int
"""The y-coordinate of the bottom-right corner of the bounding box."""
ymin: int
"""The y-coordinate of the top-left corner of the bounding box."""
@dataclass_with_extra
class ObjectDetectionOutputElement(BaseInferenceType):
"""Outputs of inference for the Object Detection task"""
box: ObjectDetectionBoundingBox
"""The predicted bounding box. Coordinates are relative to the top left corner of the input
image.
"""
label: str
"""The predicted label for the bounding box."""
score: float
"""The associated score / probability."""
kronos-venv/lib/python3.12/site-packages/huggingface_hub/inference/_generated/types/question_answering.py
+74
View File
@@ -0,0 +1,74 @@
# Inference code generated from the JSON schema spec in @huggingface/tasks.
#
# See:
# - script: https://github.com/huggingface/huggingface.js/blob/main/packages/tasks/scripts/inference-codegen.ts
# - specs: https://github.com/huggingface/huggingface.js/tree/main/packages/tasks/src/tasks.
from typing import Optional
from .base import BaseInferenceType, dataclass_with_extra
@dataclass_with_extra
class QuestionAnsweringInputData(BaseInferenceType):
"""One (context, question) pair to answer"""
context: str
"""The context to be used for answering the question"""
question: str
"""The question to be answered"""
@dataclass_with_extra
class QuestionAnsweringParameters(BaseInferenceType):
"""Additional inference parameters for Question Answering"""
align_to_words: Optional[bool] = None
"""Attempts to align the answer to real words. Improves quality on space separated
languages. Might hurt on non-space-separated languages (like Japanese or Chinese)
"""
doc_stride: Optional[int] = None
"""If the context is too long to fit with the question for the model, it will be split in
several chunks with some overlap. This argument controls the size of that overlap.
"""
handle_impossible_answer: Optional[bool] = None
"""Whether to accept impossible as an answer."""
max_answer_len: Optional[int] = None
"""The maximum length of predicted answers (e.g., only answers with a shorter length are
considered).
"""
max_question_len: Optional[int] = None
"""The maximum length of the question after tokenization. It will be truncated if needed."""
max_seq_len: Optional[int] = None
"""The maximum length of the total sentence (context + question) in tokens of each chunk
passed to the model. The context will be split in several chunks (using docStride as
overlap) if needed.
"""
top_k: Optional[int] = None
"""The number of answers to return (will be chosen by order of likelihood). Note that we
return less than topk answers if there are not enough options available within the
context.
"""
@dataclass_with_extra
class QuestionAnsweringInput(BaseInferenceType):
"""Inputs for Question Answering inference"""
inputs: QuestionAnsweringInputData
"""One (context, question) pair to answer"""
parameters: Optional[QuestionAnsweringParameters] = None
"""Additional inference parameters for Question Answering"""
@dataclass_with_extra
class QuestionAnsweringOutputElement(BaseInferenceType):
"""Outputs of inference for the Question Answering task"""
answer: str
"""The answer to the question."""
end: int
"""The character position in the input where the answer ends."""
score: float
"""The probability associated to the answer."""
start: int
"""The character position in the input where the answer begins."""
kronos-venv/lib/python3.12/site-packages/huggingface_hub/inference/_generated/types/sentence_similarity.py
+27
View File
@@ -0,0 +1,27 @@
# Inference code generated from the JSON schema spec in @huggingface/tasks.
#
# See:
# - script: https://github.com/huggingface/huggingface.js/blob/main/packages/tasks/scripts/inference-codegen.ts
# - specs: https://github.com/huggingface/huggingface.js/tree/main/packages/tasks/src/tasks.
from typing import Any, Dict, List, Optional
from .base import BaseInferenceType, dataclass_with_extra
@dataclass_with_extra
class SentenceSimilarityInputData(BaseInferenceType):
sentences: List[str]
"""A list of strings which will be compared against the source_sentence."""
source_sentence: str
"""The string that you wish to compare the other strings with. This can be a phrase,
sentence, or longer passage, depending on the model being used.
"""
@dataclass_with_extra
class SentenceSimilarityInput(BaseInferenceType):
"""Inputs for Sentence similarity inference"""
inputs: SentenceSimilarityInputData
parameters: Optional[Dict[str, Any]] = None
"""Additional inference parameters for Sentence Similarity"""
kronos-venv/lib/python3.12/site-packages/huggingface_hub/inference/_generated/types/summarization.py
+41
View File
@@ -0,0 +1,41 @@
# Inference code generated from the JSON schema spec in @huggingface/tasks.
#
# See:
# - script: https://github.com/huggingface/huggingface.js/blob/main/packages/tasks/scripts/inference-codegen.ts
# - specs: https://github.com/huggingface/huggingface.js/tree/main/packages/tasks/src/tasks.
from typing import Any, Dict, Literal, Optional
from .base import BaseInferenceType, dataclass_with_extra
SummarizationTruncationStrategy = Literal["do_not_truncate", "longest_first", "only_first", "only_second"]
@dataclass_with_extra
class SummarizationParameters(BaseInferenceType):
"""Additional inference parameters for summarization."""
clean_up_tokenization_spaces: Optional[bool] = None
"""Whether to clean up the potential extra spaces in the text output."""
generate_parameters: Optional[Dict[str, Any]] = None
"""Additional parametrization of the text generation algorithm."""
truncation: Optional["SummarizationTruncationStrategy"] = None
"""The truncation strategy to use."""
@dataclass_with_extra
class SummarizationInput(BaseInferenceType):
"""Inputs for Summarization inference"""
inputs: str
"""The input text to summarize."""
parameters: Optional[SummarizationParameters] = None
"""Additional inference parameters for summarization."""
@dataclass_with_extra
class SummarizationOutput(BaseInferenceType):
"""Outputs of inference for the Summarization task"""
summary_text: str
"""The summarized text."""
kronos-venv/lib/python3.12/site-packages/huggingface_hub/inference/_generated/types/table_question_answering.py
+62
View File
@@ -0,0 +1,62 @@
# Inference code generated from the JSON schema spec in @huggingface/tasks.
#
# See:
# - script: https://github.com/huggingface/huggingface.js/blob/main/packages/tasks/scripts/inference-codegen.ts
# - specs: https://github.com/huggingface/huggingface.js/tree/main/packages/tasks/src/tasks.
from typing import Dict, List, Literal, Optional
from .base import BaseInferenceType, dataclass_with_extra
@dataclass_with_extra
class TableQuestionAnsweringInputData(BaseInferenceType):
"""One (table, question) pair to answer"""
question: str
"""The question to be answered about the table"""
table: Dict[str, List[str]]
"""The table to serve as context for the questions"""
Padding = Literal["do_not_pad", "longest", "max_length"]
@dataclass_with_extra
class TableQuestionAnsweringParameters(BaseInferenceType):
"""Additional inference parameters for Table Question Answering"""
padding: Optional["Padding"] = None
"""Activates and controls padding."""
sequential: Optional[bool] = None
"""Whether to do inference sequentially or as a batch. Batching is faster, but models like
SQA require the inference to be done sequentially to extract relations within sequences,
given their conversational nature.
"""
truncation: Optional[bool] = None
"""Activates and controls truncation."""
@dataclass_with_extra
class TableQuestionAnsweringInput(BaseInferenceType):
"""Inputs for Table Question Answering inference"""
inputs: TableQuestionAnsweringInputData
"""One (table, question) pair to answer"""
parameters: Optional[TableQuestionAnsweringParameters] = None
"""Additional inference parameters for Table Question Answering"""
@dataclass_with_extra
class TableQuestionAnsweringOutputElement(BaseInferenceType):
"""Outputs of inference for the Table Question Answering task"""
answer: str
"""The answer of the question given the table. If there is an aggregator, the answer will be
preceded by `AGGREGATOR >`.
"""
cells: List[str]
"""List of strings made up of the answer cell values."""
coordinates: List[List[int]]
"""Coordinates of the cells of the answers."""
aggregator: Optional[str] = None
"""If the model has an aggregator, this returns the aggregator."""
kronos-venv/lib/python3.12/site-packages/huggingface_hub/inference/_generated/types/text2text_generation.py
+42
View File
@@ -0,0 +1,42 @@
# Inference code generated from the JSON schema spec in @huggingface/tasks.
#
# See:
# - script: https://github.com/huggingface/huggingface.js/blob/main/packages/tasks/scripts/inference-codegen.ts
# - specs: https://github.com/huggingface/huggingface.js/tree/main/packages/tasks/src/tasks.
from typing import Any, Dict, Literal, Optional
from .base import BaseInferenceType, dataclass_with_extra
Text2TextGenerationTruncationStrategy = Literal["do_not_truncate", "longest_first", "only_first", "only_second"]
@dataclass_with_extra
class Text2TextGenerationParameters(BaseInferenceType):
"""Additional inference parameters for Text2text Generation"""
clean_up_tokenization_spaces: Optional[bool] = None
"""Whether to clean up the potential extra spaces in the text output."""
generate_parameters: Optional[Dict[str, Any]] = None
"""Additional parametrization of the text generation algorithm"""
truncation: Optional["Text2TextGenerationTruncationStrategy"] = None
"""The truncation strategy to use"""
@dataclass_with_extra
class Text2TextGenerationInput(BaseInferenceType):
"""Inputs for Text2text Generation inference"""
inputs: str
"""The input text data"""
parameters: Optional[Text2TextGenerationParameters] = None
"""Additional inference parameters for Text2text Generation"""
@dataclass_with_extra
class Text2TextGenerationOutput(BaseInferenceType):
"""Outputs of inference for the Text2text Generation task"""
generated_text: Any
text2_text_generation_output_generated_text: Optional[str] = None
"""The generated text."""
kronos-venv/lib/python3.12/site-packages/huggingface_hub/inference/_generated/types/text_classification.py
+41
View File
@@ -0,0 +1,41 @@
# Inference code generated from the JSON schema spec in @huggingface/tasks.
#
# See:
# - script: https://github.com/huggingface/huggingface.js/blob/main/packages/tasks/scripts/inference-codegen.ts
# - specs: https://github.com/huggingface/huggingface.js/tree/main/packages/tasks/src/tasks.
from typing import Literal, Optional
from .base import BaseInferenceType, dataclass_with_extra
TextClassificationOutputTransform = Literal["sigmoid", "softmax", "none"]
@dataclass_with_extra
class TextClassificationParameters(BaseInferenceType):
"""Additional inference parameters for Text Classification"""
function_to_apply: Optional["TextClassificationOutputTransform"] = None
"""The function to apply to the model outputs in order to retrieve the scores."""
top_k: Optional[int] = None
"""When specified, limits the output to the top K most probable classes."""
@dataclass_with_extra
class TextClassificationInput(BaseInferenceType):
"""Inputs for Text Classification inference"""
inputs: str
"""The text to classify"""
parameters: Optional[TextClassificationParameters] = None
"""Additional inference parameters for Text Classification"""
@dataclass_with_extra
class TextClassificationOutputElement(BaseInferenceType):
"""Outputs of inference for the Text Classification task"""
label: str
"""The predicted class label."""
score: float
"""The corresponding probability."""
kronos-venv/lib/python3.12/site-packages/huggingface_hub/inference/_generated/types/text_generation.py
+168
View File
@@ -0,0 +1,168 @@
# Inference code generated from the JSON schema spec in @huggingface/tasks.
#
# See:
# - script: https://github.com/huggingface/huggingface.js/blob/main/packages/tasks/scripts/inference-codegen.ts
# - specs: https://github.com/huggingface/huggingface.js/tree/main/packages/tasks/src/tasks.
from typing import Any, List, Literal, Optional
from .base import BaseInferenceType, dataclass_with_extra
TypeEnum = Literal["json", "regex", "json_schema"]
@dataclass_with_extra
class TextGenerationInputGrammarType(BaseInferenceType):
type: "TypeEnum"
value: Any
"""A string that represents a [JSON Schema](https://json-schema.org/).
JSON Schema is a declarative language that allows to annotate JSON documents
with types and descriptions.
"""
@dataclass_with_extra
class TextGenerationInputGenerateParameters(BaseInferenceType):
adapter_id: Optional[str] = None
"""Lora adapter id"""
best_of: Optional[int] = None
"""Generate best_of sequences and return the one if the highest token logprobs."""
decoder_input_details: Optional[bool] = None
"""Whether to return decoder input token logprobs and ids."""
details: Optional[bool] = None
"""Whether to return generation details."""
do_sample: Optional[bool] = None
"""Activate logits sampling."""
frequency_penalty: Optional[float] = None
"""The parameter for frequency penalty. 1.0 means no penalty
Penalize new tokens based on their existing frequency in the text so far,
decreasing the model's likelihood to repeat the same line verbatim.
"""
grammar: Optional[TextGenerationInputGrammarType] = None
max_new_tokens: Optional[int] = None
"""Maximum number of tokens to generate."""
repetition_penalty: Optional[float] = None
"""The parameter for repetition penalty. 1.0 means no penalty.
See [this paper](https://arxiv.org/pdf/1909.05858.pdf) for more details.
"""
return_full_text: Optional[bool] = None
"""Whether to prepend the prompt to the generated text"""
seed: Optional[int] = None
"""Random sampling seed."""
stop: Optional[List[str]] = None
"""Stop generating tokens if a member of `stop` is generated."""
temperature: Optional[float] = None
"""The value used to module the logits distribution."""
top_k: Optional[int] = None
"""The number of highest probability vocabulary tokens to keep for top-k-filtering."""
top_n_tokens: Optional[int] = None
"""The number of highest probability vocabulary tokens to keep for top-n-filtering."""
top_p: Optional[float] = None
"""Top-p value for nucleus sampling."""
truncate: Optional[int] = None
"""Truncate inputs tokens to the given size."""
typical_p: Optional[float] = None
"""Typical Decoding mass
See [Typical Decoding for Natural Language Generation](https://arxiv.org/abs/2202.00666)
for more information.
"""
watermark: Optional[bool] = None
"""Watermarking with [A Watermark for Large Language
Models](https://arxiv.org/abs/2301.10226).
"""
@dataclass_with_extra
class TextGenerationInput(BaseInferenceType):
"""Text Generation Input.
Auto-generated from TGI specs.
For more details, check out
https://github.com/huggingface/huggingface.js/blob/main/packages/tasks/scripts/inference-tgi-import.ts.
"""
inputs: str
parameters: Optional[TextGenerationInputGenerateParameters] = None
stream: Optional[bool] = None
TextGenerationOutputFinishReason = Literal["length", "eos_token", "stop_sequence"]
@dataclass_with_extra
class TextGenerationOutputPrefillToken(BaseInferenceType):
id: int
logprob: float
text: str
@dataclass_with_extra
class TextGenerationOutputToken(BaseInferenceType):
id: int
logprob: float
special: bool
text: str
@dataclass_with_extra
class TextGenerationOutputBestOfSequence(BaseInferenceType):
finish_reason: "TextGenerationOutputFinishReason"
generated_text: str
generated_tokens: int
prefill: List[TextGenerationOutputPrefillToken]
tokens: List[TextGenerationOutputToken]
seed: Optional[int] = None
top_tokens: Optional[List[List[TextGenerationOutputToken]]] = None
@dataclass_with_extra
class TextGenerationOutputDetails(BaseInferenceType):
finish_reason: "TextGenerationOutputFinishReason"
generated_tokens: int
prefill: List[TextGenerationOutputPrefillToken]
tokens: List[TextGenerationOutputToken]
best_of_sequences: Optional[List[TextGenerationOutputBestOfSequence]] = None
seed: Optional[int] = None
top_tokens: Optional[List[List[TextGenerationOutputToken]]] = None
@dataclass_with_extra
class TextGenerationOutput(BaseInferenceType):
"""Text Generation Output.
Auto-generated from TGI specs.
For more details, check out
https://github.com/huggingface/huggingface.js/blob/main/packages/tasks/scripts/inference-tgi-import.ts.
"""
generated_text: str
details: Optional[TextGenerationOutputDetails] = None
@dataclass_with_extra
class TextGenerationStreamOutputStreamDetails(BaseInferenceType):
finish_reason: "TextGenerationOutputFinishReason"
generated_tokens: int
input_length: int
seed: Optional[int] = None
@dataclass_with_extra
class TextGenerationStreamOutputToken(BaseInferenceType):
id: int
logprob: float
special: bool
text: str
@dataclass_with_extra
class TextGenerationStreamOutput(BaseInferenceType):
"""Text Generation Stream Output.
Auto-generated from TGI specs.
For more details, check out
https://github.com/huggingface/huggingface.js/blob/main/packages/tasks/scripts/inference-tgi-import.ts.
"""
index: int
token: TextGenerationStreamOutputToken
details: Optional[TextGenerationStreamOutputStreamDetails] = None
generated_text: Optional[str] = None
top_tokens: Optional[List[TextGenerationStreamOutputToken]] = None
kronos-venv/lib/python3.12/site-packages/huggingface_hub/inference/_generated/types/text_to_audio.py
+99
View File
@@ -0,0 +1,99 @@
# Inference code generated from the JSON schema spec in @huggingface/tasks.
#
# See:
# - script: https://github.com/huggingface/huggingface.js/blob/main/packages/tasks/scripts/inference-codegen.ts
# - specs: https://github.com/huggingface/huggingface.js/tree/main/packages/tasks/src/tasks.
from typing import Any, Literal, Optional, Union
from .base import BaseInferenceType, dataclass_with_extra
TextToAudioEarlyStoppingEnum = Literal["never"]
@dataclass_with_extra
class TextToAudioGenerationParameters(BaseInferenceType):
"""Parametrization of the text generation process"""
do_sample: Optional[bool] = None
"""Whether to use sampling instead of greedy decoding when generating new tokens."""
early_stopping: Optional[Union[bool, "TextToAudioEarlyStoppingEnum"]] = None
"""Controls the stopping condition for beam-based methods."""
epsilon_cutoff: Optional[float] = None
"""If set to float strictly between 0 and 1, only tokens with a conditional probability
greater than epsilon_cutoff will be sampled. In the paper, suggested values range from
3e-4 to 9e-4, depending on the size of the model. See [Truncation Sampling as Language
Model Desmoothing](https://hf.co/papers/2210.15191) for more details.
"""
eta_cutoff: Optional[float] = None
"""Eta sampling is a hybrid of locally typical sampling and epsilon sampling. If set to
float strictly between 0 and 1, a token is only considered if it is greater than either
eta_cutoff or sqrt(eta_cutoff) * exp(-entropy(softmax(next_token_logits))). The latter
term is intuitively the expected next token probability, scaled by sqrt(eta_cutoff). In
the paper, suggested values range from 3e-4 to 2e-3, depending on the size of the model.
See [Truncation Sampling as Language Model Desmoothing](https://hf.co/papers/2210.15191)
for more details.
"""
max_length: Optional[int] = None
"""The maximum length (in tokens) of the generated text, including the input."""
max_new_tokens: Optional[int] = None
"""The maximum number of tokens to generate. Takes precedence over max_length."""
min_length: Optional[int] = None
"""The minimum length (in tokens) of the generated text, including the input."""
min_new_tokens: Optional[int] = None
"""The minimum number of tokens to generate. Takes precedence over min_length."""
num_beam_groups: Optional[int] = None
"""Number of groups to divide num_beams into in order to ensure diversity among different
groups of beams. See [this paper](https://hf.co/papers/1610.02424) for more details.
"""
num_beams: Optional[int] = None
"""Number of beams to use for beam search."""
penalty_alpha: Optional[float] = None
"""The value balances the model confidence and the degeneration penalty in contrastive
search decoding.
"""
temperature: Optional[float] = None
"""The value used to modulate the next token probabilities."""
top_k: Optional[int] = None
"""The number of highest probability vocabulary tokens to keep for top-k-filtering."""
top_p: Optional[float] = None
"""If set to float < 1, only the smallest set of most probable tokens with probabilities
that add up to top_p or higher are kept for generation.
"""
typical_p: Optional[float] = None
"""Local typicality measures how similar the conditional probability of predicting a target
token next is to the expected conditional probability of predicting a random token next,
given the partial text already generated. If set to float < 1, the smallest set of the
most locally typical tokens with probabilities that add up to typical_p or higher are
kept for generation. See [this paper](https://hf.co/papers/2202.00666) for more details.
"""
use_cache: Optional[bool] = None
"""Whether the model should use the past last key/values attentions to speed up decoding"""
@dataclass_with_extra
class TextToAudioParameters(BaseInferenceType):
"""Additional inference parameters for Text To Audio"""
generation_parameters: Optional[TextToAudioGenerationParameters] = None
"""Parametrization of the text generation process"""
@dataclass_with_extra
class TextToAudioInput(BaseInferenceType):
"""Inputs for Text To Audio inference"""
inputs: str
"""The input text data"""
parameters: Optional[TextToAudioParameters] = None
"""Additional inference parameters for Text To Audio"""
@dataclass_with_extra
class TextToAudioOutput(BaseInferenceType):
"""Outputs of inference for the Text To Audio task"""
audio: Any
"""The generated audio waveform."""
sampling_rate: float
"""The sampling rate of the generated audio waveform."""
kronos-venv/lib/python3.12/site-packages/huggingface_hub/inference/_generated/types/text_to_image.py
+50
View File
@@ -0,0 +1,50 @@
# Inference code generated from the JSON schema spec in @huggingface/tasks.
#
# See:
# - script: https://github.com/huggingface/huggingface.js/blob/main/packages/tasks/scripts/inference-codegen.ts
# - specs: https://github.com/huggingface/huggingface.js/tree/main/packages/tasks/src/tasks.
from typing import Any, Optional
from .base import BaseInferenceType, dataclass_with_extra
@dataclass_with_extra
class TextToImageParameters(BaseInferenceType):
"""Additional inference parameters for Text To Image"""
guidance_scale: Optional[float] = None
"""A higher guidance scale value encourages the model to generate images closely linked to
the text prompt, but values too high may cause saturation and other artifacts.
"""
height: Optional[int] = None
"""The height in pixels of the output image"""
negative_prompt: Optional[str] = None
"""One prompt to guide what NOT to include in image generation."""
num_inference_steps: Optional[int] = None
"""The number of denoising steps. More denoising steps usually lead to a higher quality
image at the expense of slower inference.
"""
scheduler: Optional[str] = None
"""Override the scheduler with a compatible one."""
seed: Optional[int] = None
"""Seed for the random number generator."""
width: Optional[int] = None
"""The width in pixels of the output image"""
@dataclass_with_extra
class TextToImageInput(BaseInferenceType):
"""Inputs for Text To Image inference"""
inputs: str
"""The input text data (sometimes called "prompt")"""
parameters: Optional[TextToImageParameters] = None
"""Additional inference parameters for Text To Image"""
@dataclass_with_extra
class TextToImageOutput(BaseInferenceType):
"""Outputs of inference for the Text To Image task"""
image: Any
"""The generated image returned as raw bytes in the payload."""
kronos-venv/lib/python3.12/site-packages/huggingface_hub/inference/_generated/types/text_to_speech.py
+99
View File
@@ -0,0 +1,99 @@
# Inference code generated from the JSON schema spec in @huggingface/tasks.
#
# See:
# - script: https://github.com/huggingface/huggingface.js/blob/main/packages/tasks/scripts/inference-codegen.ts
# - specs: https://github.com/huggingface/huggingface.js/tree/main/packages/tasks/src/tasks.
from typing import Any, Literal, Optional, Union
from .base import BaseInferenceType, dataclass_with_extra
TextToSpeechEarlyStoppingEnum = Literal["never"]
@dataclass_with_extra
class TextToSpeechGenerationParameters(BaseInferenceType):
"""Parametrization of the text generation process"""
do_sample: Optional[bool] = None
"""Whether to use sampling instead of greedy decoding when generating new tokens."""
early_stopping: Optional[Union[bool, "TextToSpeechEarlyStoppingEnum"]] = None
"""Controls the stopping condition for beam-based methods."""
epsilon_cutoff: Optional[float] = None
"""If set to float strictly between 0 and 1, only tokens with a conditional probability
greater than epsilon_cutoff will be sampled. In the paper, suggested values range from
3e-4 to 9e-4, depending on the size of the model. See [Truncation Sampling as Language
Model Desmoothing](https://hf.co/papers/2210.15191) for more details.
"""
eta_cutoff: Optional[float] = None
"""Eta sampling is a hybrid of locally typical sampling and epsilon sampling. If set to
float strictly between 0 and 1, a token is only considered if it is greater than either
eta_cutoff or sqrt(eta_cutoff) * exp(-entropy(softmax(next_token_logits))). The latter
term is intuitively the expected next token probability, scaled by sqrt(eta_cutoff). In
the paper, suggested values range from 3e-4 to 2e-3, depending on the size of the model.
See [Truncation Sampling as Language Model Desmoothing](https://hf.co/papers/2210.15191)
for more details.
"""
max_length: Optional[int] = None
"""The maximum length (in tokens) of the generated text, including the input."""
max_new_tokens: Optional[int] = None
"""The maximum number of tokens to generate. Takes precedence over max_length."""
min_length: Optional[int] = None
"""The minimum length (in tokens) of the generated text, including the input."""
min_new_tokens: Optional[int] = None
"""The minimum number of tokens to generate. Takes precedence over min_length."""
num_beam_groups: Optional[int] = None
"""Number of groups to divide num_beams into in order to ensure diversity among different
groups of beams. See [this paper](https://hf.co/papers/1610.02424) for more details.
"""
num_beams: Optional[int] = None
"""Number of beams to use for beam search."""
penalty_alpha: Optional[float] = None
"""The value balances the model confidence and the degeneration penalty in contrastive
search decoding.
"""
temperature: Optional[float] = None
"""The value used to modulate the next token probabilities."""
top_k: Optional[int] = None
"""The number of highest probability vocabulary tokens to keep for top-k-filtering."""
top_p: Optional[float] = None
"""If set to float < 1, only the smallest set of most probable tokens with probabilities
that add up to top_p or higher are kept for generation.
"""
typical_p: Optional[float] = None
"""Local typicality measures how similar the conditional probability of predicting a target
token next is to the expected conditional probability of predicting a random token next,
given the partial text already generated. If set to float < 1, the smallest set of the
most locally typical tokens with probabilities that add up to typical_p or higher are
kept for generation. See [this paper](https://hf.co/papers/2202.00666) for more details.
"""
use_cache: Optional[bool] = None
"""Whether the model should use the past last key/values attentions to speed up decoding"""
@dataclass_with_extra
class TextToSpeechParameters(BaseInferenceType):
"""Additional inference parameters for Text To Speech"""
generation_parameters: Optional[TextToSpeechGenerationParameters] = None
"""Parametrization of the text generation process"""
@dataclass_with_extra
class TextToSpeechInput(BaseInferenceType):
"""Inputs for Text To Speech inference"""
inputs: str
"""The input text data"""
parameters: Optional[TextToSpeechParameters] = None
"""Additional inference parameters for Text To Speech"""
@dataclass_with_extra
class TextToSpeechOutput(BaseInferenceType):
"""Outputs of inference for the Text To Speech task"""
audio: Any
"""The generated audio"""
sampling_rate: Optional[float] = None
"""The sampling rate of the generated audio waveform."""
kronos-venv/lib/python3.12/site-packages/huggingface_hub/inference/_generated/types/text_to_video.py
+46
View File
@@ -0,0 +1,46 @@
# Inference code generated from the JSON schema spec in @huggingface/tasks.
#
# See:
# - script: https://github.com/huggingface/huggingface.js/blob/main/packages/tasks/scripts/inference-codegen.ts
# - specs: https://github.com/huggingface/huggingface.js/tree/main/packages/tasks/src/tasks.
from typing import Any, List, Optional
from .base import BaseInferenceType, dataclass_with_extra
@dataclass_with_extra
class TextToVideoParameters(BaseInferenceType):
"""Additional inference parameters for Text To Video"""
guidance_scale: Optional[float] = None
"""A higher guidance scale value encourages the model to generate videos closely linked to
the text prompt, but values too high may cause saturation and other artifacts.
"""
negative_prompt: Optional[List[str]] = None
"""One or several prompt to guide what NOT to include in video generation."""
num_frames: Optional[float] = None
"""The num_frames parameter determines how many video frames are generated."""
num_inference_steps: Optional[int] = None
"""The number of denoising steps. More denoising steps usually lead to a higher quality
video at the expense of slower inference.
"""
seed: Optional[int] = None
"""Seed for the random number generator."""
@dataclass_with_extra
class TextToVideoInput(BaseInferenceType):
"""Inputs for Text To Video inference"""
inputs: str
"""The input text data (sometimes called "prompt")"""
parameters: Optional[TextToVideoParameters] = None
"""Additional inference parameters for Text To Video"""
@dataclass_with_extra
class TextToVideoOutput(BaseInferenceType):
"""Outputs of inference for the Text To Video task"""
video: Any
"""The generated video returned as raw bytes in the payload."""
kronos-venv/lib/python3.12/site-packages/huggingface_hub/inference/_generated/types/token_classification.py
+51
View File
@@ -0,0 +1,51 @@
# Inference code generated from the JSON schema spec in @huggingface/tasks.
#
# See:
# - script: https://github.com/huggingface/huggingface.js/blob/main/packages/tasks/scripts/inference-codegen.ts
# - specs: https://github.com/huggingface/huggingface.js/tree/main/packages/tasks/src/tasks.
from typing import List, Literal, Optional
from .base import BaseInferenceType, dataclass_with_extra
TokenClassificationAggregationStrategy = Literal["none", "simple", "first", "average", "max"]
@dataclass_with_extra
class TokenClassificationParameters(BaseInferenceType):
"""Additional inference parameters for Token Classification"""
aggregation_strategy: Optional["TokenClassificationAggregationStrategy"] = None
"""The strategy used to fuse tokens based on model predictions"""
ignore_labels: Optional[List[str]] = None
"""A list of labels to ignore"""
stride: Optional[int] = None
"""The number of overlapping tokens between chunks when splitting the input text."""
@dataclass_with_extra
class TokenClassificationInput(BaseInferenceType):
"""Inputs for Token Classification inference"""
inputs: str
"""The input text data"""
parameters: Optional[TokenClassificationParameters] = None
"""Additional inference parameters for Token Classification"""
@dataclass_with_extra
class TokenClassificationOutputElement(BaseInferenceType):
"""Outputs of inference for the Token Classification task"""
end: int
"""The character position in the input where this group ends."""
score: float
"""The associated score / probability"""
start: int
"""The character position in the input where this group begins."""
word: str
"""The corresponding text"""
entity: Optional[str] = None
"""The predicted label for a single token"""
entity_group: Optional[str] = None
"""The predicted label for a group of one or more tokens"""
kronos-venv/lib/python3.12/site-packages/huggingface_hub/inference/_generated/types/translation.py
+49
View File
@@ -0,0 +1,49 @@
# Inference code generated from the JSON schema spec in @huggingface/tasks.
#
# See:
# - script: https://github.com/huggingface/huggingface.js/blob/main/packages/tasks/scripts/inference-codegen.ts
# - specs: https://github.com/huggingface/huggingface.js/tree/main/packages/tasks/src/tasks.
from typing import Any, Dict, Literal, Optional
from .base import BaseInferenceType, dataclass_with_extra
TranslationTruncationStrategy = Literal["do_not_truncate", "longest_first", "only_first", "only_second"]
@dataclass_with_extra
class TranslationParameters(BaseInferenceType):
"""Additional inference parameters for Translation"""
clean_up_tokenization_spaces: Optional[bool] = None
"""Whether to clean up the potential extra spaces in the text output."""
generate_parameters: Optional[Dict[str, Any]] = None
"""Additional parametrization of the text generation algorithm."""
src_lang: Optional[str] = None
"""The source language of the text. Required for models that can translate from multiple
languages.
"""
tgt_lang: Optional[str] = None
"""Target language to translate to. Required for models that can translate to multiple
languages.
"""
truncation: Optional["TranslationTruncationStrategy"] = None
"""The truncation strategy to use."""
@dataclass_with_extra
class TranslationInput(BaseInferenceType):
"""Inputs for Translation inference"""
inputs: str
"""The text to translate."""
parameters: Optional[TranslationParameters] = None
"""Additional inference parameters for Translation"""
@dataclass_with_extra
class TranslationOutput(BaseInferenceType):
"""Outputs of inference for the Translation task"""
translation_text: str
"""The translated text."""
kronos-venv/lib/python3.12/site-packages/huggingface_hub/inference/_generated/types/video_classification.py
+45
View File
@@ -0,0 +1,45 @@
# Inference code generated from the JSON schema spec in @huggingface/tasks.
#
# See:
# - script: https://github.com/huggingface/huggingface.js/blob/main/packages/tasks/scripts/inference-codegen.ts
# - specs: https://github.com/huggingface/huggingface.js/tree/main/packages/tasks/src/tasks.
from typing import Any, Literal, Optional
from .base import BaseInferenceType, dataclass_with_extra
VideoClassificationOutputTransform = Literal["sigmoid", "softmax", "none"]
@dataclass_with_extra
class VideoClassificationParameters(BaseInferenceType):
"""Additional inference parameters for Video Classification"""
frame_sampling_rate: Optional[int] = None
"""The sampling rate used to select frames from the video."""
function_to_apply: Optional["VideoClassificationOutputTransform"] = None
"""The function to apply to the model outputs in order to retrieve the scores."""
num_frames: Optional[int] = None
"""The number of sampled frames to consider for classification."""
top_k: Optional[int] = None
"""When specified, limits the output to the top K most probable classes."""
@dataclass_with_extra
class VideoClassificationInput(BaseInferenceType):
"""Inputs for Video Classification inference"""
inputs: Any
"""The input video data"""
parameters: Optional[VideoClassificationParameters] = None
"""Additional inference parameters for Video Classification"""
@dataclass_with_extra
class VideoClassificationOutputElement(BaseInferenceType):
"""Outputs of inference for the Video Classification task"""
label: str
"""The predicted class label."""
score: float
"""The corresponding probability."""
kronos-venv/lib/python3.12/site-packages/huggingface_hub/inference/_generated/types/visual_question_answering.py
+49
View File
@@ -0,0 +1,49 @@
# Inference code generated from the JSON schema spec in @huggingface/tasks.
#
# See:
# - script: https://github.com/huggingface/huggingface.js/blob/main/packages/tasks/scripts/inference-codegen.ts
# - specs: https://github.com/huggingface/huggingface.js/tree/main/packages/tasks/src/tasks.
from typing import Any, Optional
from .base import BaseInferenceType, dataclass_with_extra
@dataclass_with_extra
class VisualQuestionAnsweringInputData(BaseInferenceType):
"""One (image, question) pair to answer"""
image: Any
"""The image."""
question: str
"""The question to answer based on the image."""
@dataclass_with_extra
class VisualQuestionAnsweringParameters(BaseInferenceType):
"""Additional inference parameters for Visual Question Answering"""
top_k: Optional[int] = None
"""The number of answers to return (will be chosen by order of likelihood). Note that we
return less than topk answers if there are not enough options available within the
context.
"""
@dataclass_with_extra
class VisualQuestionAnsweringInput(BaseInferenceType):
"""Inputs for Visual Question Answering inference"""
inputs: VisualQuestionAnsweringInputData
"""One (image, question) pair to answer"""
parameters: Optional[VisualQuestionAnsweringParameters] = None
"""Additional inference parameters for Visual Question Answering"""
@dataclass_with_extra
class VisualQuestionAnsweringOutputElement(BaseInferenceType):
"""Outputs of inference for the Visual Question Answering task"""
score: float
"""The associated score / probability"""
answer: Optional[str] = None
"""The answer to the question"""
kronos-venv/lib/python3.12/site-packages/huggingface_hub/inference/_generated/types/zero_shot_classification.py
+45
View File
@@ -0,0 +1,45 @@
# Inference code generated from the JSON schema spec in @huggingface/tasks.
#
# See:
# - script: https://github.com/huggingface/huggingface.js/blob/main/packages/tasks/scripts/inference-codegen.ts
# - specs: https://github.com/huggingface/huggingface.js/tree/main/packages/tasks/src/tasks.
from typing import List, Optional
from .base import BaseInferenceType, dataclass_with_extra
@dataclass_with_extra
class ZeroShotClassificationParameters(BaseInferenceType):
"""Additional inference parameters for Zero Shot Classification"""
candidate_labels: List[str]
"""The set of possible class labels to classify the text into."""
hypothesis_template: Optional[str] = None
"""The sentence used in conjunction with `candidate_labels` to attempt the text
classification by replacing the placeholder with the candidate labels.
"""
multi_label: Optional[bool] = None
"""Whether multiple candidate labels can be true. If false, the scores are normalized such
that the sum of the label likelihoods for each sequence is 1. If true, the labels are
considered independent and probabilities are normalized for each candidate.
"""
@dataclass_with_extra
class ZeroShotClassificationInput(BaseInferenceType):
"""Inputs for Zero Shot Classification inference"""
inputs: str
"""The text to classify"""
parameters: ZeroShotClassificationParameters
"""Additional inference parameters for Zero Shot Classification"""
@dataclass_with_extra
class ZeroShotClassificationOutputElement(BaseInferenceType):
"""Outputs of inference for the Zero Shot Classification task"""
label: str
"""The predicted class label."""
score: float
"""The corresponding probability."""
kronos-venv/lib/python3.12/site-packages/huggingface_hub/inference/_generated/types/zero_shot_image_classification.py
+40
View File
@@ -0,0 +1,40 @@
# Inference code generated from the JSON schema spec in @huggingface/tasks.
#
# See:
# - script: https://github.com/huggingface/huggingface.js/blob/main/packages/tasks/scripts/inference-codegen.ts
# - specs: https://github.com/huggingface/huggingface.js/tree/main/packages/tasks/src/tasks.
from typing import List, Optional
from .base import BaseInferenceType, dataclass_with_extra
@dataclass_with_extra
class ZeroShotImageClassificationParameters(BaseInferenceType):
"""Additional inference parameters for Zero Shot Image Classification"""
candidate_labels: List[str]
"""The candidate labels for this image"""
hypothesis_template: Optional[str] = None
"""The sentence used in conjunction with `candidate_labels` to attempt the image
classification by replacing the placeholder with the candidate labels.
"""
@dataclass_with_extra
class ZeroShotImageClassificationInput(BaseInferenceType):
"""Inputs for Zero Shot Image Classification inference"""
inputs: str
"""The input image data to classify as a base64-encoded string."""
parameters: ZeroShotImageClassificationParameters
"""Additional inference parameters for Zero Shot Image Classification"""
@dataclass_with_extra
class ZeroShotImageClassificationOutputElement(BaseInferenceType):
"""Outputs of inference for the Zero Shot Image Classification task"""
label: str
"""The predicted class label."""
score: float
"""The corresponding probability."""
kronos-venv/lib/python3.12/site-packages/huggingface_hub/inference/_generated/types/zero_shot_object_detection.py
+52
View File
@@ -0,0 +1,52 @@
# Inference code generated from the JSON schema spec in @huggingface/tasks.
#
# See:
# - script: https://github.com/huggingface/huggingface.js/blob/main/packages/tasks/scripts/inference-codegen.ts
# - specs: https://github.com/huggingface/huggingface.js/tree/main/packages/tasks/src/tasks.
from typing import List
from .base import BaseInferenceType, dataclass_with_extra
@dataclass_with_extra
class ZeroShotObjectDetectionParameters(BaseInferenceType):
"""Additional inference parameters for Zero Shot Object Detection"""
candidate_labels: List[str]
"""The candidate labels for this image"""
@dataclass_with_extra
class ZeroShotObjectDetectionInput(BaseInferenceType):
"""Inputs for Zero Shot Object Detection inference"""
inputs: str
"""The input image data as a base64-encoded string."""
parameters: ZeroShotObjectDetectionParameters
"""Additional inference parameters for Zero Shot Object Detection"""
@dataclass_with_extra
class ZeroShotObjectDetectionBoundingBox(BaseInferenceType):
"""The predicted bounding box. Coordinates are relative to the top left corner of the input
image.
"""
xmax: int
xmin: int
ymax: int
ymin: int
@dataclass_with_extra
class ZeroShotObjectDetectionOutputElement(BaseInferenceType):
"""Outputs of inference for the Zero Shot Object Detection task"""
box: ZeroShotObjectDetectionBoundingBox
"""The predicted bounding box. Coordinates are relative to the top left corner of the input
image.
"""
label: str
"""A candidate label"""
score: float
"""The associated score / probability"""
kronos-venv/lib/python3.12/site-packages/huggingface_hub/inference/_mcp/__init__.py
View File
kronos-venv/lib/python3.12/site-packages/huggingface_hub/inference/_mcp/_cli_hacks.py
+88
View File
@@ -0,0 +1,88 @@
import asyncio
import sys
from functools import partial
import typer
def _patch_anyio_open_process():
"""
Patch anyio.open_process to allow detached processes on Windows and Unix-like systems.
This is necessary to prevent the MCP client from being interrupted by Ctrl+C when running in the CLI.
"""
import subprocess
import anyio
if getattr(anyio, "_tiny_agents_patched", False):
return
anyio._tiny_agents_patched = True
original_open_process = anyio.open_process
if sys.platform == "win32":
# On Windows, we need to set the creation flags to create a new process group
async def open_process_in_new_group(*args, **kwargs):
"""
Wrapper for open_process to handle Windows-specific process creation flags.
"""
# Ensure we pass the creation flags for Windows
kwargs.setdefault("creationflags", subprocess.CREATE_NEW_PROCESS_GROUP)
return await original_open_process(*args, **kwargs)
anyio.open_process = open_process_in_new_group
else:
# For Unix-like systems, we can use setsid to create a new session
async def open_process_in_new_group(*args, **kwargs):
"""
Wrapper for open_process to handle Unix-like systems with start_new_session=True.
"""
kwargs.setdefault("start_new_session", True)
return await original_open_process(*args, **kwargs)
anyio.open_process = open_process_in_new_group
async def _async_prompt(exit_event: asyncio.Event, prompt: str = "» ") -> str:
"""
Asynchronous prompt function that reads input from stdin without blocking.
This function is designed to work in an asynchronous context, allowing the event loop to gracefully stop it (e.g. on Ctrl+C).
Alternatively, we could use https://github.com/vxgmichel/aioconsole but that would be an additional dependency.
"""
loop = asyncio.get_event_loop()
if sys.platform == "win32":
# Windows: Use run_in_executor to avoid blocking the event loop
# Degraded solution: this is not ideal as user will have to CTRL+C once more to stop the prompt (and it'll not be graceful)
return await loop.run_in_executor(None, partial(typer.prompt, prompt, prompt_suffix=" "))
else:
# UNIX-like: Use loop.add_reader for non-blocking stdin read
future = loop.create_future()
def on_input():
line = sys.stdin.readline()
loop.remove_reader(sys.stdin)
future.set_result(line)
print(prompt, end=" ", flush=True)
loop.add_reader(sys.stdin, on_input) # not supported on Windows
# Wait for user input or exit event
# Wait until either the user hits enter or exit_event is set
exit_task = asyncio.create_task(exit_event.wait())
await asyncio.wait(
[future, exit_task],
return_when=asyncio.FIRST_COMPLETED,
)
# Check which one has been triggered
if exit_event.is_set():
future.cancel()
return ""
line = await future
return line.strip()
kronos-venv/lib/python3.12/site-packages/huggingface_hub/inference/_mcp/agent.py
+102
View File
@@ -0,0 +1,102 @@
from __future__ import annotations
import asyncio
from typing import AsyncGenerator, Dict, Iterable, List, Optional, Union
from huggingface_hub import ChatCompletionInputMessage, ChatCompletionStreamOutput, MCPClient
from .._providers import PROVIDER_OR_POLICY_T
from .constants import DEFAULT_SYSTEM_PROMPT, EXIT_LOOP_TOOLS, MAX_NUM_TURNS
class Agent(MCPClient):
"""
Implementation of a Simple Agent, which is a simple while loop built right on top of an [`MCPClient`].
<Tip warning={true}>
This class is experimental and might be subject to breaking changes in the future without prior notice.
</Tip>
Args:
model (`str`, *optional*):
The model to run inference with. Can be a model id hosted on the Hugging Face Hub, e.g. `meta-llama/Meta-Llama-3-8B-Instruct`
or a URL to a deployed Inference Endpoint or other local or remote endpoint.
servers (`Iterable[Dict]`):
MCP servers to connect to. Each server is a dictionary containing a `type` key and a `config` key. The `type` key can be `"stdio"` or `"sse"`, and the `config` key is a dictionary of arguments for the server.
provider (`str`, *optional*):
Name of the provider to use for inference. Defaults to "auto" i.e. the first of the providers available for the model, sorted by the user's order in https://hf.co/settings/inference-providers.
If model is a URL or `base_url` is passed, then `provider` is not used.
base_url (`str`, *optional*):
The base URL to run inference. Defaults to None.
api_key (`str`, *optional*):
Token to use for authentication. Will default to the locally Hugging Face saved token if not provided. You can also use your own provider API key to interact directly with the provider's service.
prompt (`str`, *optional*):
The system prompt to use for the agent. Defaults to the default system prompt in `constants.py`.
"""
def __init__(
self,
*,
model: Optional[str] = None,
servers: Iterable[Dict],
provider: Optional[PROVIDER_OR_POLICY_T] = None,
base_url: Optional[str] = None,
api_key: Optional[str] = None,
prompt: Optional[str] = None,
):
super().__init__(model=model, provider=provider, base_url=base_url, api_key=api_key)
self._servers_cfg = list(servers)
self.messages: List[Union[Dict, ChatCompletionInputMessage]] = [
{"role": "system", "content": prompt or DEFAULT_SYSTEM_PROMPT}
]
async def load_tools(self) -> None:
for cfg in self._servers_cfg:
await self.add_mcp_server(cfg["type"], **cfg["config"])
async def run(
self,
user_input: str,
*,
abort_event: Optional[asyncio.Event] = None,
) -> AsyncGenerator[Union[ChatCompletionStreamOutput, ChatCompletionInputMessage], None]:
"""
Run the agent with the given user input.
Args:
user_input (`str`):
The user input to run the agent with.
abort_event (`asyncio.Event`, *optional*):
An event that can be used to abort the agent. If the event is set, the agent will stop running.
"""
self.messages.append({"role": "user", "content": user_input})
num_turns: int = 0
next_turn_should_call_tools = True
while True:
if abort_event and abort_event.is_set():
return
async for item in self.process_single_turn_with_tools(
self.messages,
exit_loop_tools=EXIT_LOOP_TOOLS,
exit_if_first_chunk_no_tool=(num_turns > 0 and next_turn_should_call_tools),
):
yield item
num_turns += 1
last = self.messages[-1]
if last.get("role") == "tool" and last.get("name") in {t.function.name for t in EXIT_LOOP_TOOLS}:
return
if last.get("role") != "tool" and num_turns > MAX_NUM_TURNS:
return
if last.get("role") != "tool" and next_turn_should_call_tools:
return
next_turn_should_call_tools = last.get("role") != "tool"
kronos-venv/lib/python3.12/site-packages/huggingface_hub/inference/_mcp/cli.py
+236
View File
@@ -0,0 +1,236 @@
import asyncio
import os
import signal
import traceback
from typing import Optional
import typer
from rich import print
from ._cli_hacks import _async_prompt, _patch_anyio_open_process
from .agent import Agent
from .utils import _load_agent_config
app = typer.Typer(
rich_markup_mode="rich",
help="A squad of lightweight composable AI applications built on Hugging Face's Inference Client and MCP stack.",
)
run_cli = typer.Typer(
name="run",
help="Run the Agent in the CLI",
invoke_without_command=True,
)
app.add_typer(run_cli, name="run")
async def run_agent(
agent_path: Optional[str],
) -> None:
"""
Tiny Agent loop.
Args:
agent_path (`str`, *optional*):
Path to a local folder containing an `agent.json` and optionally a custom `PROMPT.md` file or a built-in agent stored in a Hugging Face dataset.
"""
_patch_anyio_open_process() # Hacky way to prevent stdio connections to be stopped by Ctrl+C
config, prompt = _load_agent_config(agent_path)
inputs = config.get("inputs", [])
servers = config.get("servers", [])
abort_event = asyncio.Event()
exit_event = asyncio.Event()
first_sigint = True
loop = asyncio.get_running_loop()
original_sigint_handler = signal.getsignal(signal.SIGINT)
def _sigint_handler() -> None:
nonlocal first_sigint
if first_sigint:
first_sigint = False
abort_event.set()
print("\n[red]Interrupted. Press Ctrl+C again to quit.[/red]", flush=True)
return
print("\n[red]Exiting...[/red]", flush=True)
exit_event.set()
try:
sigint_registered_in_loop = False
try:
loop.add_signal_handler(signal.SIGINT, _sigint_handler)
sigint_registered_in_loop = True
except (AttributeError, NotImplementedError):
# Windows (or any loop that doesn't support it) : fall back to sync
signal.signal(signal.SIGINT, lambda *_: _sigint_handler())
# Handle inputs (i.e. env variables injection)
if len(inputs) > 0:
print(
"[bold blue]Some initial inputs are required by the agent. "
"Please provide a value or leave empty to load from env.[/bold blue]"
)
for input_item in inputs:
input_id = input_item["id"]
description = input_item["description"]
env_special_value = "${input:" + input_id + "}" # Special value to indicate env variable injection
# Check env variables that will use this input
input_vars = set()
for server in servers:
# Check stdio's "env" and http/sse's "headers" mappings
env_or_headers = (
server["config"].get("env", {})
if server["type"] == "stdio"
else server["config"].get("options", {}).get("requestInit", {}).get("headers", {})
)
for key, value in env_or_headers.items():
if env_special_value in value:
input_vars.add(key)
if not input_vars:
print(f"[yellow]Input {input_id} defined in config but not used by any server.[/yellow]")
continue
# Prompt user for input
print(
f"[blue] • {input_id}[/blue]: {description}. (default: load from {', '.join(sorted(input_vars))}).",
end=" ",
)
user_input = (await _async_prompt(exit_event=exit_event)).strip()
if exit_event.is_set():
return
# Inject user input (or env variable) into stdio's env or http/sse's headers
for server in servers:
env_or_headers = (
server["config"].get("env", {})
if server["type"] == "stdio"
else server["config"].get("options", {}).get("requestInit", {}).get("headers", {})
)
for key, value in env_or_headers.items():
if env_special_value in value:
if user_input:
env_or_headers[key] = env_or_headers[key].replace(env_special_value, user_input)
else:
value_from_env = os.getenv(key, "")
env_or_headers[key] = env_or_headers[key].replace(env_special_value, value_from_env)
if value_from_env:
print(f"[green]Value successfully loaded from '{key}'[/green]")
else:
print(
f"[yellow]No value found for '{key}' in environment variables. Continuing.[/yellow]"
)
print()
# Main agent loop
async with Agent(
provider=config.get("provider"), # type: ignore[arg-type]
model=config.get("model"),
base_url=config.get("endpointUrl"), # type: ignore[arg-type]
servers=servers, # type: ignore[arg-type]
prompt=prompt,
) as agent:
await agent.load_tools()
print(f"[bold blue]Agent loaded with {len(agent.available_tools)} tools:[/bold blue]")
for t in agent.available_tools:
print(f"[blue] • {t.function.name}[/blue]")
while True:
abort_event.clear()
# Check if we should exit
if exit_event.is_set():
return
try:
user_input = await _async_prompt(exit_event=exit_event)
first_sigint = True
except EOFError:
print("\n[red]EOF received, exiting.[/red]", flush=True)
break
except KeyboardInterrupt:
if not first_sigint and abort_event.is_set():
continue
else:
print("\n[red]Keyboard interrupt during input processing.[/red]", flush=True)
break
try:
async for chunk in agent.run(user_input, abort_event=abort_event):
if abort_event.is_set() and not first_sigint:
break
if exit_event.is_set():
return
if hasattr(chunk, "choices"):
delta = chunk.choices[0].delta
if delta.content:
print(delta.content, end="", flush=True)
if delta.tool_calls:
for call in delta.tool_calls:
if call.id:
print(f"<Tool {call.id}>", end="")
if call.function.name:
print(f"{call.function.name}", end=" ")
if call.function.arguments:
print(f"{call.function.arguments}", end="")
else:
print(
f"\n\n[green]Tool[{chunk.name}] {chunk.tool_call_id}\n{chunk.content}[/green]\n",
flush=True,
)
print()
except Exception as e:
tb_str = traceback.format_exc()
print(f"\n[bold red]Error during agent run: {e}\n{tb_str}[/bold red]", flush=True)
first_sigint = True # Allow graceful interrupt for the next command
except Exception as e:
tb_str = traceback.format_exc()
print(f"\n[bold red]An unexpected error occurred: {e}\n{tb_str}[/bold red]", flush=True)
raise e
finally:
if sigint_registered_in_loop:
try:
loop.remove_signal_handler(signal.SIGINT)
except (AttributeError, NotImplementedError):
pass
else:
signal.signal(signal.SIGINT, original_sigint_handler)
@run_cli.callback()
def run(
path: Optional[str] = typer.Argument(
None,
help=(
"Path to a local folder containing an agent.json file or a built-in agent "
"stored in the 'tiny-agents/tiny-agents' Hugging Face dataset "
"(https://huggingface.co/datasets/tiny-agents/tiny-agents)"
),
show_default=False,
),
):
try:
asyncio.run(run_agent(path))
except KeyboardInterrupt:
print("\n[red]Application terminated by KeyboardInterrupt.[/red]", flush=True)
raise typer.Exit(code=130)
except Exception as e:
print(f"\n[bold red]An unexpected error occurred: {e}[/bold red]", flush=True)
raise e
if __name__ == "__main__":
app()
kronos-venv/lib/python3.12/site-packages/huggingface_hub/inference/_mcp/constants.py
+80
View File
@@ -0,0 +1,80 @@
from __future__ import annotations
import sys
from pathlib import Path
from typing import List
from huggingface_hub import ChatCompletionInputTool
FILENAME_CONFIG = "agent.json"
FILENAME_PROMPT = "PROMPT.md"
DEFAULT_AGENT = {
"model": "Qwen/Qwen2.5-72B-Instruct",
"provider": "nebius",
"servers": [
{
"type": "stdio",
"config": {
"command": "npx",
"args": [
"-y",
"@modelcontextprotocol/server-filesystem",
str(Path.home() / ("Desktop" if sys.platform == "darwin" else "")),
],
},
},
{
"type": "stdio",
"config": {
"command": "npx",
"args": ["@playwright/mcp@latest"],
},
},
],
}
DEFAULT_SYSTEM_PROMPT = """
You are an agent - please keep going until the user’s query is completely
resolved, before ending your turn and yielding back to the user. Only terminate
your turn when you are sure that the problem is solved, or if you need more
info from the user to solve the problem.
If you are not sure about anything pertaining to the user’s request, use your
tools to read files and gather the relevant information: do NOT guess or make
up an answer.
You MUST plan extensively before each function call, and reflect extensively
on the outcomes of the previous function calls. DO NOT do this entire process
by making function calls only, as this can impair your ability to solve the
problem and think insightfully.
""".strip()
MAX_NUM_TURNS = 10
TASK_COMPLETE_TOOL: ChatCompletionInputTool = ChatCompletionInputTool.parse_obj( # type: ignore[assignment]
{
"type": "function",
"function": {
"name": "task_complete",
"description": "Call this tool when the task given by the user is complete",
"parameters": {"type": "object", "properties": {}},
},
}
)
ASK_QUESTION_TOOL: ChatCompletionInputTool = ChatCompletionInputTool.parse_obj( # type: ignore[assignment]
{
"type": "function",
"function": {
"name": "ask_question",
"description": "Ask the user for more info required to solve or clarify their problem.",
"parameters": {"type": "object", "properties": {}},
},
}
)
EXIT_LOOP_TOOLS: List[ChatCompletionInputTool] = [TASK_COMPLETE_TOOL, ASK_QUESTION_TOOL]
DEFAULT_REPO_ID = "tiny-agents/tiny-agents"
kronos-venv/lib/python3.12/site-packages/huggingface_hub/inference/_mcp/mcp_client.py
+369
View File
@@ -0,0 +1,369 @@
import json
import logging
from contextlib import AsyncExitStack
from datetime import timedelta
from pathlib import Path
from typing import TYPE_CHECKING, Any, AsyncIterable, Dict, List, Literal, Optional, Union, overload
from typing_extensions import NotRequired, TypeAlias, TypedDict, Unpack
from ...utils._runtime import get_hf_hub_version
from .._generated._async_client import AsyncInferenceClient
from .._generated.types import (
ChatCompletionInputMessage,
ChatCompletionInputTool,
ChatCompletionStreamOutput,
ChatCompletionStreamOutputDeltaToolCall,
)
from .._providers import PROVIDER_OR_POLICY_T
from .utils import format_result
if TYPE_CHECKING:
from mcp import ClientSession
logger = logging.getLogger(__name__)
# Type alias for tool names
ToolName: TypeAlias = str
ServerType: TypeAlias = Literal["stdio", "sse", "http"]
class StdioServerParameters_T(TypedDict):
command: str
args: NotRequired[List[str]]
env: NotRequired[Dict[str, str]]
cwd: NotRequired[Union[str, Path, None]]
class SSEServerParameters_T(TypedDict):
url: str
headers: NotRequired[Dict[str, Any]]
timeout: NotRequired[float]
sse_read_timeout: NotRequired[float]
class StreamableHTTPParameters_T(TypedDict):
url: str
headers: NotRequired[dict[str, Any]]
timeout: NotRequired[timedelta]
sse_read_timeout: NotRequired[timedelta]
terminate_on_close: NotRequired[bool]
class MCPClient:
"""
Client for connecting to one or more MCP servers and processing chat completions with tools.
<Tip warning={true}>
This class is experimental and might be subject to breaking changes in the future without prior notice.
</Tip>
Args:
model (`str`, `optional`):
The model to run inference with. Can be a model id hosted on the Hugging Face Hub, e.g. `meta-llama/Meta-Llama-3-8B-Instruct`
or a URL to a deployed Inference Endpoint or other local or remote endpoint.
provider (`str`, *optional*):
Name of the provider to use for inference. Defaults to "auto" i.e. the first of the providers available for the model, sorted by the user's order in https://hf.co/settings/inference-providers.
If model is a URL or `base_url` is passed, then `provider` is not used.
base_url (`str`, *optional*):
The base URL to run inference. Defaults to None.
api_key (`str`, `optional`):
Token to use for authentication. Will default to the locally Hugging Face saved token if not provided. You can also use your own provider API key to interact directly with the provider's service.
"""
def __init__(
self,
*,
model: Optional[str] = None,
provider: Optional[PROVIDER_OR_POLICY_T] = None,
base_url: Optional[str] = None,
api_key: Optional[str] = None,
):
# Initialize MCP sessions as a dictionary of ClientSession objects
self.sessions: Dict[ToolName, "ClientSession"] = {}
self.exit_stack = AsyncExitStack()
self.available_tools: List[ChatCompletionInputTool] = []
# To be able to send the model in the payload if `base_url` is provided
if model is None and base_url is None:
raise ValueError("At least one of `model` or `base_url` should be set in `MCPClient`.")
self.payload_model = model
self.client = AsyncInferenceClient(
model=None if base_url is not None else model,
provider=provider,
api_key=api_key,
base_url=base_url,
)
async def __aenter__(self):
"""Enter the context manager"""
await self.client.__aenter__()
await self.exit_stack.__aenter__()
return self
async def __aexit__(self, exc_type, exc_val, exc_tb):
"""Exit the context manager"""
await self.client.__aexit__(exc_type, exc_val, exc_tb)
await self.cleanup()
async def cleanup(self):
"""Clean up resources"""
await self.client.close()
await self.exit_stack.aclose()
@overload
async def add_mcp_server(self, type: Literal["stdio"], **params: Unpack[StdioServerParameters_T]): ...
@overload
async def add_mcp_server(self, type: Literal["sse"], **params: Unpack[SSEServerParameters_T]): ...
@overload
async def add_mcp_server(self, type: Literal["http"], **params: Unpack[StreamableHTTPParameters_T]): ...
async def add_mcp_server(self, type: ServerType, **params: Any):
"""Connect to an MCP server
Args:
type (`str`):
Type of the server to connect to. Can be one of:
- "stdio": Standard input/output server (local)
- "sse": Server-sent events (SSE) server
- "http": StreamableHTTP server
**params (`Dict[str, Any]`):
Server parameters that can be either:
- For stdio servers:
- command (str): The command to run the MCP server
- args (List[str], optional): Arguments for the command
- env (Dict[str, str], optional): Environment variables for the command
- cwd (Union[str, Path, None], optional): Working directory for the command
- For SSE servers:
- url (str): The URL of the SSE server
- headers (Dict[str, Any], optional): Headers for the SSE connection
- timeout (float, optional): Connection timeout
- sse_read_timeout (float, optional): SSE read timeout
- For StreamableHTTP servers:
- url (str): The URL of the StreamableHTTP server
- headers (Dict[str, Any], optional): Headers for the StreamableHTTP connection
- timeout (timedelta, optional): Connection timeout
- sse_read_timeout (timedelta, optional): SSE read timeout
- terminate_on_close (bool, optional): Whether to terminate on close
"""
from mcp import ClientSession, StdioServerParameters
from mcp import types as mcp_types
# Determine server type and create appropriate parameters
if type == "stdio":
# Handle stdio server
from mcp.client.stdio import stdio_client
logger.info(f"Connecting to stdio MCP server with command: {params['command']} {params.get('args', [])}")
client_kwargs = {"command": params["command"]}
for key in ["args", "env", "cwd"]:
if params.get(key) is not None:
client_kwargs[key] = params[key]
server_params = StdioServerParameters(**client_kwargs)
read, write = await self.exit_stack.enter_async_context(stdio_client(server_params))
elif type == "sse":
# Handle SSE server
from mcp.client.sse import sse_client
logger.info(f"Connecting to SSE MCP server at: {params['url']}")
client_kwargs = {"url": params["url"]}
for key in ["headers", "timeout", "sse_read_timeout"]:
if params.get(key) is not None:
client_kwargs[key] = params[key]
read, write = await self.exit_stack.enter_async_context(sse_client(**client_kwargs))
elif type == "http":
# Handle StreamableHTTP server
from mcp.client.streamable_http import streamablehttp_client
logger.info(f"Connecting to StreamableHTTP MCP server at: {params['url']}")
client_kwargs = {"url": params["url"]}
for key in ["headers", "timeout", "sse_read_timeout", "terminate_on_close"]:
if params.get(key) is not None:
client_kwargs[key] = params[key]
read, write, _ = await self.exit_stack.enter_async_context(streamablehttp_client(**client_kwargs))
# ^ TODO: should be handle `get_session_id_callback`? (function to retrieve the current session ID)
else:
raise ValueError(f"Unsupported server type: {type}")
session = await self.exit_stack.enter_async_context(
ClientSession(
read_stream=read,
write_stream=write,
client_info=mcp_types.Implementation(
name="huggingface_hub.MCPClient",
version=get_hf_hub_version(),
),
)
)
logger.debug("Initializing session...")
await session.initialize()
# List available tools
response = await session.list_tools()
logger.debug("Connected to server with tools:", [tool.name for tool in response.tools])
for tool in response.tools:
if tool.name in self.sessions:
logger.warning(f"Tool '{tool.name}' already defined by another server. Skipping.")
continue
# Map tool names to their server for later lookup
self.sessions[tool.name] = session
# Add tool to the list of available tools (for use in chat completions)
self.available_tools.append(
ChatCompletionInputTool.parse_obj_as_instance(
{
"type": "function",
"function": {
"name": tool.name,
"description": tool.description,
"parameters": tool.inputSchema,
},
}
)
)
async def process_single_turn_with_tools(
self,
messages: List[Union[Dict, ChatCompletionInputMessage]],
exit_loop_tools: Optional[List[ChatCompletionInputTool]] = None,
exit_if_first_chunk_no_tool: bool = False,
) -> AsyncIterable[Union[ChatCompletionStreamOutput, ChatCompletionInputMessage]]:
"""Process a query using `self.model` and available tools, yielding chunks and tool outputs.
Args:
messages (`List[Dict]`):
List of message objects representing the conversation history
exit_loop_tools (`List[ChatCompletionInputTool]`, *optional*):
List of tools that should exit the generator when called
exit_if_first_chunk_no_tool (`bool`, *optional*):
Exit if no tool is present in the first chunks. Default to False.
Yields:
[`ChatCompletionStreamOutput`] chunks or [`ChatCompletionInputMessage`] objects
"""
# Prepare tools list based on options
tools = self.available_tools
if exit_loop_tools is not None:
tools = [*exit_loop_tools, *self.available_tools]
# Create the streaming request
response = await self.client.chat.completions.create(
model=self.payload_model,
messages=messages,
tools=tools,
tool_choice="auto",
stream=True,
)
message: Dict[str, Any] = {"role": "unknown", "content": ""}
final_tool_calls: Dict[int, ChatCompletionStreamOutputDeltaToolCall] = {}
num_of_chunks = 0
# Read from stream
async for chunk in response:
num_of_chunks += 1
delta = chunk.choices[0].delta if chunk.choices and len(chunk.choices) > 0 else None
if not delta:
continue
# Process message
if delta.role:
message["role"] = delta.role
if delta.content:
message["content"] += delta.content
# Process tool calls
if delta.tool_calls:
for tool_call in delta.tool_calls:
# Aggregate chunks into tool calls
if tool_call.index not in final_tool_calls:
if (
tool_call.function.arguments is None or tool_call.function.arguments == "{}"
): # Corner case (depends on provider)
tool_call.function.arguments = ""
final_tool_calls[tool_call.index] = tool_call
elif tool_call.function.arguments:
final_tool_calls[tool_call.index].function.arguments += tool_call.function.arguments
# Optionally exit early if no tools in first chunks
if exit_if_first_chunk_no_tool and num_of_chunks <= 2 and len(final_tool_calls) == 0:
return
# Yield each chunk to caller
yield chunk
# Add the assistant message with tool calls (if any) to messages
if message["content"] or final_tool_calls:
# if the role is unknown, set it to assistant
if message.get("role") == "unknown":
message["role"] = "assistant"
# Convert final_tool_calls to the format expected by OpenAI
if final_tool_calls:
tool_calls_list: List[Dict[str, Any]] = []
for tc in final_tool_calls.values():
tool_calls_list.append(
{
"id": tc.id,
"type": "function",
"function": {
"name": tc.function.name,
"arguments": tc.function.arguments or "{}",
},
}
)
message["tool_calls"] = tool_calls_list
messages.append(message)
# Process tool calls one by one
for tool_call in final_tool_calls.values():
function_name = tool_call.function.name
try:
function_args = json.loads(tool_call.function.arguments or "{}")
except json.JSONDecodeError as err:
tool_message = {
"role": "tool",
"tool_call_id": tool_call.id,
"name": function_name,
"content": f"Invalid JSON generated by the model: {err}",
}
tool_message_as_obj = ChatCompletionInputMessage.parse_obj_as_instance(tool_message)
messages.append(tool_message_as_obj)
yield tool_message_as_obj
continue # move to next tool call
tool_message = {"role": "tool", "tool_call_id": tool_call.id, "content": "", "name": function_name}
# Check if this is an exit loop tool
if exit_loop_tools and function_name in [t.function.name for t in exit_loop_tools]:
tool_message_as_obj = ChatCompletionInputMessage.parse_obj_as_instance(tool_message)
messages.append(tool_message_as_obj)
yield tool_message_as_obj
return
# Execute tool call with the appropriate session
session = self.sessions.get(function_name)
if session is not None:
try:
result = await session.call_tool(function_name, function_args)
tool_message["content"] = format_result(result)
except Exception as err:
tool_message["content"] = f"Error: MCP tool call failed with error message: {err}"
else:
tool_message["content"] = f"Error: No session found for tool: {function_name}"
# Yield tool message
tool_message_as_obj = ChatCompletionInputMessage.parse_obj_as_instance(tool_message)
messages.append(tool_message_as_obj)
yield tool_message_as_obj
kronos-venv/lib/python3.12/site-packages/huggingface_hub/inference/_mcp/types.py
+65
View File
@@ -0,0 +1,65 @@
from typing import Dict, List, Literal, TypedDict, Union
# Input config
class InputConfig(TypedDict, total=False):
id: str
description: str
type: str
password: bool
# stdio server config
class StdioServerConfig(TypedDict, total=False):
command: str
args: List[str]
env: Dict[str, str]
cwd: str
class StdioServer(TypedDict):
type: Literal["stdio"]
config: StdioServerConfig
# http server config
class HTTPRequestInit(TypedDict, total=False):
headers: Dict[str, str]
class HTTPServerOptions(TypedDict, total=False):
requestInit: HTTPRequestInit
sessionId: str
class HTTPServerConfig(TypedDict, total=False):
url: str
options: HTTPServerOptions
class HTTPServer(TypedDict):
type: Literal["http"]
config: HTTPServerConfig
# sse server config
class SSEServerOptions(TypedDict, total=False):
requestInit: HTTPRequestInit
class SSEServerConfig(TypedDict):
url: str
options: SSEServerOptions
class SSEServer(TypedDict):
type: Literal["sse"]
config: SSEServerConfig
# AgentConfig root object
class AgentConfig(TypedDict):
model: str
provider: str
inputs: List[InputConfig]
servers: List[Union[StdioServer, HTTPServer, SSEServer]]
kronos-venv/lib/python3.12/site-packages/huggingface_hub/inference/_mcp/utils.py
+124
View File
@@ -0,0 +1,124 @@
"""
Utility functions for MCPClient and Tiny Agents.
Formatting utilities taken from the JS SDK: https://github.com/huggingface/huggingface.js/blob/main/packages/mcp-client/src/ResultFormatter.ts.
"""
import json
from pathlib import Path
from typing import TYPE_CHECKING, List, Optional, Tuple
from huggingface_hub import snapshot_download
from huggingface_hub.errors import EntryNotFoundError
from .constants import DEFAULT_AGENT, DEFAULT_REPO_ID, FILENAME_CONFIG, FILENAME_PROMPT
from .types import AgentConfig
if TYPE_CHECKING:
from mcp import types as mcp_types
def format_result(result: "mcp_types.CallToolResult") -> str:
"""
Formats a mcp.types.CallToolResult content into a human-readable string.
Args:
result (CallToolResult)
Object returned by mcp.ClientSession.call_tool.
Returns:
str
A formatted string representing the content of the result.
"""
content = result.content
if len(content) == 0:
return "[No content]"
formatted_parts: List[str] = []
for item in content:
if item.type == "text":
formatted_parts.append(item.text)
elif item.type == "image":
formatted_parts.append(
f"[Binary Content: Image {item.mimeType}, {_get_base64_size(item.data)} bytes]\n"
f"The task is complete and the content accessible to the User"
)
elif item.type == "audio":
formatted_parts.append(
f"[Binary Content: Audio {item.mimeType}, {_get_base64_size(item.data)} bytes]\n"
f"The task is complete and the content accessible to the User"
)
elif item.type == "resource":
resource = item.resource
if hasattr(resource, "text"):
formatted_parts.append(resource.text)
elif hasattr(resource, "blob"):
formatted_parts.append(
f"[Binary Content ({resource.uri}): {resource.mimeType}, {_get_base64_size(resource.blob)} bytes]\n"
f"The task is complete and the content accessible to the User"
)
return "\n".join(formatted_parts)
def _get_base64_size(base64_str: str) -> int:
"""Estimate the byte size of a base64-encoded string."""
# Remove any prefix like "data:image/png;base64,"
if "," in base64_str:
base64_str = base64_str.split(",")[1]
padding = 0
if base64_str.endswith("=="):
padding = 2
elif base64_str.endswith("="):
padding = 1
return (len(base64_str) * 3) // 4 - padding
def _load_agent_config(agent_path: Optional[str]) -> Tuple[AgentConfig, Optional[str]]:
"""Load server config and prompt."""
def _read_dir(directory: Path) -> Tuple[AgentConfig, Optional[str]]:
cfg_file = directory / FILENAME_CONFIG
if not cfg_file.exists():
raise FileNotFoundError(f" Config file not found in {directory}! Please make sure it exists locally")
config: AgentConfig = json.loads(cfg_file.read_text(encoding="utf-8"))
prompt_file = directory / FILENAME_PROMPT
prompt: Optional[str] = prompt_file.read_text(encoding="utf-8") if prompt_file.exists() else None
return config, prompt
if agent_path is None:
return DEFAULT_AGENT, None # type: ignore[return-value]
path = Path(agent_path).expanduser()
if path.is_file():
return json.loads(path.read_text(encoding="utf-8")), None
if path.is_dir():
return _read_dir(path)
# fetch from the Hub
try:
repo_dir = Path(
snapshot_download(
repo_id=DEFAULT_REPO_ID,
allow_patterns=f"{agent_path}/*",
repo_type="dataset",
)
)
return _read_dir(repo_dir / agent_path)
except Exception as err:
raise EntryNotFoundError(
f" Agent {agent_path} not found in tiny-agents/tiny-agents! Please make sure it exists in https://huggingface.co/datasets/tiny-agents/tiny-agents."
) from err
kronos-venv/lib/python3.12/site-packages/huggingface_hub/inference/_providers/__init__.py
+205
View File
@@ -0,0 +1,205 @@
from typing import Dict, Literal, Optional, Union
from huggingface_hub.inference._providers.featherless_ai import (
FeatherlessConversationalTask,
FeatherlessTextGenerationTask,
)
from huggingface_hub.utils import logging
from ._common import TaskProviderHelper, _fetch_inference_provider_mapping
from .black_forest_labs import BlackForestLabsTextToImageTask
from .cerebras import CerebrasConversationalTask
from .cohere import CohereConversationalTask
from .fal_ai import (
FalAIAutomaticSpeechRecognitionTask,
FalAITextToImageTask,
FalAITextToSpeechTask,
FalAITextToVideoTask,
)
from .fireworks_ai import FireworksAIConversationalTask
from .groq import GroqConversationalTask
from .hf_inference import (
HFInferenceBinaryInputTask,
HFInferenceConversational,
HFInferenceFeatureExtractionTask,
HFInferenceTask,
)
from .hyperbolic import HyperbolicTextGenerationTask, HyperbolicTextToImageTask
from .nebius import (
NebiusConversationalTask,
NebiusFeatureExtractionTask,
NebiusTextGenerationTask,
NebiusTextToImageTask,
)
from .novita import NovitaConversationalTask, NovitaTextGenerationTask, NovitaTextToVideoTask
from .nscale import NscaleConversationalTask, NscaleTextToImageTask
from .openai import OpenAIConversationalTask
from .replicate import ReplicateTask, ReplicateTextToImageTask, ReplicateTextToSpeechTask
from .sambanova import SambanovaConversationalTask, SambanovaFeatureExtractionTask
from .together import TogetherConversationalTask, TogetherTextGenerationTask, TogetherTextToImageTask
logger = logging.get_logger(__name__)
PROVIDER_T = Literal[
"black-forest-labs",
"cerebras",
"cohere",
"fal-ai",
"featherless-ai",
"fireworks-ai",
"groq",
"hf-inference",
"hyperbolic",
"nebius",
"novita",
"nscale",
"openai",
"replicate",
"sambanova",
"together",
]
PROVIDER_OR_POLICY_T = Union[PROVIDER_T, Literal["auto"]]
PROVIDERS: Dict[PROVIDER_T, Dict[str, TaskProviderHelper]] = {
"black-forest-labs": {
"text-to-image": BlackForestLabsTextToImageTask(),
},
"cerebras": {
"conversational": CerebrasConversationalTask(),
},
"cohere": {
"conversational": CohereConversationalTask(),
},
"fal-ai": {
"automatic-speech-recognition": FalAIAutomaticSpeechRecognitionTask(),
"text-to-image": FalAITextToImageTask(),
"text-to-speech": FalAITextToSpeechTask(),
"text-to-video": FalAITextToVideoTask(),
},
"featherless-ai": {
"conversational": FeatherlessConversationalTask(),
"text-generation": FeatherlessTextGenerationTask(),
},
"fireworks-ai": {
"conversational": FireworksAIConversationalTask(),
},
"groq": {
"conversational": GroqConversationalTask(),
},
"hf-inference": {
"text-to-image": HFInferenceTask("text-to-image"),
"conversational": HFInferenceConversational(),
"text-generation": HFInferenceTask("text-generation"),
"text-classification": HFInferenceTask("text-classification"),
"question-answering": HFInferenceTask("question-answering"),
"audio-classification": HFInferenceBinaryInputTask("audio-classification"),
"automatic-speech-recognition": HFInferenceBinaryInputTask("automatic-speech-recognition"),
"fill-mask": HFInferenceTask("fill-mask"),
"feature-extraction": HFInferenceFeatureExtractionTask(),
"image-classification": HFInferenceBinaryInputTask("image-classification"),
"image-segmentation": HFInferenceBinaryInputTask("image-segmentation"),
"document-question-answering": HFInferenceTask("document-question-answering"),
"image-to-text": HFInferenceBinaryInputTask("image-to-text"),
"object-detection": HFInferenceBinaryInputTask("object-detection"),
"audio-to-audio": HFInferenceBinaryInputTask("audio-to-audio"),
"zero-shot-image-classification": HFInferenceBinaryInputTask("zero-shot-image-classification"),
"zero-shot-classification": HFInferenceTask("zero-shot-classification"),
"image-to-image": HFInferenceBinaryInputTask("image-to-image"),
"sentence-similarity": HFInferenceTask("sentence-similarity"),
"table-question-answering": HFInferenceTask("table-question-answering"),
"tabular-classification": HFInferenceTask("tabular-classification"),
"text-to-speech": HFInferenceTask("text-to-speech"),
"token-classification": HFInferenceTask("token-classification"),
"translation": HFInferenceTask("translation"),
"summarization": HFInferenceTask("summarization"),
"visual-question-answering": HFInferenceBinaryInputTask("visual-question-answering"),
},
"hyperbolic": {
"text-to-image": HyperbolicTextToImageTask(),
"conversational": HyperbolicTextGenerationTask("conversational"),
"text-generation": HyperbolicTextGenerationTask("text-generation"),
},
"nebius": {
"text-to-image": NebiusTextToImageTask(),
"conversational": NebiusConversationalTask(),
"text-generation": NebiusTextGenerationTask(),
"feature-extraction": NebiusFeatureExtractionTask(),
},
"novita": {
"text-generation": NovitaTextGenerationTask(),
"conversational": NovitaConversationalTask(),
"text-to-video": NovitaTextToVideoTask(),
},
"nscale": {
"conversational": NscaleConversationalTask(),
"text-to-image": NscaleTextToImageTask(),
},
"openai": {
"conversational": OpenAIConversationalTask(),
},
"replicate": {
"text-to-image": ReplicateTextToImageTask(),
"text-to-speech": ReplicateTextToSpeechTask(),
"text-to-video": ReplicateTask("text-to-video"),
},
"sambanova": {
"conversational": SambanovaConversationalTask(),
"feature-extraction": SambanovaFeatureExtractionTask(),
},
"together": {
"text-to-image": TogetherTextToImageTask(),
"conversational": TogetherConversationalTask(),
"text-generation": TogetherTextGenerationTask(),
},
}
def get_provider_helper(
provider: Optional[PROVIDER_OR_POLICY_T], task: str, model: Optional[str]
) -> TaskProviderHelper:
"""Get provider helper instance by name and task.
Args:
provider (`str`, *optional*): name of the provider, or "auto" to automatically select the provider for the model.
task (`str`): Name of the task
model (`str`, *optional*): Name of the model
Returns:
TaskProviderHelper: Helper instance for the specified provider and task
Raises:
ValueError: If provider or task is not supported
"""
if (model is None and provider in (None, "auto")) or (
model is not None and model.startswith(("http://", "https://"))
):
provider = "hf-inference"
if provider is None:
logger.info(
"Defaulting to 'auto' which will select the first provider available for the model, sorted by the user's order in https://hf.co/settings/inference-providers."
)
provider = "auto"
if provider == "auto":
if model is None:
raise ValueError("Specifying a model is required when provider is 'auto'")
provider_mapping = _fetch_inference_provider_mapping(model)
provider = next(iter(provider_mapping)).provider
provider_tasks = PROVIDERS.get(provider) # type: ignore
if provider_tasks is None:
raise ValueError(
f"Provider '{provider}' not supported. Available values: 'auto' or any provider from {list(PROVIDERS.keys())}."
"Passing 'auto' (default value) will automatically select the first provider available for the model, sorted "
"by the user's order in https://hf.co/settings/inference-providers."
)
if task not in provider_tasks:
raise ValueError(
f"Task '{task}' not supported for provider '{provider}'. Available tasks: {list(provider_tasks.keys())}"
)
return provider_tasks[task]
kronos-venv/lib/python3.12/site-packages/huggingface_hub/inference/_providers/_common.py
+299
View File
@@ -0,0 +1,299 @@
from functools import lru_cache
from typing import Any, Dict, List, Optional, Union, overload
from huggingface_hub import constants
from huggingface_hub.hf_api import InferenceProviderMapping
from huggingface_hub.inference._common import RequestParameters
from huggingface_hub.inference._generated.types.chat_completion import ChatCompletionInputMessage
from huggingface_hub.utils import build_hf_headers, get_token, logging
logger = logging.get_logger(__name__)
# Dev purposes only.
# If you want to try to run inference for a new model locally before it's registered on huggingface.co
# for a given Inference Provider, you can add it to the following dictionary.
HARDCODED_MODEL_INFERENCE_MAPPING: Dict[str, Dict[str, InferenceProviderMapping]] = {
# "HF model ID" => InferenceProviderMapping object initialized with "Model ID on Inference Provider's side"
#
# Example:
# "Qwen/Qwen2.5-Coder-32B-Instruct": InferenceProviderMapping(hf_model_id="Qwen/Qwen2.5-Coder-32B-Instruct",
# provider_id="Qwen2.5-Coder-32B-Instruct",
# task="conversational",
# status="live")
"cerebras": {},
"cohere": {},
"fal-ai": {},
"fireworks-ai": {},
"groq": {},
"hf-inference": {},
"hyperbolic": {},
"nebius": {},
"nscale": {},
"replicate": {},
"sambanova": {},
"together": {},
}
@overload
def filter_none(obj: Dict[str, Any]) -> Dict[str, Any]: ...
@overload
def filter_none(obj: List[Any]) -> List[Any]: ...
def filter_none(obj: Union[Dict[str, Any], List[Any]]) -> Union[Dict[str, Any], List[Any]]:
if isinstance(obj, dict):
cleaned: Dict[str, Any] = {}
for k, v in obj.items():
if v is None:
continue
if isinstance(v, (dict, list)):
v = filter_none(v)
# remove empty nested dicts
if isinstance(v, dict) and not v:
continue
cleaned[k] = v
return cleaned
if isinstance(obj, list):
return [filter_none(v) if isinstance(v, (dict, list)) else v for v in obj]
raise ValueError(f"Expected dict or list, got {type(obj)}")
class TaskProviderHelper:
"""Base class for task-specific provider helpers."""
def __init__(self, provider: str, base_url: str, task: str) -> None:
self.provider = provider
self.task = task
self.base_url = base_url
def prepare_request(
self,
*,
inputs: Any,
parameters: Dict[str, Any],
headers: Dict,
model: Optional[str],
api_key: Optional[str],
extra_payload: Optional[Dict[str, Any]] = None,
) -> RequestParameters:
"""
Prepare the request to be sent to the provider.
Each step (api_key, model, headers, url, payload) can be customized in subclasses.
"""
# api_key from user, or local token, or raise error
api_key = self._prepare_api_key(api_key)
# mapped model from HF model ID
provider_mapping_info = self._prepare_mapping_info(model)
# default HF headers + user headers (to customize in subclasses)
headers = self._prepare_headers(headers, api_key)
# routed URL if HF token, or direct URL (to customize in '_prepare_route' in subclasses)
url = self._prepare_url(api_key, provider_mapping_info.provider_id)
# prepare payload (to customize in subclasses)
payload = self._prepare_payload_as_dict(inputs, parameters, provider_mapping_info=provider_mapping_info)
if payload is not None:
payload = recursive_merge(payload, extra_payload or {})
# body data (to customize in subclasses)
data = self._prepare_payload_as_bytes(inputs, parameters, provider_mapping_info, extra_payload)
# check if both payload and data are set and return
if payload is not None and data is not None:
raise ValueError("Both payload and data cannot be set in the same request.")
if payload is None and data is None:
raise ValueError("Either payload or data must be set in the request.")
return RequestParameters(
url=url, task=self.task, model=provider_mapping_info.provider_id, json=payload, data=data, headers=headers
)
def get_response(
self,
response: Union[bytes, Dict],
request_params: Optional[RequestParameters] = None,
) -> Any:
"""
Return the response in the expected format.
Override this method in subclasses for customized response handling."""
return response
def _prepare_api_key(self, api_key: Optional[str]) -> str:
"""Return the API key to use for the request.
Usually not overwritten in subclasses."""
if api_key is None:
api_key = get_token()
if api_key is None:
raise ValueError(
f"You must provide an api_key to work with {self.provider} API or log in with `huggingface-cli login`."
)
return api_key
def _prepare_mapping_info(self, model: Optional[str]) -> InferenceProviderMapping:
"""Return the mapped model ID to use for the request.
Usually not overwritten in subclasses."""
if model is None:
raise ValueError(f"Please provide an HF model ID supported by {self.provider}.")
# hardcoded mapping for local testing
if HARDCODED_MODEL_INFERENCE_MAPPING.get(self.provider, {}).get(model):
return HARDCODED_MODEL_INFERENCE_MAPPING[self.provider][model]
provider_mapping = None
for mapping in _fetch_inference_provider_mapping(model):
if mapping.provider == self.provider:
provider_mapping = mapping
break
if provider_mapping is None:
raise ValueError(f"Model {model} is not supported by provider {self.provider}.")
if provider_mapping.task != self.task:
raise ValueError(
f"Model {model} is not supported for task {self.task} and provider {self.provider}. "
f"Supported task: {provider_mapping.task}."
)
if provider_mapping.status == "staging":
logger.warning(
f"Model {model} is in staging mode for provider {self.provider}. Meant for test purposes only."
)
if provider_mapping.status == "error":
logger.warning(
f"Our latest automated health check on model '{model}' for provider '{self.provider}' did not complete successfully. "
"Inference call might fail."
)
return provider_mapping
def _prepare_headers(self, headers: Dict, api_key: str) -> Dict:
"""Return the headers to use for the request.
Override this method in subclasses for customized headers.
"""
return {**build_hf_headers(token=api_key), **headers}
def _prepare_url(self, api_key: str, mapped_model: str) -> str:
"""Return the URL to use for the request.
Usually not overwritten in subclasses."""
base_url = self._prepare_base_url(api_key)
route = self._prepare_route(mapped_model, api_key)
return f"{base_url.rstrip('/')}/{route.lstrip('/')}"
def _prepare_base_url(self, api_key: str) -> str:
"""Return the base URL to use for the request.
Usually not overwritten in subclasses."""
# Route to the proxy if the api_key is a HF TOKEN
if api_key.startswith("hf_"):
logger.info(f"Calling '{self.provider}' provider through Hugging Face router.")
return constants.INFERENCE_PROXY_TEMPLATE.format(provider=self.provider)
else:
logger.info(f"Calling '{self.provider}' provider directly.")
return self.base_url
def _prepare_route(self, mapped_model: str, api_key: str) -> str:
"""Return the route to use for the request.
Override this method in subclasses for customized routes.
"""
return ""
def _prepare_payload_as_dict(
self, inputs: Any, parameters: Dict, provider_mapping_info: InferenceProviderMapping
) -> Optional[Dict]:
"""Return the payload to use for the request, as a dict.
Override this method in subclasses for customized payloads.
Only one of `_prepare_payload_as_dict` and `_prepare_payload_as_bytes` should return a value.
"""
return None
def _prepare_payload_as_bytes(
self,
inputs: Any,
parameters: Dict,
provider_mapping_info: InferenceProviderMapping,
extra_payload: Optional[Dict],
) -> Optional[bytes]:
"""Return the body to use for the request, as bytes.
Override this method in subclasses for customized body data.
Only one of `_prepare_payload_as_dict` and `_prepare_payload_as_bytes` should return a value.
"""
return None
class BaseConversationalTask(TaskProviderHelper):
"""
Base class for conversational (chat completion) tasks.
The schema follows the OpenAI API format defined here: https://platform.openai.com/docs/api-reference/chat
"""
def __init__(self, provider: str, base_url: str):
super().__init__(provider=provider, base_url=base_url, task="conversational")
def _prepare_route(self, mapped_model: str, api_key: str) -> str:
return "/v1/chat/completions"
def _prepare_payload_as_dict(
self,
inputs: List[Union[Dict, ChatCompletionInputMessage]],
parameters: Dict,
provider_mapping_info: InferenceProviderMapping,
) -> Optional[Dict]:
return filter_none({"messages": inputs, **parameters, "model": provider_mapping_info.provider_id})
class BaseTextGenerationTask(TaskProviderHelper):
"""
Base class for text-generation (completion) tasks.
The schema follows the OpenAI API format defined here: https://platform.openai.com/docs/api-reference/completions
"""
def __init__(self, provider: str, base_url: str):
super().__init__(provider=provider, base_url=base_url, task="text-generation")
def _prepare_route(self, mapped_model: str, api_key: str) -> str:
return "/v1/completions"
def _prepare_payload_as_dict(
self, inputs: Any, parameters: Dict, provider_mapping_info: InferenceProviderMapping
) -> Optional[Dict]:
return {"prompt": inputs, **filter_none(parameters), "model": provider_mapping_info.provider_id}
@lru_cache(maxsize=None)
def _fetch_inference_provider_mapping(model: str) -> List["InferenceProviderMapping"]:
"""
Fetch provider mappings for a model from the Hub.
"""
from huggingface_hub.hf_api import HfApi
info = HfApi().model_info(model, expand=["inferenceProviderMapping"])
provider_mapping = info.inference_provider_mapping
if provider_mapping is None:
raise ValueError(f"No provider mapping found for model {model}")
return provider_mapping
def recursive_merge(dict1: Dict, dict2: Dict) -> Dict:
return {
**dict1,
**{
key: recursive_merge(dict1[key], value)
if (key in dict1 and isinstance(dict1[key], dict) and isinstance(value, dict))
else value
for key, value in dict2.items()
},
}
kronos-venv/lib/python3.12/site-packages/huggingface_hub/inference/_providers/black_forest_labs.py
+69
View File
@@ -0,0 +1,69 @@
import time
from typing import Any, Dict, Optional, Union
from huggingface_hub.hf_api import InferenceProviderMapping
from huggingface_hub.inference._common import RequestParameters, _as_dict
from huggingface_hub.inference._providers._common import TaskProviderHelper, filter_none
from huggingface_hub.utils import logging
from huggingface_hub.utils._http import get_session
logger = logging.get_logger(__name__)
MAX_POLLING_ATTEMPTS = 6
POLLING_INTERVAL = 1.0
class BlackForestLabsTextToImageTask(TaskProviderHelper):
def __init__(self):
super().__init__(provider="black-forest-labs", base_url="https://api.us1.bfl.ai", task="text-to-image")
def _prepare_headers(self, headers: Dict, api_key: str) -> Dict:
headers = super()._prepare_headers(headers, api_key)
if not api_key.startswith("hf_"):
_ = headers.pop("authorization")
headers["X-Key"] = api_key
return headers
def _prepare_route(self, mapped_model: str, api_key: str) -> str:
return f"/v1/{mapped_model}"
def _prepare_payload_as_dict(
self, inputs: Any, parameters: Dict, provider_mapping_info: InferenceProviderMapping
) -> Optional[Dict]:
parameters = filter_none(parameters)
if "num_inference_steps" in parameters:
parameters["steps"] = parameters.pop("num_inference_steps")
if "guidance_scale" in parameters:
parameters["guidance"] = parameters.pop("guidance_scale")
return {"prompt": inputs, **parameters}
def get_response(self, response: Union[bytes, Dict], request_params: Optional[RequestParameters] = None) -> Any:
"""
Polling mechanism for Black Forest Labs since the API is asynchronous.
"""
url = _as_dict(response).get("polling_url")
session = get_session()
for _ in range(MAX_POLLING_ATTEMPTS):
time.sleep(POLLING_INTERVAL)
response = session.get(url, headers={"Content-Type": "application/json"}) # type: ignore
response.raise_for_status() # type: ignore
response_json: Dict = response.json() # type: ignore
status = response_json.get("status")
logger.info(
f"Polling generation result from {url}. Current status: {status}. "
f"Will retry after {POLLING_INTERVAL} seconds if not ready."
)
if (
status == "Ready"
and isinstance(response_json.get("result"), dict)
and (sample_url := response_json["result"].get("sample"))
):
image_resp = session.get(sample_url)
image_resp.raise_for_status()
return image_resp.content
raise TimeoutError(f"Failed to get the image URL after {MAX_POLLING_ATTEMPTS} attempts.")
kronos-venv/lib/python3.12/site-packages/huggingface_hub/inference/_providers/cerebras.py
+6
View File
@@ -0,0 +1,6 @@
from ._common import BaseConversationalTask
class CerebrasConversationalTask(BaseConversationalTask):
def __init__(self):
super().__init__(provider="cerebras", base_url="https://api.cerebras.ai")
kronos-venv/lib/python3.12/site-packages/huggingface_hub/inference/_providers/cohere.py
+32
View File
@@ -0,0 +1,32 @@
from typing import Any, Dict, Optional
from huggingface_hub.hf_api import InferenceProviderMapping
from ._common import BaseConversationalTask
_PROVIDER = "cohere"
_BASE_URL = "https://api.cohere.com"
class CohereConversationalTask(BaseConversationalTask):
def __init__(self):
super().__init__(provider=_PROVIDER, base_url=_BASE_URL)
def _prepare_route(self, mapped_model: str, api_key: str) -> str:
return "/compatibility/v1/chat/completions"
def _prepare_payload_as_dict(
self, inputs: Any, parameters: Dict, provider_mapping_info: InferenceProviderMapping
) -> Optional[Dict]:
payload = super()._prepare_payload_as_dict(inputs, parameters, provider_mapping_info)
response_format = parameters.get("response_format")
if isinstance(response_format, dict) and response_format.get("type") == "json_schema":
json_schema_details = response_format.get("json_schema")
if isinstance(json_schema_details, dict) and "schema" in json_schema_details:
payload["response_format"] = { # type: ignore [index]
"type": "json_object",
"schema": json_schema_details["schema"],
}
return payload
kronos-venv/lib/python3.12/site-packages/huggingface_hub/inference/_providers/fal_ai.py
+172
View File
@@ -0,0 +1,172 @@
import base64
import time
from abc import ABC
from typing import Any, Dict, Optional, Union
from urllib.parse import urlparse
from huggingface_hub import constants
from huggingface_hub.hf_api import InferenceProviderMapping
from huggingface_hub.inference._common import RequestParameters, _as_dict
from huggingface_hub.inference._providers._common import TaskProviderHelper, filter_none
from huggingface_hub.utils import get_session, hf_raise_for_status
from huggingface_hub.utils.logging import get_logger
logger = get_logger(__name__)
# Arbitrary polling interval
_POLLING_INTERVAL = 0.5
class FalAITask(TaskProviderHelper, ABC):
def __init__(self, task: str):
super().__init__(provider="fal-ai", base_url="https://fal.run", task=task)
def _prepare_headers(self, headers: Dict, api_key: str) -> Dict:
headers = super()._prepare_headers(headers, api_key)
if not api_key.startswith("hf_"):
headers["authorization"] = f"Key {api_key}"
return headers
def _prepare_route(self, mapped_model: str, api_key: str) -> str:
return f"/{mapped_model}"
class FalAIAutomaticSpeechRecognitionTask(FalAITask):
def __init__(self):
super().__init__("automatic-speech-recognition")
def _prepare_payload_as_dict(
self, inputs: Any, parameters: Dict, provider_mapping_info: InferenceProviderMapping
) -> Optional[Dict]:
if isinstance(inputs, str) and inputs.startswith(("http://", "https://")):
# If input is a URL, pass it directly
audio_url = inputs
else:
# If input is a file path, read it first
if isinstance(inputs, str):
with open(inputs, "rb") as f:
inputs = f.read()
audio_b64 = base64.b64encode(inputs).decode()
content_type = "audio/mpeg"
audio_url = f"data:{content_type};base64,{audio_b64}"
return {"audio_url": audio_url, **filter_none(parameters)}
def get_response(self, response: Union[bytes, Dict], request_params: Optional[RequestParameters] = None) -> Any:
text = _as_dict(response)["text"]
if not isinstance(text, str):
raise ValueError(f"Unexpected output format from FalAI API. Expected string, got {type(text)}.")
return text
class FalAITextToImageTask(FalAITask):
def __init__(self):
super().__init__("text-to-image")
def _prepare_payload_as_dict(
self, inputs: Any, parameters: Dict, provider_mapping_info: InferenceProviderMapping
) -> Optional[Dict]:
payload: Dict[str, Any] = {
"prompt": inputs,
**filter_none(parameters),
}
if "width" in payload and "height" in payload:
payload["image_size"] = {
"width": payload.pop("width"),
"height": payload.pop("height"),
}
if provider_mapping_info.adapter_weights_path is not None:
lora_path = constants.HUGGINGFACE_CO_URL_TEMPLATE.format(
repo_id=provider_mapping_info.hf_model_id,
revision="main",
filename=provider_mapping_info.adapter_weights_path,
)
payload["loras"] = [{"path": lora_path, "scale": 1}]
if provider_mapping_info.provider_id == "fal-ai/lora":
# little hack: fal requires the base model for stable-diffusion-based loras but not for flux-based
# See payloads in https://fal.ai/models/fal-ai/lora/api vs https://fal.ai/models/fal-ai/flux-lora/api
payload["model_name"] = "stabilityai/stable-diffusion-xl-base-1.0"
return payload
def get_response(self, response: Union[bytes, Dict], request_params: Optional[RequestParameters] = None) -> Any:
url = _as_dict(response)["images"][0]["url"]
return get_session().get(url).content
class FalAITextToSpeechTask(FalAITask):
def __init__(self):
super().__init__("text-to-speech")
def _prepare_payload_as_dict(
self, inputs: Any, parameters: Dict, provider_mapping_info: InferenceProviderMapping
) -> Optional[Dict]:
return {"text": inputs, **filter_none(parameters)}
def get_response(self, response: Union[bytes, Dict], request_params: Optional[RequestParameters] = None) -> Any:
url = _as_dict(response)["audio"]["url"]
return get_session().get(url).content
class FalAITextToVideoTask(FalAITask):
def __init__(self):
super().__init__("text-to-video")
def _prepare_base_url(self, api_key: str) -> str:
if api_key.startswith("hf_"):
return super()._prepare_base_url(api_key)
else:
logger.info(f"Calling '{self.provider}' provider directly.")
return "https://queue.fal.run"
def _prepare_route(self, mapped_model: str, api_key: str) -> str:
if api_key.startswith("hf_"):
# Use the queue subdomain for HF routing
return f"/{mapped_model}?_subdomain=queue"
return f"/{mapped_model}"
def _prepare_payload_as_dict(
self, inputs: Any, parameters: Dict, provider_mapping_info: InferenceProviderMapping
) -> Optional[Dict]:
return {"prompt": inputs, **filter_none(parameters)}
def get_response(
self,
response: Union[bytes, Dict],
request_params: Optional[RequestParameters] = None,
) -> Any:
response_dict = _as_dict(response)
request_id = response_dict.get("request_id")
if not request_id:
raise ValueError("No request ID found in the response")
if request_params is None:
raise ValueError(
"A `RequestParameters` object should be provided to get text-to-video responses with Fal AI."
)
# extract the base url and query params
parsed_url = urlparse(request_params.url)
# a bit hacky way to concatenate the provider name without parsing `parsed_url.path`
base_url = f"{parsed_url.scheme}://{parsed_url.netloc}{'/fal-ai' if parsed_url.netloc == 'router.huggingface.co' else ''}"
query_param = f"?{parsed_url.query}" if parsed_url.query else ""
# extracting the provider model id for status and result urls
# from the response as it might be different from the mapped model in `request_params.url`
model_id = urlparse(response_dict.get("response_url")).path
status_url = f"{base_url}{str(model_id)}/status{query_param}"
result_url = f"{base_url}{str(model_id)}{query_param}"
status = response_dict.get("status")
logger.info("Generating the video.. this can take several minutes.")
while status != "COMPLETED":
time.sleep(_POLLING_INTERVAL)
status_response = get_session().get(status_url, headers=request_params.headers)
hf_raise_for_status(status_response)
status = status_response.json().get("status")
response = get_session().get(result_url, headers=request_params.headers).json()
url = _as_dict(response)["video"]["url"]
return get_session().get(url).content
kronos-venv/lib/python3.12/site-packages/huggingface_hub/inference/_providers/featherless_ai.py
+38
View File
@@ -0,0 +1,38 @@
from typing import Any, Dict, Optional, Union
from huggingface_hub.hf_api import InferenceProviderMapping
from huggingface_hub.inference._common import RequestParameters, _as_dict
from ._common import BaseConversationalTask, BaseTextGenerationTask, filter_none
_PROVIDER = "featherless-ai"
_BASE_URL = "https://api.featherless.ai"
class FeatherlessTextGenerationTask(BaseTextGenerationTask):
def __init__(self):
super().__init__(provider=_PROVIDER, base_url=_BASE_URL)
def _prepare_payload_as_dict(
self, inputs: Any, parameters: Dict, provider_mapping_info: InferenceProviderMapping
) -> Optional[Dict]:
params = filter_none(parameters.copy())
params["max_tokens"] = params.pop("max_new_tokens", None)
return {"prompt": inputs, **params, "model": provider_mapping_info.provider_id}
def get_response(self, response: Union[bytes, Dict], request_params: Optional[RequestParameters] = None) -> Any:
output = _as_dict(response)["choices"][0]
return {
"generated_text": output["text"],
"details": {
"finish_reason": output.get("finish_reason"),
"seed": output.get("seed"),
},
}
class FeatherlessConversationalTask(BaseConversationalTask):
def __init__(self):
super().__init__(provider=_PROVIDER, base_url=_BASE_URL)
kronos-venv/lib/python3.12/site-packages/huggingface_hub/inference/_providers/fireworks_ai.py
+27
View File
@@ -0,0 +1,27 @@
from typing import Any, Dict, Optional
from huggingface_hub.hf_api import InferenceProviderMapping
from ._common import BaseConversationalTask
class FireworksAIConversationalTask(BaseConversationalTask):
def __init__(self):
super().__init__(provider="fireworks-ai", base_url="https://api.fireworks.ai")
def _prepare_route(self, mapped_model: str, api_key: str) -> str:
return "/inference/v1/chat/completions"
def _prepare_payload_as_dict(
self, inputs: Any, parameters: Dict, provider_mapping_info: InferenceProviderMapping
) -> Optional[Dict]:
payload = super()._prepare_payload_as_dict(inputs, parameters, provider_mapping_info)
response_format = parameters.get("response_format")
if isinstance(response_format, dict) and response_format.get("type") == "json_schema":
json_schema_details = response_format.get("json_schema")
if isinstance(json_schema_details, dict) and "schema" in json_schema_details:
payload["response_format"] = { # type: ignore [index]
"type": "json_object",
"schema": json_schema_details["schema"],
}
return payload
kronos-venv/lib/python3.12/site-packages/huggingface_hub/inference/_providers/groq.py
+9
View File
@@ -0,0 +1,9 @@
from ._common import BaseConversationalTask
class GroqConversationalTask(BaseConversationalTask):
def __init__(self):
super().__init__(provider="groq", base_url="https://api.groq.com")
def _prepare_route(self, mapped_model: str, api_key: str) -> str:
return "/openai/v1/chat/completions"
kronos-venv/lib/python3.12/site-packages/huggingface_hub/inference/_providers/hf_inference.py
+212
View File
@@ -0,0 +1,212 @@
import json
from functools import lru_cache
from pathlib import Path
from typing import Any, Dict, Optional, Union
from huggingface_hub import constants
from huggingface_hub.hf_api import InferenceProviderMapping
from huggingface_hub.inference._common import RequestParameters, _b64_encode, _bytes_to_dict, _open_as_binary
from huggingface_hub.inference._providers._common import TaskProviderHelper, filter_none
from huggingface_hub.utils import build_hf_headers, get_session, get_token, hf_raise_for_status
class HFInferenceTask(TaskProviderHelper):
"""Base class for HF Inference API tasks."""
def __init__(self, task: str):
super().__init__(
provider="hf-inference",
base_url=constants.INFERENCE_PROXY_TEMPLATE.format(provider="hf-inference"),
task=task,
)
def _prepare_api_key(self, api_key: Optional[str]) -> str:
# special case: for HF Inference we allow not providing an API key
return api_key or get_token() # type: ignore[return-value]
def _prepare_mapping_info(self, model: Optional[str]) -> InferenceProviderMapping:
if model is not None and model.startswith(("http://", "https://")):
return InferenceProviderMapping(
provider="hf-inference", providerId=model, hf_model_id=model, task=self.task, status="live"
)
model_id = model if model is not None else _fetch_recommended_models().get(self.task)
if model_id is None:
raise ValueError(
f"Task {self.task} has no recommended model for HF Inference. Please specify a model"
" explicitly. Visit https://huggingface.co/tasks for more info."
)
_check_supported_task(model_id, self.task)
return InferenceProviderMapping(
provider="hf-inference", providerId=model_id, hf_model_id=model_id, task=self.task, status="live"
)
def _prepare_url(self, api_key: str, mapped_model: str) -> str:
# hf-inference provider can handle URLs (e.g. Inference Endpoints or TGI deployment)
if mapped_model.startswith(("http://", "https://")):
return mapped_model
return (
# Feature-extraction and sentence-similarity are the only cases where we handle models with several tasks.
f"{self.base_url}/models/{mapped_model}/pipeline/{self.task}"
if self.task in ("feature-extraction", "sentence-similarity")
# Otherwise, we use the default endpoint
else f"{self.base_url}/models/{mapped_model}"
)
def _prepare_payload_as_dict(
self, inputs: Any, parameters: Dict, provider_mapping_info: InferenceProviderMapping
) -> Optional[Dict]:
if isinstance(inputs, bytes):
raise ValueError(f"Unexpected binary input for task {self.task}.")
if isinstance(inputs, Path):
raise ValueError(f"Unexpected path input for task {self.task} (got {inputs})")
return {"inputs": inputs, "parameters": filter_none(parameters)}
class HFInferenceBinaryInputTask(HFInferenceTask):
def _prepare_payload_as_dict(
self, inputs: Any, parameters: Dict, provider_mapping_info: InferenceProviderMapping
) -> Optional[Dict]:
return None
def _prepare_payload_as_bytes(
self,
inputs: Any,
parameters: Dict,
provider_mapping_info: InferenceProviderMapping,
extra_payload: Optional[Dict],
) -> Optional[bytes]:
parameters = filter_none(parameters)
extra_payload = extra_payload or {}
has_parameters = len(parameters) > 0 or len(extra_payload) > 0
# Raise if not a binary object or a local path or a URL.
if not isinstance(inputs, (bytes, Path)) and not isinstance(inputs, str):
raise ValueError(f"Expected binary inputs or a local path or a URL. Got {inputs}")
# Send inputs as raw content when no parameters are provided
if not has_parameters:
with _open_as_binary(inputs) as data:
data_as_bytes = data if isinstance(data, bytes) else data.read()
return data_as_bytes
# Otherwise encode as b64
return json.dumps({"inputs": _b64_encode(inputs), "parameters": parameters, **extra_payload}).encode("utf-8")
class HFInferenceConversational(HFInferenceTask):
def __init__(self):
super().__init__("conversational")
def _prepare_payload_as_dict(
self, inputs: Any, parameters: Dict, provider_mapping_info: InferenceProviderMapping
) -> Optional[Dict]:
payload = filter_none(parameters)
mapped_model = provider_mapping_info.provider_id
payload_model = parameters.get("model") or mapped_model
if payload_model is None or payload_model.startswith(("http://", "https://")):
payload_model = "dummy"
response_format = parameters.get("response_format")
if isinstance(response_format, dict) and response_format.get("type") == "json_schema":
payload["response_format"] = {
"type": "json_object",
"value": response_format["json_schema"]["schema"],
}
return {**payload, "model": payload_model, "messages": inputs}
def _prepare_url(self, api_key: str, mapped_model: str) -> str:
base_url = (
mapped_model
if mapped_model.startswith(("http://", "https://"))
else f"{constants.INFERENCE_PROXY_TEMPLATE.format(provider='hf-inference')}/models/{mapped_model}"
)
return _build_chat_completion_url(base_url)
def _build_chat_completion_url(model_url: str) -> str:
# Strip trailing /
model_url = model_url.rstrip("/")
# Append /chat/completions if not already present
if model_url.endswith("/v1"):
model_url += "/chat/completions"
# Append /v1/chat/completions if not already present
if not model_url.endswith("/chat/completions"):
model_url += "/v1/chat/completions"
return model_url
@lru_cache(maxsize=1)
def _fetch_recommended_models() -> Dict[str, Optional[str]]:
response = get_session().get(f"{constants.ENDPOINT}/api/tasks", headers=build_hf_headers())
hf_raise_for_status(response)
return {task: next(iter(details["widgetModels"]), None) for task, details in response.json().items()}
@lru_cache(maxsize=None)
def _check_supported_task(model: str, task: str) -> None:
from huggingface_hub.hf_api import HfApi
model_info = HfApi().model_info(model)
pipeline_tag = model_info.pipeline_tag
tags = model_info.tags or []
is_conversational = "conversational" in tags
if task in ("text-generation", "conversational"):
if pipeline_tag == "text-generation":
# text-generation + conversational tag -> both tasks allowed
if is_conversational:
return
# text-generation without conversational tag -> only text-generation allowed
if task == "text-generation":
return
raise ValueError(f"Model '{model}' doesn't support task '{task}'.")
if pipeline_tag == "text2text-generation":
if task == "text-generation":
return
raise ValueError(f"Model '{model}' doesn't support task '{task}'.")
if pipeline_tag == "image-text-to-text":
if is_conversational and task == "conversational":
return # Only conversational allowed if tagged as conversational
raise ValueError("Non-conversational image-text-to-text task is not supported.")
if (
task in ("feature-extraction", "sentence-similarity")
and pipeline_tag in ("feature-extraction", "sentence-similarity")
and task in tags
):
# feature-extraction and sentence-similarity are interchangeable for HF Inference
return
# For all other tasks, just check pipeline tag
if pipeline_tag != task:
raise ValueError(
f"Model '{model}' doesn't support task '{task}'. Supported tasks: '{pipeline_tag}', got: '{task}'"
)
return
class HFInferenceFeatureExtractionTask(HFInferenceTask):
def __init__(self):
super().__init__("feature-extraction")
def _prepare_payload_as_dict(
self, inputs: Any, parameters: Dict, provider_mapping_info: InferenceProviderMapping
) -> Optional[Dict]:
if isinstance(inputs, bytes):
raise ValueError(f"Unexpected binary input for task {self.task}.")
if isinstance(inputs, Path):
raise ValueError(f"Unexpected path input for task {self.task} (got {inputs})")
# Parameters are sent at root-level for feature-extraction task
# See specs: https://github.com/huggingface/huggingface.js/blob/main/packages/tasks/src/tasks/feature-extraction/spec/input.json
return {"inputs": inputs, **filter_none(parameters)}
def get_response(self, response: Union[bytes, Dict], request_params: Optional[RequestParameters] = None) -> Any:
if isinstance(response, bytes):
return _bytes_to_dict(response)
return response
kronos-venv/lib/python3.12/site-packages/huggingface_hub/inference/_providers/hyperbolic.py
+47
View File
@@ -0,0 +1,47 @@
import base64
from typing import Any, Dict, Optional, Union
from huggingface_hub.hf_api import InferenceProviderMapping
from huggingface_hub.inference._common import RequestParameters, _as_dict
from huggingface_hub.inference._providers._common import BaseConversationalTask, TaskProviderHelper, filter_none
class HyperbolicTextToImageTask(TaskProviderHelper):
def __init__(self):
super().__init__(provider="hyperbolic", base_url="https://api.hyperbolic.xyz", task="text-to-image")
def _prepare_route(self, mapped_model: str, api_key: str) -> str:
return "/v1/images/generations"
def _prepare_payload_as_dict(
self, inputs: Any, parameters: Dict, provider_mapping_info: InferenceProviderMapping
) -> Optional[Dict]:
mapped_model = provider_mapping_info.provider_id
parameters = filter_none(parameters)
if "num_inference_steps" in parameters:
parameters["steps"] = parameters.pop("num_inference_steps")
if "guidance_scale" in parameters:
parameters["cfg_scale"] = parameters.pop("guidance_scale")
# For Hyperbolic, the width and height are required parameters
if "width" not in parameters:
parameters["width"] = 512
if "height" not in parameters:
parameters["height"] = 512
return {"prompt": inputs, "model_name": mapped_model, **parameters}
def get_response(self, response: Union[bytes, Dict], request_params: Optional[RequestParameters] = None) -> Any:
response_dict = _as_dict(response)
return base64.b64decode(response_dict["images"][0]["image"])
class HyperbolicTextGenerationTask(BaseConversationalTask):
"""
Special case for Hyperbolic, where text-generation task is handled as a conversational task.
"""
def __init__(self, task: str):
super().__init__(
provider="hyperbolic",
base_url="https://api.hyperbolic.xyz",
)
self.task = task
kronos-venv/lib/python3.12/site-packages/huggingface_hub/inference/_providers/nebius.py
+83
View File
@@ -0,0 +1,83 @@
import base64
from typing import Any, Dict, Optional, Union
from huggingface_hub.hf_api import InferenceProviderMapping
from huggingface_hub.inference._common import RequestParameters, _as_dict
from huggingface_hub.inference._providers._common import (
BaseConversationalTask,
BaseTextGenerationTask,
TaskProviderHelper,
filter_none,
)
class NebiusTextGenerationTask(BaseTextGenerationTask):
def __init__(self):
super().__init__(provider="nebius", base_url="https://api.studio.nebius.ai")
def get_response(self, response: Union[bytes, Dict], request_params: Optional[RequestParameters] = None) -> Any:
output = _as_dict(response)["choices"][0]
return {
"generated_text": output["text"],
"details": {
"finish_reason": output.get("finish_reason"),
"seed": output.get("seed"),
},
}
class NebiusConversationalTask(BaseConversationalTask):
def __init__(self):
super().__init__(provider="nebius", base_url="https://api.studio.nebius.ai")
def _prepare_payload_as_dict(
self, inputs: Any, parameters: Dict, provider_mapping_info: InferenceProviderMapping
) -> Optional[Dict]:
payload = super()._prepare_payload_as_dict(inputs, parameters, provider_mapping_info)
response_format = parameters.get("response_format")
if isinstance(response_format, dict) and response_format.get("type") == "json_schema":
json_schema_details = response_format.get("json_schema")
if isinstance(json_schema_details, dict) and "schema" in json_schema_details:
payload["guided_json"] = json_schema_details["schema"] # type: ignore [index]
return payload
class NebiusTextToImageTask(TaskProviderHelper):
def __init__(self):
super().__init__(task="text-to-image", provider="nebius", base_url="https://api.studio.nebius.ai")
def _prepare_route(self, mapped_model: str, api_key: str) -> str:
return "/v1/images/generations"
def _prepare_payload_as_dict(
self, inputs: Any, parameters: Dict, provider_mapping_info: InferenceProviderMapping
) -> Optional[Dict]:
mapped_model = provider_mapping_info.provider_id
parameters = filter_none(parameters)
if "guidance_scale" in parameters:
parameters.pop("guidance_scale")
if parameters.get("response_format") not in ("b64_json", "url"):
parameters["response_format"] = "b64_json"
return {"prompt": inputs, **parameters, "model": mapped_model}
def get_response(self, response: Union[bytes, Dict], request_params: Optional[RequestParameters] = None) -> Any:
response_dict = _as_dict(response)
return base64.b64decode(response_dict["data"][0]["b64_json"])
class NebiusFeatureExtractionTask(TaskProviderHelper):
def __init__(self):
super().__init__(task="feature-extraction", provider="nebius", base_url="https://api.studio.nebius.ai")
def _prepare_route(self, mapped_model: str, api_key: str) -> str:
return "/v1/embeddings"
def _prepare_payload_as_dict(
self, inputs: Any, parameters: Dict, provider_mapping_info: InferenceProviderMapping
) -> Optional[Dict]:
return {"input": inputs, "model": provider_mapping_info.provider_id}
def get_response(self, response: Union[bytes, Dict], request_params: Optional[RequestParameters] = None) -> Any:
embeddings = _as_dict(response)["data"]
return [embedding["embedding"] for embedding in embeddings]
kronos-venv/lib/python3.12/site-packages/huggingface_hub/inference/_providers/novita.py
+69
View File
@@ -0,0 +1,69 @@
from typing import Any, Dict, Optional, Union
from huggingface_hub.hf_api import InferenceProviderMapping
from huggingface_hub.inference._common import RequestParameters, _as_dict
from huggingface_hub.inference._providers._common import (
BaseConversationalTask,
BaseTextGenerationTask,
TaskProviderHelper,
filter_none,
)
from huggingface_hub.utils import get_session
_PROVIDER = "novita"
_BASE_URL = "https://api.novita.ai"
class NovitaTextGenerationTask(BaseTextGenerationTask):
def __init__(self):
super().__init__(provider=_PROVIDER, base_url=_BASE_URL)
def _prepare_route(self, mapped_model: str, api_key: str) -> str:
# there is no v1/ route for novita
return "/v3/openai/completions"
def get_response(self, response: Union[bytes, Dict], request_params: Optional[RequestParameters] = None) -> Any:
output = _as_dict(response)["choices"][0]
return {
"generated_text": output["text"],
"details": {
"finish_reason": output.get("finish_reason"),
"seed": output.get("seed"),
},
}
class NovitaConversationalTask(BaseConversationalTask):
def __init__(self):
super().__init__(provider=_PROVIDER, base_url=_BASE_URL)
def _prepare_route(self, mapped_model: str, api_key: str) -> str:
# there is no v1/ route for novita
return "/v3/openai/chat/completions"
class NovitaTextToVideoTask(TaskProviderHelper):
def __init__(self):
super().__init__(provider=_PROVIDER, base_url=_BASE_URL, task="text-to-video")
def _prepare_route(self, mapped_model: str, api_key: str) -> str:
return f"/v3/hf/{mapped_model}"
def _prepare_payload_as_dict(
self, inputs: Any, parameters: Dict, provider_mapping_info: InferenceProviderMapping
) -> Optional[Dict]:
return {"prompt": inputs, **filter_none(parameters)}
def get_response(self, response: Union[bytes, Dict], request_params: Optional[RequestParameters] = None) -> Any:
response_dict = _as_dict(response)
if not (
isinstance(response_dict, dict)
and "video" in response_dict
and isinstance(response_dict["video"], dict)
and "video_url" in response_dict["video"]
):
raise ValueError("Expected response format: { 'video': { 'video_url': string } }")
video_url = response_dict["video"]["video_url"]
return get_session().get(video_url).content
kronos-venv/lib/python3.12/site-packages/huggingface_hub/inference/_providers/nscale.py
+44
View File
@@ -0,0 +1,44 @@
import base64
from typing import Any, Dict, Optional, Union
from huggingface_hub.hf_api import InferenceProviderMapping
from huggingface_hub.inference._common import RequestParameters, _as_dict
from ._common import BaseConversationalTask, TaskProviderHelper, filter_none
class NscaleConversationalTask(BaseConversationalTask):
def __init__(self):
super().__init__(provider="nscale", base_url="https://inference.api.nscale.com")
class NscaleTextToImageTask(TaskProviderHelper):
def __init__(self):
super().__init__(provider="nscale", base_url="https://inference.api.nscale.com", task="text-to-image")
def _prepare_route(self, mapped_model: str, api_key: str) -> str:
return "/v1/images/generations"
def _prepare_payload_as_dict(
self, inputs: Any, parameters: Dict, provider_mapping_info: InferenceProviderMapping
) -> Optional[Dict]:
mapped_model = provider_mapping_info.provider_id
# Combine all parameters except inputs and parameters
parameters = filter_none(parameters)
if "width" in parameters and "height" in parameters:
parameters["size"] = f"{parameters.pop('width')}x{parameters.pop('height')}"
if "num_inference_steps" in parameters:
parameters.pop("num_inference_steps")
if "cfg_scale" in parameters:
parameters.pop("cfg_scale")
payload = {
"response_format": "b64_json",
"prompt": inputs,
"model": mapped_model,
**parameters,
}
return payload
def get_response(self, response: Union[bytes, Dict], request_params: Optional[RequestParameters] = None) -> Any:
response_dict = _as_dict(response)
return base64.b64decode(response_dict["data"][0]["b64_json"])
kronos-venv/lib/python3.12/site-packages/huggingface_hub/inference/_providers/openai.py
+25
View File
@@ -0,0 +1,25 @@
from typing import Optional
from huggingface_hub.hf_api import InferenceProviderMapping
from huggingface_hub.inference._providers._common import BaseConversationalTask
class OpenAIConversationalTask(BaseConversationalTask):
def __init__(self):
super().__init__(provider="openai", base_url="https://api.openai.com")
def _prepare_api_key(self, api_key: Optional[str]) -> str:
if api_key is None:
raise ValueError("You must provide an api_key to work with OpenAI API.")
if api_key.startswith("hf_"):
raise ValueError(
"OpenAI provider is not available through Hugging Face routing, please use your own OpenAI API key."
)
return api_key
def _prepare_mapping_info(self, model: Optional[str]) -> InferenceProviderMapping:
if model is None:
raise ValueError("Please provide an OpenAI model ID, e.g. `gpt-4o` or `o1`.")
return InferenceProviderMapping(
provider="openai", providerId=model, task="conversational", status="live", hf_model_id=model
)
kronos-venv/lib/python3.12/site-packages/huggingface_hub/inference/_providers/replicate.py
+72
View File
@@ -0,0 +1,72 @@
from typing import Any, Dict, Optional, Union
from huggingface_hub.hf_api import InferenceProviderMapping
from huggingface_hub.inference._common import RequestParameters, _as_dict
from huggingface_hub.inference._providers._common import TaskProviderHelper, filter_none
from huggingface_hub.utils import get_session
_PROVIDER = "replicate"
_BASE_URL = "https://api.replicate.com"
class ReplicateTask(TaskProviderHelper):
def __init__(self, task: str):
super().__init__(provider=_PROVIDER, base_url=_BASE_URL, task=task)
def _prepare_headers(self, headers: Dict, api_key: str) -> Dict:
headers = super()._prepare_headers(headers, api_key)
headers["Prefer"] = "wait"
return headers
def _prepare_route(self, mapped_model: str, api_key: str) -> str:
if ":" in mapped_model:
return "/v1/predictions"
return f"/v1/models/{mapped_model}/predictions"
def _prepare_payload_as_dict(
self, inputs: Any, parameters: Dict, provider_mapping_info: InferenceProviderMapping
) -> Optional[Dict]:
mapped_model = provider_mapping_info.provider_id
payload: Dict[str, Any] = {"input": {"prompt": inputs, **filter_none(parameters)}}
if ":" in mapped_model:
version = mapped_model.split(":", 1)[1]
payload["version"] = version
return payload
def get_response(self, response: Union[bytes, Dict], request_params: Optional[RequestParameters] = None) -> Any:
response_dict = _as_dict(response)
if response_dict.get("output") is None:
raise TimeoutError(
f"Inference request timed out after 60 seconds. No output generated for model {response_dict.get('model')}"
"The model might be in cold state or starting up. Please try again later."
)
output_url = (
response_dict["output"] if isinstance(response_dict["output"], str) else response_dict["output"][0]
)
return get_session().get(output_url).content
class ReplicateTextToImageTask(ReplicateTask):
def __init__(self):
super().__init__("text-to-image")
def _prepare_payload_as_dict(
self, inputs: Any, parameters: Dict, provider_mapping_info: InferenceProviderMapping
) -> Optional[Dict]:
payload: Dict = super()._prepare_payload_as_dict(inputs, parameters, provider_mapping_info) # type: ignore[assignment]
if provider_mapping_info.adapter_weights_path is not None:
payload["input"]["lora_weights"] = f"https://huggingface.co/{provider_mapping_info.hf_model_id}"
return payload
class ReplicateTextToSpeechTask(ReplicateTask):
def __init__(self):
super().__init__("text-to-speech")
def _prepare_payload_as_dict(
self, inputs: Any, parameters: Dict, provider_mapping_info: InferenceProviderMapping
) -> Optional[Dict]:
payload: Dict = super()._prepare_payload_as_dict(inputs, parameters, provider_mapping_info) # type: ignore[assignment]
payload["input"]["text"] = payload["input"].pop("prompt") # rename "prompt" to "text" for TTS
return payload
kronos-venv/lib/python3.12/site-packages/huggingface_hub/inference/_providers/sambanova.py
+42
View File
@@ -0,0 +1,42 @@
from typing import Any, Dict, Optional, Union
from huggingface_hub.hf_api import InferenceProviderMapping
from huggingface_hub.inference._common import RequestParameters, _as_dict
from huggingface_hub.inference._providers._common import BaseConversationalTask, TaskProviderHelper, filter_none
class SambanovaConversationalTask(BaseConversationalTask):
def __init__(self):
super().__init__(provider="sambanova", base_url="https://api.sambanova.ai")
def _prepare_payload_as_dict(
self, inputs: Any, parameters: Dict, provider_mapping_info: InferenceProviderMapping
) -> Optional[Dict]:
response_format_config = parameters.get("response_format")
if isinstance(response_format_config, dict):
if response_format_config.get("type") == "json_schema":
json_schema_config = response_format_config.get("json_schema", {})
strict = json_schema_config.get("strict")
if isinstance(json_schema_config, dict) and (strict is True or strict is None):
json_schema_config["strict"] = False
payload = super()._prepare_payload_as_dict(inputs, parameters, provider_mapping_info)
return payload
class SambanovaFeatureExtractionTask(TaskProviderHelper):
def __init__(self):
super().__init__(provider="sambanova", base_url="https://api.sambanova.ai", task="feature-extraction")
def _prepare_route(self, mapped_model: str, api_key: str) -> str:
return "/v1/embeddings"
def _prepare_payload_as_dict(
self, inputs: Any, parameters: Dict, provider_mapping_info: InferenceProviderMapping
) -> Optional[Dict]:
parameters = filter_none(parameters)
return {"input": inputs, "model": provider_mapping_info.provider_id, **parameters}
def get_response(self, response: Union[bytes, Dict], request_params: Optional[RequestParameters] = None) -> Any:
embeddings = _as_dict(response)["data"]
return [embedding["embedding"] for embedding in embeddings]

Some files were not shown because too many files have changed in this diff Show More