added support for datasets with a toggle

app.py

@@ -4,19 +4,18 @@ import pandas as pd
 import numpy as np
 import requests
 from datetime import datetime
-from typing import Dict, List, Optional
+from typing import Dict, List, Optional, Literal
 
 
 class HFDownloadsCalculator:
-
     BASE_URL = "https://huggingface.co/api"
     
     def __init__(self, token: Optional[str] = None):
         self.headers = {"Authorization": f"Bearer {token}"} if token else {}
     
-    def get_user_models_with_all_time_downloads(self, username: str) -> List[Dict]:
+    def get_user_items(self, username: str, item_type: Literal["models", "datasets"]) -> List[Dict]:
         response = requests.get(
-            f"{self.BASE_URL}/models",
+            f"{self.BASE_URL}/{item_type}",
             params={
                 "author": username,
                 "limit": 1000,
@@ -27,36 +26,37 @@ class HFDownloadsCalculator:
         response.raise_for_status()
         return response.json()
     
-    def calculate_total_downloads(self, username: str) -> Dict:
-        models = self.get_user_models_with_all_time_downloads(username)
+    def calculate_total_downloads(self, username: str, item_type: Literal["models", "datasets"]) -> Dict:
+        items = self.get_user_items(username, item_type)
         
         total_all_time = 0
         total_monthly = 0
-        model_stats = []
+        item_stats = []
         
-        for model in models:
-            model_id = model.get("modelId") or model.get("id") or model.get("_id", "unknown")
-            all_time = model.get("downloadsAllTime", 0)
-            monthly = model.get("downloads", 0)
+        for item in items:
+            item_id = item.get(f"{item_type[:-1]}Id") or item.get("id") or item.get("_id", "unknown")
+            all_time = item.get("downloadsAllTime", 0)
+            monthly = item.get("downloads", 0)
             
             total_all_time += all_time
             total_monthly += monthly
             
             if all_time > 0:
-                model_stats.append({
-                    "name": model_id,
+                item_stats.append({
+                    "name": item_id,
                     "downloads_all_time": all_time,
                     "downloads_monthly": monthly
                 })
         
-        model_stats.sort(key=lambda x: x["downloads_all_time"], reverse=True)
+        item_stats.sort(key=lambda x: x["downloads_all_time"], reverse=True)
         
         return {
             "total_downloads_all_time": total_all_time,
             "total_downloads_monthly": total_monthly,
-            "model_count": len(models),
-            "models_with_downloads": len(model_stats),
-            "top_models": model_stats
+            "item_count": len(items),
+            "items_with_downloads": len(item_stats),
+            "top_items": item_stats,
+            "item_type": item_type
         }
 
 
@@ -64,8 +64,8 @@ class HFDashboard:
     def __init__(self):
         self.calculator = HFDownloadsCalculator()
     
-    def get_model_timeseries(self, model_id: str, days: int = 30) -> pd.DataFrame:
-        response = requests.get(f"https://huggingface.co/api/models/{model_id}")
+    def get_item_timeseries(self, item_id: str, item_type: str, days: int = 30) -> pd.DataFrame:
+        response = requests.get(f"https://huggingface.co/api/{item_type}/{item_id}")
         data = response.json()
         
         avg_daily = data.get('downloads', 0) / 30
@@ -78,12 +78,13 @@ class HFDashboard:
             'downloads': daily_downloads
         })
     
-    def create_dashboard(self, username: str):
+    def create_dashboard(self, username: str, item_type: str):
         if not username:
             return None, None, None, "Please enter a username"
         
         try:
-            stats = self.calculator.calculate_total_downloads(username)
+            stats = self.calculator.calculate_total_downloads(username, item_type)
+            type_label = item_type.capitalize()
             
             # Metrics HTML
             metrics_html = f"""
@@ -97,8 +98,8 @@ class HFDashboard:
                     <p style="margin: 5px 0; color: #a8a8b8;">Monthly Downloads</p>
                 </div>
                 <div style="text-align: center; padding: 20px; background: linear-gradient(135deg, #1e1e2e 0%, #2d2d44 100%); border-radius: 10px; flex: 1; margin: 0 10px; border: 1px solid #3d3d5c;">
-                    <h2 style="margin: 0; color: #fff;">{stats['model_count']}</h2>
-                    <p style="margin: 5px 0; color: #a8a8b8;">Total Models</p>
+                    <h2 style="margin: 0; color: #fff;">{stats['item_count']}</h2>
+                    <p style="margin: 5px 0; color: #a8a8b8;">Total {type_label}</p>
                 </div>
             </div>
             """
@@ -108,15 +109,15 @@ class HFDashboard:
             colors = ['#6366f1', '#10b981', '#f59e0b', '#ef4444', '#00b4d8']
             colors_rgba = [f'rgba({int(c[1:3],16)}, {int(c[3:5],16)}, {int(c[5:7],16)}, 0.1)' for c in colors]
             
-            for i, model in enumerate(stats['top_models'][:5]):
-                ts_data = self.get_model_timeseries(model['name'])
+            for i, item in enumerate(stats['top_items'][:5]):
+                ts_data = self.get_item_timeseries(item['name'], item_type)
                 color_idx = i % len(colors)
                 
                 fig_line.add_trace(go.Scatter(
                     x=ts_data['date'],
                     y=ts_data['downloads'],
                     mode='lines',
-                    name=model['name'].split('/')[-1],
+                    name=item['name'].split('/')[-1],
                     line=dict(color=colors[color_idx], width=3),
                     hovertemplate='%{y} downloads<br>%{x|%b %d}',
                     fill='tozeroy',
@@ -125,7 +126,7 @@ class HFDashboard:
             
             fig_line.update_layout(
                 height=400,
-                title=dict(text="Top 5 Models - Daily Download Trends", font=dict(size=18), x=0.5, xanchor='center'),
+                title=dict(text=f"Top 5 {type_label} - Daily Download Trends", font=dict(size=18), x=0.5, xanchor='center'),
                 xaxis_title="Date",
                 yaxis_title="Daily Downloads",
                 hovermode='x unified',
@@ -141,7 +142,7 @@ class HFDashboard:
             
             # Bar chart for download distribution
             fig_bar = go.Figure()
-            top_10 = stats['top_models'][:10]
+            top_10 = stats['top_items'][:10]
             
             fig_bar.add_trace(go.Bar(
                 x=[m['name'].split('/')[-1] for m in top_10],
@@ -161,8 +162,8 @@ class HFDashboard:
             
             fig_bar.update_layout(
                 height=400,
-                title=dict(text="Top 10 Models - Download Distribution", font=dict(size=18), x=0.5, xanchor='center'),
-                xaxis_title="Model",
+                title=dict(text=f"Top 10 {type_label} - Download Distribution", font=dict(size=18), x=0.5, xanchor='center'),
+                xaxis_title=type_label[:-1],
                 yaxis_title="Downloads",
                 barmode='group',
                 template='plotly_dark',
@@ -180,13 +181,13 @@ class HFDashboard:
             # Create table
             df = pd.DataFrame([
                 [
-                    model['name'],
-                    f"{model['downloads_all_time']:,}",
-                    f"{model['downloads_monthly']:,}",
-                    f"{(model['downloads_monthly'] / model['downloads_all_time'] * 100):.1f}%" if model['downloads_all_time'] > 0 else "0%"
+                    item['name'],
+                    f"{item['downloads_all_time']:,}",
+                    f"{item['downloads_monthly']:,}",
+                    f"{(item['downloads_monthly'] / item['downloads_all_time'] * 100):.1f}%" if item['downloads_all_time'] > 0 else "0%"
                 ]
-                for model in stats['top_models']
-            ], columns=["Model", "All-Time Downloads", "Monthly Downloads", "Monthly %"])
+                for item in stats['top_items']
+            ], columns=[type_label[:-1], "All-Time Downloads", "Monthly Downloads", "Monthly %"])
             
             return metrics_html, fig_line, fig_bar, df
             
@@ -207,37 +208,51 @@ def main():
         )
     ) as app:
         gr.Markdown("# 🤗 HuggingFace Downloads Dashboard")
-        gr.Markdown("Track your model downloads and visualize trends over time")
+        gr.Markdown("Track your model and dataset downloads and visualize trends over time")
         
         with gr.Row():
-            with gr.Column():
+            with gr.Column(scale=3):
                 username_input = gr.Textbox(
                     label="HuggingFace Username",
                     placeholder="Enter username (e.g., macadeliccc)",
                     value="macadeliccc"
                 )
                 refresh_btn = gr.Button("Load Dashboard", variant="primary", size="lg")
+            
+            with gr.Column(scale=1):
+                type_selector = gr.Radio(
+                    ["models", "datasets"],
+                    value="models",
+                    label="Select Type",
+                    info="Choose between models or datasets"
+                )
         
         metrics_display = gr.HTML()
         line_plot = gr.Plot()
         bar_plot = gr.Plot()
         table_output = gr.Dataframe(
-            headers=["Model", "All-Time Downloads", "Monthly Downloads", "Monthly %"],
-            label="All Models with Downloads"
+            headers=["Item", "All-Time Downloads", "Monthly Downloads", "Monthly %"],
+            label="All Items with Downloads"
         )
         
-        def update_dashboard(username):
-            return dashboard.create_dashboard(username)
+        def update_dashboard(username, item_type):
+            return dashboard.create_dashboard(username, item_type)
         
         refresh_btn.click(
             fn=update_dashboard,
-            inputs=[username_input],
+            inputs=[username_input, type_selector],
+            outputs=[metrics_display, line_plot, bar_plot, table_output]
+        )
+        
+        type_selector.change(
+            fn=update_dashboard,
+            inputs=[username_input, type_selector],
             outputs=[metrics_display, line_plot, bar_plot, table_output]
         )
         
         app.load(
             fn=update_dashboard,
-            inputs=[username_input],
+            inputs=[username_input, type_selector],
             outputs=[metrics_display, line_plot, bar_plot, table_output]
         )