File size: 620 Bytes
a0a4d30
 
aee1d4b
 
 
 
 
 
 
a0a4d30
 
 
 
 
 
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
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