Update app.py

app.py

@@ -1,4 +1,5 @@
 import time
+import threading
 from datetime import datetime
 from functools import lru_cache
 from concurrent.futures import ThreadPoolExecutor, as_completed
@@ -7,125 +8,178 @@ import gradio as gr
 import pandas as pd
 import matplotlib.pyplot as plt
 import seaborn as sns
-
 from huggingface_hub import HfApi
 
 
-api = HfApi()
+# ============================================================
+# API
+# ============================================================
+def get_api(token: str | None = None) -> HfApi:
+    return HfApi(token=token) if token else HfApi()
+
+
+# ============================================================
+# Rate limiter (thread-safe)
+# ============================================================
+class RateLimiter:
+    def __init__(self, calls_per_second=10):
+        self.calls_per_second = calls_per_second
+        self.last_call = 0.0
+        self.lock = threading.Lock()
+
+    def wait(self):
+        with self.lock:
+            current_time = time.time()
+            elapsed = current_time - self.last_call
+            min_interval = 1.0 / self.calls_per_second
+            if elapsed < min_interval:
+                time.sleep(min_interval - elapsed)
+            self.last_call = time.time()
+
+
+rate_limiter = RateLimiter(calls_per_second=10)
 
 
-# -----------------------------
-# Cached HF API helpers
-# -----------------------------
+# ============================================================
+# Public cache helpers
+# Note:
+# - We cache only public calls.
+# - For authenticated/private access, we bypass cache so we don't
+#   mix results across users/tokens.
+# ============================================================
 @lru_cache(maxsize=1000)
-def cached_repo_info(repo_id, repo_type):
-    return api.repo_info(repo_id=repo_id, repo_type=repo_type)
+def cached_repo_info_public(repo_id: str, repo_type: str):
+    return HfApi().repo_info(repo_id=repo_id, repo_type=repo_type)
 
 
 @lru_cache(maxsize=1000)
-def cached_list_commits(repo_id, repo_type):
-    return list(api.list_repo_commits(repo_id=repo_id, repo_type=repo_type))
+def cached_list_commits_public(repo_id: str, repo_type: str):
+    return list(HfApi().list_repo_commits(repo_id=repo_id, repo_type=repo_type))
 
 
-@lru_cache(maxsize=100)
-def cached_list_items(username, kind):
+@lru_cache(maxsize=200)
+def cached_list_items_public(username: str, kind: str):
+    api = HfApi()
     if kind == "model":
         return list(api.list_models(author=username))
-    elif kind == "dataset":
+    if kind == "dataset":
         return list(api.list_datasets(author=username))
-    elif kind == "space":
+    if kind == "space":
         return list(api.list_spaces(author=username))
     return []
 
 
-# -----------------------------
-# Simple rate limiter
-# -----------------------------
-class RateLimiter:
-    def __init__(self, calls_per_second=10):
-        self.calls_per_second = calls_per_second
-        self.last_call = 0
-
-    def wait(self):
-        current_time = time.time()
-        time_since_last_call = current_time - self.last_call
-        min_interval = 1.0 / self.calls_per_second
-        if time_since_last_call < min_interval:
-            time.sleep(min_interval - time_since_last_call)
-        self.last_call = time.time()
-
-
-rate_limiter = RateLimiter()
-
-
-# -----------------------------
-# Data fetching
-# -----------------------------
-def fetch_commits_for_repo(repo_id, repo_type, selected_year):
+def list_items(username: str, kind: str, token: str | None = None):
+    if token:
+        api = get_api(token)
+        if kind == "model":
+            return list(api.list_models(author=username))
+        if kind == "dataset":
+            return list(api.list_datasets(author=username))
+        if kind == "space":
+            return list(api.list_spaces(author=username))
+        return []
+    return cached_list_items_public(username, kind)
+
+
+def repo_info(repo_id: str, repo_type: str, token: str | None = None):
+    if token:
+        return get_api(token).repo_info(repo_id=repo_id, repo_type=repo_type)
+    return cached_repo_info_public(repo_id, repo_type)
+
+
+def list_commits(repo_id: str, repo_type: str, token: str | None = None):
+    if token:
+        return list(get_api(token).list_repo_commits(repo_id=repo_id, repo_type=repo_type))
+    return cached_list_commits_public(repo_id, repo_type)
+
+
+# ============================================================
+# Data fetch
+# ============================================================
+def fetch_commits_for_repo(
+    repo_id: str,
+    repo_type: str,
+    selected_year: int,
+    token: str | None = None,
+):
     try:
         rate_limiter.wait()
 
