File size: 2,978 Bytes
d4fbda3 |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 |
from huggingface_hub import HfApi
import csv
import os
import gradio as gr
import pandas as pd
from pathlib import Path
import tempfile
def get_model_stats(search_term):
# Initialize the API
api = HfApi()
# Create a temporary file for the CSV
temp_dir = tempfile.mkdtemp()
output_file = Path(temp_dir) / f"{search_term}_models_alltime.csv"
# Get the generator of models with the working sort parameter
print(f"Fetching {search_term} models with all-time download statistics...")
models_generator = api.list_models(
search=search_term,
expand=["downloadsAllTime"],
sort="_id" # Sort by ID to avoid timeout issues
)
# Create and write to CSV
with open(output_file, 'w', newline='', encoding='utf-8') as csvfile:
csv_writer = csv.writer(csvfile)
# Write header
csv_writer.writerow(["Model ID", "Downloads (30 days)", "Downloads (All Time)", "Trending Score"])
# Process models
model_count = 0
for model in models_generator:
# Write to CSV
csv_writer.writerow([
getattr(model, 'id', "Unknown"),
getattr(model, 'downloads', 0), # Last 30 days downloads
getattr(model, 'downloads_all_time', 0),
getattr(model, 'trending_score', 0)
])
model_count += 1
# Read the CSV file into a pandas DataFrame
df = pd.read_csv(output_file)
# Return both the DataFrame, status message, and the CSV file path
return df, f"Found {model_count} models for search term '{search_term}'", str(output_file)
# Create the Gradio interface
with gr.Blocks(title="Hugging Face Model Statistics") as demo:
gr.Markdown("# Hugging Face Model Statistics")
gr.Markdown("Enter a search term to find model statistics from Hugging Face Hub")
with gr.Row():
search_input = gr.Textbox(
label="Search Term",
placeholder="Enter a model name or keyword (e.g., 'gemma', 'llama')",
value="gemma"
)
search_button = gr.Button("Search")
with gr.Row():
output_table = gr.Dataframe(
headers=["Model ID", "Downloads (30 days)", "Downloads (All Time)", "Trending Score"],
datatype=["str", "number", "number", "number"],
label="Model Statistics"
)
status_message = gr.Textbox(label="Status")
with gr.Row():
download_button = gr.Button("Download CSV")
csv_file = gr.File(label="CSV File", visible=False)
# Store the CSV file path in a state
csv_path = gr.State()
search_button.click(
fn=get_model_stats,
inputs=search_input,
outputs=[output_table, status_message, csv_path]
)
download_button.click(
fn=lambda x: x,
inputs=csv_path,
outputs=csv_file
)
if __name__ == "__main__":
demo.launch() |