File size: 621 Bytes
c0fead9
 
 
 
 
4c20621
c0fead9
 
 
8e27464
c0fead9
 
f415d88
 
 
c0fead9
 
8e27464
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
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 Hugging Face Hub. " 
        "It takes the name of the category (such as text-classification, depth-estimation, etc), and " 
        "returns the name of the checkpoint"
    )

    inputs = ['text']
    outputs = ['text']

    def __call__(self, task: str):
        model = next(iter(list_models(filter=task, sort='downloads', direction=-1)))
        return model.id