-        repo_info = cached_repo_info(repo_id, repo_type)
-
-        # Ignore private / gated repos
-        if getattr(repo_info, "private", False) or getattr(repo_info, "gated", False):
-            return [], 0
+        info = repo_info(repo_id, repo_type, token=token)
 
         commit_dates = []
         commit_count = 0
 
-        # Initial creation date
-        if getattr(repo_info, "created_at", None) is not None:
-            initial_commit_date = pd.to_datetime(repo_info.created_at).tz_localize(None).date()
-            if initial_commit_date.year == selected_year:
-                commit_dates.append(initial_commit_date)
+        created_at = getattr(info, "created_at", None)
+        if created_at is not None:
+            created = pd.to_datetime(created_at).tz_localize(None).date()
+            if created.year == selected_year:
+                commit_dates.append(created)
                 commit_count += 1
 
-        # Repo commits
-        commits = cached_list_commits(repo_id, repo_type)
+        commits = list_commits(repo_id, repo_type, token=token)
         for commit in commits:
-            if getattr(commit, "created_at", None) is None:
+            commit_created_at = getattr(commit, "created_at", None)
+            if commit_created_at is None:
                 continue
-            commit_date = pd.to_datetime(commit.created_at).tz_localize(None).date()
-            if commit_date.year == selected_year:
-                commit_dates.append(commit_date)
+
+            dt = pd.to_datetime(commit_created_at).tz_localize(None).date()
+            if dt.year == selected_year:
+                commit_dates.append(dt)
                 commit_count += 1
 
-        return commit_dates, commit_count
+        return commit_dates, commit_count, None
 
-    except Exception:
-        return [], 0
+    except Exception as e:
+        return [], 0, f"{repo_type}:{repo_id} -> {str(e)}"
 
 
-def collect_kind_data(username, kind, selected_year):
+def collect_kind_data(
+    username: str,
+    kind: str,
+    selected_year: int,
+    token: str | None = None,
+):
     try:
-        items = cached_list_items(username, kind)
+        items = list_items(username, kind, token=token)
         repo_ids = [item.id for item in items]
 
         chunk_size = 5
         total_commits = 0
         all_commit_dates = []
+        errors = []
 
         for i in range(0, len(repo_ids), chunk_size):
             chunk = repo_ids[i:i + chunk_size]
+            if not chunk:
+                continue
 
             with ThreadPoolExecutor(max_workers=min(5, len(chunk))) as executor:
-                future_to_repo = {
-                    executor.submit(fetch_commits_for_repo, repo_id, kind, selected_year): repo_id
+                futures = {
+                    executor.submit(
+                        fetch_commits_for_repo,
+                        repo_id,
+                        kind,
+                        selected_year,
+                        token,
+                    ): repo_id
                     for repo_id in chunk
                 }
 
-                for future in as_completed(future_to_repo):
-                    repo_commits, repo_count = future.result()
+                for future in as_completed(futures):
+                    repo_commits, repo_count, err = future.result()
                     if repo_commits:
                         all_commit_dates.extend(repo_commits)
                         total_commits += repo_count
+                    if err:
+                        errors.append(err)
 
         return {
             "total_items": len(items),
             "commit_dates": all_commit_dates,
             "commit_count": total_commits,
-            "error": None,
+            "errors": errors,
         }
 
     except Exception as e:
@@ -133,19 +187,28 @@ def collect_kind_data(username, kind, selected_year):
             "total_items": 0,
             "commit_dates": [],
             "commit_count": 0,
-            "error": f"Error fetching {kind}s for {username}: {str(e)}",
+            "errors": [f"Error fetching {kind}s for {username}: {str(e)}"],
         }
 
 
-# -----------------------------
+# ============================================================
 # Plot helpers
-# -----------------------------
-def make_calendar_heatmap(df, title, year):
-    fig, ax = plt.subplots(figsize=(12, 1.3))
+# ============================================================
+def make_calendar_heatmap(df: pd.DataFrame, title: str, year: int, wide: bool = False):
+    figsize = (18, 3.0) if wide else (10, 3.4)
+    dpi = 180
+
+    fig, ax = plt.subplots(figsize=figsize, dpi=dpi)
 
     if df.empty:
-        ax.text(0.5, 0.5, f"No {title.lower()} found for {year}.",
-                ha="center", va="center", fontsize=11)
+        ax.text(
+            0.5,
+            0.5,
+            f"No {title.lower()} found for {year}.",
+            ha="center",
+            va="center",
+            fontsize=12,
+        )
         ax.axis("off")
         fig.tight_layout()
         return fig
