Enhance data handling and UI for OCRQA exploration tool

- Added error handling for empty dataframes.
- Improved ranking calculations for newspapers.
- Updated newspaper selection logic to include random choices.
- Enhanced UI components for better user interaction.

app.py

@@ -1,4 +1,5 @@
 import json
+import random
 from urllib.request import urlopen
 
 import pandas as pd
@@ -51,31 +52,117 @@ for media in data.get("media_list", []):
 
 df = pd.DataFrame(rows).sort_values(["provider", "newspaper", "year"])
 
+if df.empty:
+    raise ValueError("No yearly OCRQA data found.")
+
 provider_options = ["All"] + sorted(df["provider"].dropna().unique().tolist())
 
+# -------------------------------------------------------------------
+# Rankings
+# -------------------------------------------------------------------
+
+ranking_by_provider = (
+    df.groupby(["provider", "newspaper"], as_index=False)["avg_ocrqa"]
+    .mean()
+    .rename(columns={"avg_ocrqa": "mean_ocrqa"})
+)
 
-def newspapers_for_provider(provider):
-    subset = df if provider == "All" else df[df["provider"] == provider]
-    ranking = (
-        subset.groupby("newspaper", as_index=False)["avg_ocrqa"]
-        .mean()
-        .sort_values("avg_ocrqa", ascending=False)
-    )
-    return ranking["newspaper"].tolist()
+ranking_global = (
+    df.groupby("newspaper", as_index=False)["avg_ocrqa"]
+    .mean()
+    .rename(columns={"avg_ocrqa": "mean_ocrqa"})
+)
+
+
+def get_ranked_df(provider="All", query=""):
+    if provider == "All":
+        ranked = ranking_global.copy()
+    else:
+        ranked = ranking_by_provider.loc[
+            ranking_by_provider["provider"] == provider, ["newspaper", "mean_ocrqa"]
+        ].copy()
+
+    ranked = ranked.sort_values(
+        ["mean_ocrqa", "newspaper"], ascending=[False, True]
+    ).reset_index(drop=True)
+
+    if query:
+        q = query.strip().lower()
+        ranked = ranked[
+            ranked["newspaper"].str.lower().str.contains(q, na=False)
+        ].reset_index(drop=True)
+
+    return ranked
+
+
+def choose_newspapers(ranked, n_best, n_worst, n_random, seed=13):
+    ranked_names = ranked["newspaper"].tolist()
 
+    best = ranked_names[: int(n_best)] if n_best > 0 else []
+    worst = ranked_names[-int(n_worst) :] if n_worst > 0 else []
 
-def update_newspaper_choices(provider):
-    choices = newspapers_for_provider(provider)
-    return gr.update(choices=choices, value=choices[:10])
+    remaining_for_random = [
+        n for n in ranked_names if n not in set(best) and n not in set(worst)
+    ]
+
+    rng = random.Random(seed)
+    n_random = min(int(n_random), len(remaining_for_random))
+    random_pick = rng.sample(remaining_for_random, n_random) if n_random > 0 else []
+
+    selected = best + worst + random_pick
+
+    # Deduplicate while preserving order
+    selected = list(dict.fromkeys(selected))
+
+    # Choices should remain OCRQA-ranked, not in selection order
+    choices = ranked_names
+
+    return choices, selected
+
+
+def update_newspapers(provider, query, n_best, n_worst, n_random):
+    ranked = get_ranked_df(provider, query)
+    choices, selected = choose_newspapers(ranked, n_best, n_worst, n_random)
+    return gr.update(choices=choices, value=selected)
 
 
 def make_plot(provider, selected_newspapers):
