Surface modality filter on the UI: family column + status line showing required capability and eligible-candidate count

app.py

@@ -12,9 +12,14 @@ import os
 import traceback
 
 import gradio as gr
+import numpy as np
 import pandas as pd
 
-from recommend import default_recommender
+from recommend import (
+    caps_bits_to_labels,
+    default_recommender,
+    _task_required_caps_bits,
+)
 
 
 # Load once at module import time so the model is warm before the first request.
@@ -59,11 +64,14 @@ def _format_size(size_b: float) -> str:
     return f"{size_b * 1_000_000:.0f}K"
 
 
+_TABLE_COLS = ["rank", "model", "family", "score", "size", "popularity", "link"]
+
+
 def recommend_ui(dataset_description: str, task: str, metric: str, top_k: int,
                  min_size: float, max_size: float, official_only: bool, hf_only: bool,
                  api_key: str):
     if not (dataset_description or "").strip():
-        return pd.DataFrame(columns=["rank", "model", "score", "size", "popularity", "link"]), \
+        return pd.DataFrame(columns=_TABLE_COLS), \
                "Please enter a dataset description."
 
     api_key = (api_key or "").strip()
@@ -104,13 +112,26 @@ def recommend_ui(dataset_description: str, task: str, metric: str, top_k: int,
         rows.append({
             "rank": r.rank,
             "model": r.model_name,
+            "family": r.family or "—",
             "score": round(r.score, 4),
             "size": _format_size(r.size_b),
             "popularity": r.popularity,
             "link": link,
         })
-    df = pd.DataFrame(rows, columns=["rank", "model", "score", "size", "popularity", "link"])
-    return df, f"Returned top-{len(rows)} of {len(RECOMMENDER.model_names)} candidates."
+    df = pd.DataFrame(rows, columns=_TABLE_COLS)
+
+    # Surface what the modality filter is doing so the user can see why
+    # certain candidates were (or weren't) eligible.
+    caps_bits = _task_required_caps_bits(task)
+    caps_labels = caps_bits_to_labels(caps_bits)
+    n_compat = int(((RECOMMENDER.model_caps_bits & caps_bits) != 0).sum())
+    caps_str = ", ".join(f"`{c}`" for c in caps_labels) if caps_labels else "any"
+    status = (
+        f"Task **{task}** requires model capability: {caps_str} → "
+        f"**{n_compat:,}** of {len(RECOMMENDER.model_names):,} candidates eligible. "
+        f"Returned top-{len(rows)}."
+    )
+    return df, status
 
 
 with gr.Blocks(title="ModelLens · Finding the Best Model for Your Task", theme=gr.themes.Soft()) as demo:
@@ -121,6 +142,11 @@ with gr.Blocks(title="ModelLens · Finding the Best Model for Your Task", theme=
         the top candidates from a pool of **47k+** HuggingFace models. Backed by the
         ablation_no_id MLPMetric checkpoint trained on `unified_augmented`.
 
+        Results are post-filtered by a modality sanity check so that e.g.
+        *Image Generation* won't surface text-only LLMs. The status line below
+        the table shows which capability your task requires and how many
+        candidates passed the filter.
+
         > **BYO OpenAI key.** This Space embeds your dataset description with
         > `text-embedding-3-small`.
         """
@@ -178,10 +204,10 @@ with gr.Blocks(title="ModelLens · Finding the Best Model for Your Task", theme=
         with gr.Column(scale=3):
             status = gr.Markdown("")
             table = gr.Dataframe(
-                headers=["rank", "model", "score", "size", "popularity", "link"],
+                headers=_TABLE_COLS,
                 interactive=False,
                 wrap=True,
-                datatype=["number", "str", "number", "str", "number", "markdown"],
+                datatype=["number", "str", "str", "number", "str", "number", "markdown"],
             )
 
     run_btn.click(

recommend.py

@@ -400,6 +400,21 @@ def _task_required_caps_bits(task: str) -> int:
     return _DEFAULT_TASK_CAPS
 
 
+def caps_bits_to_labels(bits: int) -> list[str]:
+    """Human-readable capability labels (e.g., 'vision-generation') from a bitmask."""
+    pretty = {
+        "text": "text",
+        "text_embedding": "text-embedding",
+        "vision": "vision",
+        "vision_generate": "vision-generation",
+        "vision_language": "vision-language",
+        "audio": "audio",
+        "video": "video",
+        "document": "document",
+    }
+    return [pretty[c] for c in _ALL_CAPS if bits & _CAP_BIT[c]]
+
+
 def _slug(s: str) -> str:
     return re.sub(r"[^a-z0-9]+", "", str(s).strip().lower())
 
@@ -428,6 +443,7 @@ class Recommendation:
     size_bucket: int
     size_b: float  # raw size in billions of params; NaN if unknown
     family_id: int
+    family: str
     popularity: int
     hf_url: str
 
@@ -731,6 +747,7 @@ class Recommender:
                     size_bucket=int(self.size_ids[i]),
                     size_b=float(self.sizes_b[i]),
                     family_id=int(self.family_ids[i]),
+                    family=self._id2family.get(int(self.family_ids[i]), ""),
                     popularity=int(self.popularities[i]),
                     hf_url=self.urls[i],
                 )