@@ -162,7 +225,7 @@ def make_calendar_heatmap(df, title, year):
     heatmap_data = pd.DataFrame({"date": all_days, "count": 0})
     heatmap_data = heatmap_data.merge(df, on="date", how="left", suffixes=("", "_y"))
     heatmap_data["count"] = heatmap_data["count_y"].fillna(0)
-    heatmap_data = heatmap_data.drop("count_y", axis=1)
+    heatmap_data = heatmap_data.drop(columns=["count_y"])
 
     heatmap_data["dow"] = heatmap_data["date"].dt.dayofweek
     heatmap_data["week"] = (heatmap_data["date"] - start).dt.days // 7
@@ -174,8 +237,9 @@ def make_calendar_heatmap(df, title, year):
     month_positions = month_starts.map(lambda x: (x - start).days // 7)
 
     from matplotlib.colors import ListedColormap, BoundaryNorm
-    colors = ['#ebedf0', '#9be9a8', '#40c463', '#30a14e', '#216e39']
-    bounds = [0, 1, 3, 11, 31, float('inf')]
+
+    colors = ["#ebedf0", "#9be9a8", "#40c463", "#30a14e", "#216e39"]
+    bounds = [0, 1, 3, 11, 31, float("inf")]
     cmap = ListedColormap(colors)
     norm = BoundaryNorm(bounds, cmap.N)
 
@@ -184,64 +248,86 @@ def make_calendar_heatmap(df, title, year):
         ax=ax,
         cmap=cmap,
         norm=norm,
-        linewidths=0.5,
+        linewidths=0.6,
         linecolor="white",
         square=True,
         cbar=False,
         yticklabels=["M", "T", "W", "T", "F", "S", "S"],
     )
 
-    ax.set_title(title, fontsize=12, pad=10)
+    ax.set_title(title, fontsize=14, pad=12)
     ax.set_xlabel("")
     ax.set_ylabel("")
-    ax.set_xticks(month_positions)
-    ax.set_xticklabels(month_labels, fontsize=8)
-    ax.set_yticklabels(ax.get_yticklabels(), rotation=0, fontsize=8)
+    ax.set_xticks(month_positions + 0.5)
+    ax.set_xticklabels(month_labels, fontsize=10)
+    ax.set_yticklabels(ax.get_yticklabels(), rotation=0, fontsize=10)
 
     fig.tight_layout()
     return fig
 
 
-def metric_md(label, value):
-    return f"### {label}\n# {value}"
+def metric_card(label: str, value):
+    return f"""
+    <div style="
+        padding: 14px 16px;
+        border: 1px solid #e5e7eb;
+        border-radius: 14px;
+        background: white;
+    ">
+      <div style="font-size: 14px; opacity: 0.75;">{label}</div>
+      <div style="font-size: 30px; font-weight: 700; margin-top: 4px;">{value}</div>
+    </div>
+    """
 
 
-# -----------------------------
+# ============================================================
 # Main callback
-# -----------------------------
-def run_app(selected_username, custom_username, selected_year):
+# Special Gradio injected params:
+# - profile: gr.OAuthProfile | None
+# - oauth_token: gr.OAuthToken | None
+# ============================================================
+def run_app(
+    selected_username,
+    custom_username,
+    selected_year,
+    profile: gr.OAuthProfile | None,
+    oauth_token: gr.OAuthToken | None,
+):
     username = (custom_username or "").strip() or selected_username
+    selected_year = int(selected_year)
 
     if not username:
-        empty_fig = make_calendar_heatmap(pd.DataFrame(), "All Commits", selected_year)
+        empty_all = make_calendar_heatmap(pd.DataFrame(), "All Commits", selected_year, wide=True)
+        empty_kind = make_calendar_heatmap(pd.DataFrame(), "Commits", selected_year, wide=False)
         return (
             "## Please provide a username or org name.",
-            metric_md("Total Commits", 0),
-            empty_fig,
-            metric_md("🧠 Models", 0),
-            metric_md(f"Commits in {selected_year}", 0),
-            make_calendar_heatmap(pd.DataFrame(), "Models Commits", selected_year),
-            metric_md("📦 Datasets", 0),
-            metric_md(f"Commits in {selected_year}", 0),
-            make_calendar_heatmap(pd.DataFrame(), "Datasets Commits", selected_year),
-            metric_md("🚀 Spaces", 0),
-            metric_md(f"Commits in {selected_year}", 0),
-            make_calendar_heatmap(pd.DataFrame(), "Spaces Commits", selected_year),
-            "No username provided."
+            metric_card("Total Commits", 0),
+            empty_all,
+            metric_card("🧠 Models", 0),
+            metric_card(f"Commits in {selected_year}", 0),
+            empty_kind,
+            metric_card("📦 Datasets", 0),
+            metric_card(f"Commits in {selected_year}", 0),
+            empty_kind,
+            metric_card("🚀 Spaces", 0),
+            metric_card(f"Commits in {selected_year}", 0),
+            empty_kind,
+            "No username provided.",
         )
 
