saidsef's picture
Upload tool
f1f51f8 verified
raw
history blame contribute delete
728 Bytes
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