from transformers.tools import Tool from huggingface_hub import list_models class HFModelDownloadsTool(Tool): name = 'model-download-counter' description = ( "This is a tool that returns the most downloaded model of a given task on the HF中国镜像站 Hub. " "It takes the name of the category (such as text-classification, depth-estimation, etc), and " "returns the name of the checkpoint alongside its 30-day download numbers." ) inputs = ['text'] outputs = ['text'] def __call__(self, task: str): model = next(iter(list_models(filter=task, sort='downloads', direction=-1))) return f"{model.id} with {model.downloads} downloads over the past 30 days."