+    token = oauth_token.token if oauth_token is not None else None
+
     commits_by_type = {}
     commit_counts_by_type = {}
     totals_by_type = {}
     warnings = []
 
     for kind in ["model", "dataset", "space"]:
-        result = collect_kind_data(username, kind, int(selected_year))
+        result = collect_kind_data(username, kind, selected_year, token=token)
         commits_by_type[kind] = result["commit_dates"]
         commit_counts_by_type[kind] = result["commit_count"]
         totals_by_type[kind] = result["total_items"]
-        if result["error"]:
-            warnings.append(result["error"])
+        warnings.extend(result["errors"])
 
     total_commits = sum(commit_counts_by_type.values())
 
@@ -254,109 +340,128 @@ def run_app(selected_username, custom_username, selected_year):
         all_df = all_df.drop_duplicates()
 
     model_df = pd.DataFrame(commits_by_type["model"], columns=["date"])
+    dataset_df = pd.DataFrame(commits_by_type["dataset"], columns=["date"])
+    space_df = pd.DataFrame(commits_by_type["space"], columns=["date"])
+
     if not model_df.empty:
         model_df = model_df.drop_duplicates()
-
-    dataset_df = pd.DataFrame(commits_by_type["dataset"], columns=["date"])
     if not dataset_df.empty:
         dataset_df = dataset_df.drop_duplicates()
-
-    space_df = pd.DataFrame(commits_by_type["space"], columns=["date"])
     if not space_df.empty:
         space_df = space_df.drop_duplicates()
 
-    header = f"## {username}'s Activity in {selected_year}"
-    warnings_text = "\n".join(warnings) if warnings else "Done."
+    auth_status = (
+        f"Logged in as **{profile.username}**. Private repos accessible if your token has access."
+        if profile is not None
+        else "Not logged in. Public repos only."
+    )
 
     return (
-        header,
-        metric_md("Total Commits", total_commits),
-        make_calendar_heatmap(all_df, "All Commits", int(selected_year)),
+        f"## {username}'s Activity in {selected_year}\n\n{auth_status}",
+        metric_card("Total Commits", total_commits),
+        make_calendar_heatmap(all_df, "All Commits", selected_year, wide=True),
 
-        metric_md("🧠 Models", totals_by_type["model"]),
-        metric_md(f"Commits in {selected_year}", commit_counts_by_type["model"]),
-        make_calendar_heatmap(model_df, "Models Commits", int(selected_year)),
+        metric_card("🧠 Models", totals_by_type["model"]),
+        metric_card(f"Commits in {selected_year}", commit_counts_by_type["model"]),
+        make_calendar_heatmap(model_df, "Models Commits", selected_year, wide=False),
 
-        metric_md("📦 Datasets", totals_by_type["dataset"]),
-        metric_md(f"Commits in {selected_year}", commit_counts_by_type["dataset"]),
-        make_calendar_heatmap(dataset_df, "Datasets Commits", int(selected_year)),
+        metric_card("📦 Datasets", totals_by_type["dataset"]),
+        metric_card(f"Commits in {selected_year}", commit_counts_by_type["dataset"]),
+        make_calendar_heatmap(dataset_df, "Datasets Commits", selected_year, wide=False),
 
-        metric_md("🚀 Spaces", totals_by_type["space"]),
-        metric_md(f"Commits in {selected_year}", commit_counts_by_type["space"]),
-        make_calendar_heatmap(space_df, "Spaces Commits", int(selected_year)),
+        metric_card("🚀 Spaces", totals_by_type["space"]),
+        metric_card(f"Commits in {selected_year}", commit_counts_by_type["space"]),
+        make_calendar_heatmap(space_df, "Spaces Commits", selected_year, wide=False),
 
-        warnings_text,
+        "\n".join(warnings) if warnings else "Done.",
     )
 
 
-# -----------------------------
-# Gradio UI
-# -----------------------------
+# ============================================================
+# UI
+# ============================================================
 year_options = list(range(datetime.now().year, 2017, -1))
