| from smolagents 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 | |
| model_downloads_tool = HFModelDownloadsTool() |