Spaces:
Running
Running
File size: 8,500 Bytes
c770c2e |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 146 147 148 149 150 151 152 153 154 155 156 157 158 159 160 161 162 163 164 165 166 167 168 169 170 171 172 173 174 175 176 177 178 179 180 181 182 183 184 185 186 187 188 189 190 191 192 193 194 195 196 197 198 199 200 201 202 203 204 205 206 207 208 209 210 211 212 213 214 215 216 217 218 219 220 221 |
import json
from typing import Literal
import gradio as gr
from huggingface_hub import list_models, model_info, hf_hub_download
def search_models(
search: str = None,
library: str = None,
tags: str = None,
pipeline_tag: str = None,
sort: Literal[
"trending_score", "last_modified", "created_at", "downloads", "likes"
] = "trending_score",
direction: Literal["descending", "ascending"] = "descending",
limit: int = 20,
) -> str:
"""
Search models on Hugging Face Hub.
Use this tool to search for models by name, tags, or other filters, and to get a list of model IDs.
This is the first step when you need to find a specific model before retrieving its details.
Parameters:
search (str, optional): A string to search for in model IDs or names (e.g., "deepseek").
library (list[str], optional): List of libraries the models use (e.g., ["pytorch", "tensorflow"]).
tags (list[str], optional): List of tags to filter models by (e.g., ["text-generation", "llama"]).
pipeline_tag (str, optional): Filter by pipeline tag (e.g., "text-generation").
sort (Literal["trending_score", "last_modified", "created_at", "downloads", "likes"], default="trending_score"): Sort models by the specified key.
direction (int, default=-1): Sort direction: -1 for descending, 1 for ascending.
limit (int, default=20): Maximum number of models to return.
Returns:
list[str]: A list of model IDs matching the search criteria.
Examples:
- To find trending models: search_models(sort="trending_score", limit=10)
- To search for models related to "deepseek": search_models(search="deepseek", sort="likes", limit=5)
- To filter by tag: search_models(tags=["text-generation"], pipeline_tag="text-generation")
"""
try:
library = library.split(",") if library else None
tags = tags.split(",") if tags else None
direction = -1 if direction == "descending" else 1
models = list_models(
library=library,
tags=tags,
search=search,
pipeline_tag=pipeline_tag,
sort=sort,
direction=direction,
limit=limit,
)
return json.dumps([model.modelId for model in models])
except Exception as e:
return f"Error: {e}"
def get_model_info(model_id: str) -> dict:
"""
Get structured metadata about a model on the Hugging Face Hub.
Use this when you need specific fields like downloads, tags, or other metadata.
For comprehensive model information, use `get_model_card`.
This tool requires the exact model ID, which can be obtained using `search_models`.
If you have a partial name or tag, use `search_models` first to find the exact ID.
Parameters:
model_id (str): The exact model ID in the format "organization/model-name" (e.g., "DeepSeek/DeepSeek-R1").
Returns:
dict: A dictionary containing model information including available fields such as:
- id: The model ID
- author: The author of the model
- created_at: The creation date
- last_modified: The last modified date
- downloads: Number of downloads
- likes: Number of likes
- tags: List of tags
- pipeline_tag: The pipeline tag
- library_name: The library name
- license: The model license
- base_model: The base model (if available)
- siblings: List of repository files (if available)
- datasets: Datasets used to train the model (if available)
- spaces: List of spaces using this model (if available)
- xet_enabled: Whether XET is enabled (if available)
Raises:
Exception: If the model_id is invalid or not found. Use search_models to find the correct ID.
Example:
- First, find the model ID: search_models(search="deepseek", sort="likes", limit=1)
- Then, get the model info: get_model_info("DeepSeek/DeepSeek-R1")
"""
try:
model = model_info(model_id)
result = {}
if hasattr(model, "id") and model.id is not None:
result["id"] = model.id
if hasattr(model, "author") and model.author is not None:
result["author"] = model.author
if hasattr(model, "created_at") and model.created_at is not None:
result["created_at"] = model.created_at
if hasattr(model, "last_modified") and model.last_modified is not None:
result["last_modified"] = model.last_modified
if hasattr(model, "downloads") and model.downloads is not None:
result["downloads"] = model.downloads
if hasattr(model, "likes") and model.likes is not None:
result["likes"] = model.likes
if hasattr(model, "tags") and model.tags is not None:
result["tags"] = model.tags
if hasattr(model, "pipeline_tag") and model.pipeline_tag is not None:
result["pipeline_tag"] = model.pipeline_tag
if hasattr(model, "library_name") and model.library_name is not None:
result["library_name"] = model.library_name
if hasattr(model, "card_data") and model.card_data is not None:
if (
hasattr(model.card_data, "license")
and model.card_data.license is not None
):
result["license"] = model.card_data.license
if (
hasattr(model.card_data, "base_model")
and model.card_data.base_model is not None
):
result["base_model"] = model.card_data.base_model
if (
hasattr(model.card_data, "datasets")
and model.card_data.datasets is not None
):
result["datasets"] = model.card_data.datasets
if hasattr(model, "siblings") and model.siblings is not None:
result["siblings"] = model.siblings
if hasattr(model, "spaces") and model.spaces is not None:
result["spaces"] = model.spaces
if hasattr(model, "xet_enabled") and model.xet_enabled is not None:
result["xet_enabled"] = model.xet_enabled
return json.dumps(result)
except Exception as e:
return f"Error: {e}"
def get_model_card(model_id: str) -> str:
"""
Get the complete model card (README.md) for a specific model on Hugging Face Hub.
Use this when you need comprehensive model documentation including usage examples, model limitations, etc.
For only structured metadata, use `get_model_info` instead.
This tool requires the exact model ID, which can be obtained using `search_models`.
If you have a partial name or tag, use `search_models` first to find the exact ID.
Args:
model_id (str): The model ID in the format "organization/model-name" (e.g., "DeepSeek/DeepSeek-R1").
Returns:
str: The markdown content of the model card.
Example:
- First, find the model ID: search_models(search="deepseek", sort="likes", limit=1)
- Then, get the model card: get_model_card("DeepSeek/DeepSeek-R1")
"""
try:
filepath = hf_hub_download(model_id, "README.md")
with open(filepath, "r", encoding="utf-8") as f:
content = f.read()
return content
except Exception as e:
return f"Error: {e}"
search_models = gr.Interface(
fn=search_models,
inputs=[
gr.Textbox(label="search", value=""),
gr.Textbox(label="library", value=""),
gr.Textbox(label="tags", value=""),
gr.Textbox(label="pipeline_tag", value=""),
gr.Radio(label="sort", choices=["trending_score", "last_modified", "created_at", "downloads", "likes"], value="trending_score"),
gr.Radio(label="direction", choices=["descending", "ascending"], value="descending"),
gr.Number(label="limit", value=20),
],
outputs="text")
get_model_info = gr.Interface(
fn=get_model_info,
inputs=[
gr.Textbox(label="model_id", value=""),
],
outputs="text")
get_model_card = gr.Interface(
fn=get_model_card,
inputs=[
gr.Textbox(label="model_id", value=""),
],
outputs="text")
demo = gr.TabbedInterface([search_models, get_model_info, get_model_card], ["search_models", "get_model_info", "get_model_card"])
demo.launch(mcp_server=True)
|