models.py 4.0 KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133
  1. # Copyright 2026 The HuggingFace Team. All rights reserved.
  2. #
  3. # Licensed under the Apache License, Version 2.0 (the "License");
  4. # you may not use this file except in compliance with the License.
  5. # You may obtain a copy of the License at
  6. #
  7. # http://www.apache.org/licenses/LICENSE-2.0
  8. #
  9. # Unless required by applicable law or agreed to in writing, software
  10. # distributed under the License is distributed on an "AS IS" BASIS,
  11. # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
  12. # See the License for the specific language governing permissions and
  13. # limitations under the License.
  14. """Contains commands to interact with models on the Hugging Face Hub.
  15. Usage:
  16. # list models on the Hub
  17. hf models ls
  18. # list models with a search query
  19. hf models ls --search "llama"
  20. # get info about a model
  21. hf models info Lightricks/LTX-2
  22. """
  23. import enum
  24. from typing import Annotated, get_args
  25. import typer
  26. from huggingface_hub.errors import CLIError, RepositoryNotFoundError, RevisionNotFoundError
  27. from huggingface_hub.hf_api import ExpandModelProperty_T, ModelSort_T
  28. from ._cli_utils import (
  29. AuthorOpt,
  30. FilterOpt,
  31. FormatWithAutoOpt,
  32. LimitOpt,
  33. RevisionOpt,
  34. SearchOpt,
  35. TokenOpt,
  36. api_object_to_dict,
  37. get_hf_api,
  38. make_expand_properties_parser,
  39. typer_factory,
  40. )
  41. from ._output import OutputFormatWithAuto, out
  42. _EXPAND_PROPERTIES = sorted(get_args(ExpandModelProperty_T))
  43. _SORT_OPTIONS = get_args(ModelSort_T)
  44. ModelSortEnum = enum.Enum("ModelSortEnum", {s: s for s in _SORT_OPTIONS}, type=str) # type: ignore[misc]
  45. ExpandOpt = Annotated[
  46. str | None,
  47. typer.Option(
  48. help=f"Comma-separated properties to return. When used, only the listed properties (and id) are returned. Example: '--expand=downloads,likes,tags'. Valid: {', '.join(_EXPAND_PROPERTIES)}.",
  49. callback=make_expand_properties_parser(_EXPAND_PROPERTIES),
  50. ),
  51. ]
  52. models_cli = typer_factory(help="Interact with models on the Hub.")
  53. @models_cli.command(
  54. "list | ls",
  55. examples=[
  56. "hf models ls --sort downloads --limit 10",
  57. 'hf models ls --search "llama" --author meta-llama',
  58. "hf models ls --num-parameters min:6B,max:128B --sort likes",
  59. ],
  60. )
  61. def models_ls(
  62. search: SearchOpt = None,
  63. author: AuthorOpt = None,
  64. filter: FilterOpt = None,
  65. num_parameters: Annotated[
  66. str | None,
  67. typer.Option(help="Filter by parameter count, e.g. 'min:6B,max:128B'."),
  68. ] = None,
  69. sort: Annotated[
  70. ModelSortEnum | None,
  71. typer.Option(help="Sort results."),
  72. ] = None,
  73. limit: LimitOpt = 10,
  74. expand: ExpandOpt = None,
  75. format: FormatWithAutoOpt = OutputFormatWithAuto.auto,
  76. token: TokenOpt = None,
  77. ) -> None:
  78. """List models on the Hub."""
  79. api = get_hf_api(token=token)
  80. sort_key = sort.value if sort else None
  81. results = [
  82. api_object_to_dict(model_info)
  83. for model_info in api.list_models(
  84. filter=filter,
  85. author=author,
  86. search=search,
  87. num_parameters=num_parameters,
  88. sort=sort_key,
  89. limit=limit,
  90. expand=expand, # type: ignore
  91. )
  92. ]
  93. out.table(results)
  94. @models_cli.command(
  95. "info",
  96. examples=[
  97. "hf models info meta-llama/Llama-3.2-1B-Instruct",
  98. "hf models info Qwen/Qwen3.5-9B --expand downloads,likes,tags",
  99. ],
  100. )
  101. def models_info(
  102. model_id: Annotated[str, typer.Argument(help="The model ID (e.g. `username/repo-name`).")],
  103. revision: RevisionOpt = None,
  104. expand: ExpandOpt = None,
  105. format: FormatWithAutoOpt = OutputFormatWithAuto.auto,
  106. token: TokenOpt = None,
  107. ) -> None:
  108. """Get info about a model on the Hub."""
  109. api = get_hf_api(token=token)
  110. try:
  111. info = api.model_info(repo_id=model_id, revision=revision, expand=expand) # type: ignore
  112. except RepositoryNotFoundError as e:
  113. raise CLIError(f"Model '{model_id}' not found.") from e
  114. except RevisionNotFoundError as e:
  115. raise CLIError(f"Revision '{revision}' not found on '{model_id}'.") from e
  116. out.dict(info)