tool.py

from typing import Any, Optional
from smolagents.tools import Tool
import huggingface_hub

class HFModelDownloadsTool(Tool):
    name = "model_download_counter"
    description = """
    这是一个工具，它会在Hugging Face Hub上返回特定任务的下载量最高的模型。
它返回的是检查点的名称。"""
    inputs = {'task': {'type': 'string', 'description': '任务类别（如文本分类、深度估计等）'}}
    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