-default_users = ["ritvik77", "facebook", "google", "stabilityai", "Salesforce", "tiiuae", "bigscience"]
-
-with gr.Blocks(title="HF Contributions") as demo:
+default_users = [
+    "ritvik77",
+    "facebook",
+    "google",
+    "stabilityai",
+    "Salesforce",
+    "tiiuae",
+    "bigscience",
+]
+
+with gr.Blocks(title="HF Contributions", fill_width=True) as demo:
     gr.Markdown("# 🤗 Hugging Face Contributions")
 
     with gr.Sidebar():
+        gr.LoginButton("Sign in with Hugging Face")
+
         gr.Markdown("## 👤 Contributor")
         selected_username = gr.Dropdown(
             choices=default_users,
             value="ritvik77",
-            label="Select a username/org"
+            label="Select a username/org",
         )
-        gr.Markdown("**OR**")
         custom_username = gr.Textbox(
             label="Custom username/org",
-            placeholder="Enter custom username/org"
+            placeholder="Enter custom username/org",
         )
         selected_year = gr.Dropdown(
             choices=year_options,
             value=year_options[0],
-            label="🗓️ Year"
+            label="🗓️ Year",
         )
         run_button = gr.Button("Fetch commit data", variant="primary")
 
     status_header = gr.Markdown("## Ready")
-    total_commits_md = gr.Markdown("### Total Commits\n# 0")
-    all_plot = gr.Plot(label="All Commits")
-
-    with gr.Row():
-        with gr.Column():
-            models_total_md = gr.Markdown("### 🧠 Models\n# 0")
-            models_commits_md = gr.Markdown(f"### Commits in {year_options[0]}\n# 0")
-            models_plot = gr.Plot(label="Models Commits")
-
-        with gr.Column():
-            datasets_total_md = gr.Markdown("### 📦 Datasets\n# 0")
-            datasets_commits_md = gr.Markdown(f"### Commits in {year_options[0]}\n# 0")
-            datasets_plot = gr.Plot(label="Datasets Commits")
-
-        with gr.Column():
-            spaces_total_md = gr.Markdown("### 🚀 Spaces\n# 0")
-            spaces_commits_md = gr.Markdown(f"### Commits in {year_options[0]}\n# 0")
-            spaces_plot = gr.Plot(label="Spaces Commits")
-
-    logs_box = gr.Textbox(label="Logs / Warnings", interactive=False, lines=6)
+    total_commits_html = gr.HTML(metric_card("Total Commits", 0))
+    all_plot = gr.Plot(format="png")
+
+    with gr.Row(equal_height=False):
+        with gr.Column(scale=1, min_width=420):
+            models_total_html = gr.HTML(metric_card("🧠 Models", 0))
+            models_commits_html = gr.HTML(metric_card(f"Commits in {year_options[0]}", 0))
+            models_plot = gr.Plot(format="png")
+
+        with gr.Column(scale=1, min_width=420):
+            datasets_total_html = gr.HTML(metric_card("📦 Datasets", 0))
+            datasets_commits_html = gr.HTML(metric_card(f"Commits in {year_options[0]}", 0))
+            datasets_plot = gr.Plot(format="png")
+
+        with gr.Column(scale=1, min_width=420):
+            spaces_total_html = gr.HTML(metric_card("🚀 Spaces", 0))
+            spaces_commits_html = gr.HTML(metric_card(f"Commits in {year_options[0]}", 0))
+            spaces_plot = gr.Plot(format="png")
+
+    logs_box = gr.Textbox(label="Logs / Warnings", interactive=False, lines=8)
+
+    outputs = [
+        status_header,
+        total_commits_html,
+        all_plot,
+        models_total_html,
+        models_commits_html,
+        models_plot,
+        datasets_total_html,
+        datasets_commits_html,
+        datasets_plot,
+        spaces_total_html,
+        spaces_commits_html,
+        spaces_plot,
+        logs_box,
+    ]
 
     run_button.click(
         fn=run_app,
         inputs=[selected_username, custom_username, selected_year],
-        outputs=[
-            status_header,
-            total_commits_md,
-            all_plot,
-            models_total_md,
-            models_commits_md,
-            models_plot,
-            datasets_total_md,
-            datasets_commits_md,
-            datasets_plot,
-            spaces_total_md,
-            spaces_commits_md,
-            spaces_plot,
-            logs_box,
-        ],
+        outputs=outputs,
+    )
+
+    demo.load(
+        fn=run_app,
+        inputs=[selected_username, custom_username, selected_year],
+        outputs=outputs,
     )
 
 demo.queue().launch()