hf-model-downloads / model_downloads.py
m-ric's picture
Upload tool
987ce08 verified
raw
history blame
671 Bytes
from transformers import Tool
from huggingface_hub import list_models
class HFModelDownloadsTool(Tool):
task = "model_download_counter"
description = (
"This is a tool that returns the most downloaded model of a given task on the Hugging Face Hub. "
"It returns the task of the checkpoint."
)
inputs = {
"name": {
"type": str,
"description": "the name of the category (such as text-classification, depth-estimation, etc)",
}
}
output_type = str
def __call__(self, task: str):
model = next(iter(list_models(filter=task, sort="downloads", direction=-1)))
return model.id