most-downloaded-model / most_downloaeded_model.py
Chris4K's picture
Rename text_generator.py to most_downloaeded_model.py
99a0ef6
raw
history blame
No virus
1.19 kB
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)