File size: 728 Bytes
f1f51f8
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
from smolagents.tools import Tool
import huggingface_hub

class HFModelsTrendingTool(Tool):
    name = "model_trending_counter"
    description = """
    This is a tool that returns the most trending models 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(task=task, sort="trending_score", direction=-1)))
        return model.id


    def __init__(self, *args, **kwargs):
        self.is_initialized = False