m-ric's picture
m-ric HF staff
Upload tool
91cf03d verified
raw
history blame
652 Bytes
from agents import Tool
class HFModelDownloadsTool(Tool):
name = "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 name of the checkpoint."""
inputs = {"task":{"type":"string","description":"the task category (such as text-classification, depth-estimation, etc)"}}
output_type = "string"
def forward(self, task: str):
from huggingface_hub import list_models
model = next(iter(list_models(filter=task, sort="downloads", direction=-1)))
return model.id