Spaces:
Running
on
Zero
Running
on
Zero
from smolagents import Tool | |
from huggingface_hub import HfApi | |
class HFAPITool(Tool): | |
name = "hf_api" | |
description = "Use the HuggingFace API to search for models" | |
inputs = { | |
"prompt": { | |
"type": "string", | |
"description": "The prompt to search for models", | |
}, | |
} | |
output_type = "object" | |
def __init__(self): | |
super().__init__() | |
self.api = HfApi() | |
def forward(self, prompt: str): | |
models = self.api.list_models( | |
library=["transformers"], pipeline_tag="object-detection", fetch_config=True | |
) | |
print(models) | |