Spaces:
Sleeping
Sleeping
from transformers 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 | |
# Push the tool to the Hub | |
tool = HFModelDownloadsTool() | |
tool.push_to_hub("hf-model-downloads") | |
# Load the tool from the Hub | |
loaded_tool = Tool.from_hub("hf-model-downloads") | |
# Instantiate the HfAgent with the additional tool | |
from transformers import HfAgent | |
agent = HfAgent("https://api-inference.huggingface.co/models/bigcode/starcoder", additional_tools=[loaded_tool]) | |
# Run the agent with the new tool | |
result = agent.run( | |
"Can you read out loud the name of the model that has the most downloads in the 'text-to-video' task on the Hugging Face Hub?" | |
) | |
print(result) | |