Spaces:
Running
Running
from smolagents.tools import Tool | |
import huggingface_hub | |
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 | |
def __init__(self, *args, **kwargs): | |
self.is_initialized = False | |