-    subset = df if provider == "All" else df[df["provider"] == provider]
+    if not selected_newspapers:
+        fig = go.Figure()
+        fig.update_layout(
+            title="Select one or more newspapers",
+            xaxis_title="Year",
+            yaxis_title="Average OCRQA",
+            yaxis=dict(range=[0, 1.05]),
+            template="plotly_white",
+            height=650,
+        )
+        return fig
+
+    subset = df.copy() if provider == "All" else df[df["provider"] == provider].copy()
     subset = subset[subset["newspaper"].isin(selected_newspapers)]
 
+    if subset.empty:
+        fig = go.Figure()
+        fig.update_layout(
+            title="No data for the current selection",
+            xaxis_title="Year",
+            yaxis_title="Average OCRQA",
+            yaxis=dict(range=[0, 1.05]),
+            template="plotly_white",
+            height=650,
+        )
+        return fig
+
+    # Preserve ranking order in legend/traces
+    ranked = get_ranked_df(provider, "")
+    ranked_order = [
+        n for n in ranked["newspaper"].tolist() if n in set(selected_newspapers)
+    ]
+
     fig = go.Figure()
 
-    for newspaper in selected_newspapers:
+    for newspaper in ranked_order:
         dfn = subset[subset["newspaper"] == newspaper].sort_values("year")
         if dfn.empty:
             continue
@@ -109,34 +196,85 @@ def make_plot(provider, selected_newspapers):
     return fig
 
 
+# -------------------------------------------------------------------
+# Initial state
+# -------------------------------------------------------------------
+
+initial_provider = "All"
+initial_query = ""
+initial_best = 10
+initial_worst = 0
+initial_random = 0
+
+initial_ranked = get_ranked_df(initial_provider, initial_query)
+initial_choices, initial_selected = choose_newspapers(
+    initial_ranked, initial_best, initial_worst, initial_random
+)
+
+# -------------------------------------------------------------------
+# UI
+# -------------------------------------------------------------------
+
 with gr.Blocks() as demo:
     gr.Markdown("## OCRQA exploration")
 
-    provider = gr.Dropdown(
-        choices=provider_options,
-        value="All",
-        label="Provider",
-    )
+    with gr.Row():
+        provider = gr.Dropdown(
+            choices=provider_options,
+            value=initial_provider,
+            label="Provider",
+        )
+        query = gr.Textbox(
+            value=initial_query,
+            label="Filter newspapers",
+            placeholder="Type part of a newspaper title",
+        )
 
-    newspaper = gr.CheckboxGroup(
-        choices=newspapers_for_provider("All"),
-        value=newspapers_for_provider("All")[:10],
-        label="Newspapers (ranked by mean OCRQA)",
+    with gr.Row():
+        n_best = gr.Slider(
+            minimum=0,
+            maximum=400,
+            value=initial_best,
+            step=1,
+            label="Best OCRQA",
+        )
+        n_worst = gr.Slider(
+            minimum=0,
+            maximum=400,
+            value=initial_worst,
+            step=1,
+            label="Worst OCRQA",
+        )
+        n_random = gr.Slider(
+            minimum=0,
+            maximum=400,
+            value=initial_random,
+            step=1,
+            label="Random OCRQA",
+        )
+
+    newspaper = gr.Dropdown(
+        choices=initial_choices,
+        value=initial_selected,
+        multiselect=True,
+        label="Newspapers (filtered and ranked)",
     )
 
     plot = gr.Plot()
 
-    provider.change(
-        fn=update_newspaper_choices,
-        inputs=provider,
-        outputs=newspaper,
-    )
+    selector_inputs = [provider, query, n_best, n_worst, n_random]
 
-    provider.change(
-        fn=make_plot,
-        inputs=[provider, newspaper],
-        outputs=plot,
-    )
+    for trigger in selector_inputs:
+        trigger.change(
+            fn=update_newspapers,
+            inputs=selector_inputs,
+            outputs=newspaper,
+        )
+        trigger.change(
+            fn=lambda provider, newspaper: make_plot(provider, newspaper),
+            inputs=[provider, newspaper],
+            outputs=plot,
+        )
 
     newspaper.change(
         fn=make_plot,