feat: DoctoBERT Fill-Mask demo

Switch to doctolib-lab/doctomodernbert-fr-base (ModernBERT, [MASK]
token), translate the UI to English, and replace the raw JSON output
with a styled prediction card matching the DoctoBERT Diffusion demo's
dark palette and accent gradient. Auto-run predictions on load and on
example click instead of leaving the card blank.

Runs on ZeroGPU when deployed; resolves the device from
torch.cuda.is_available() rather than guessing from whether the
spaces package merely imported, which adapts correctly across
cpu-basic, ZeroGPU, and local MPS/CPU with no crashes.

README.md

@@ -1,19 +1,28 @@
 ---
-title: DoctoBERT-fr-base
-emoji: 🏥
-colorFrom: yellow
+title: DoctoBERT Fill-Mask
+emoji: 🩺
+colorFrom: indigo
 colorTo: purple
 sdk: gradio
 sdk_version: 6.19.0
 app_file: app.py
-short_description: Fill-mask demo for French medical RoBERTa
+short_description: Fill-mask demo for a French medical ModernBERT
 python_version: "3.12"
 startup_duration_timeout: 30m
 ---
 
-# DoctoBERT-fr-base
+# 🩺 DoctoBERT Fill-Mask
 
-French medical RoBERTa (111M parameters) fill-mask demo. Enter a French medical sentence
-containing a `<mask>` token and the model predicts the most likely completions.
+Fill-mask demo of [doctolib-lab/doctomodernbert-fr-base](https://huggingface.co/doctolib-lab/doctomodernbert-fr-base),
+a French medical **ModernBERT** encoder. Write a sentence containing a `[MASK]` token and the
+model predicts the most likely words, ranked by probability.
 
-Model: [doctolib-lab/doctobert-fr-base](https://huggingface.co/doctolib-lab/doctobert-fr-base)
+Runs on **ZeroGPU** when deployed (via `@spaces.GPU`); falls back to MPS/CPU when run locally.
+
+## Run locally
+
+```bash
+uv venv --python 3.12 .venv
+uv pip install torch transformers gradio
+python app.py
+```

app.py

@@ -1,136 +1,172 @@
-import spaces  # MUST come before torch / any CUDA-touching import
-import torch
-import gradio as gr
-from transformers import AutoTokenizer, AutoModelForMaskedLM
-import torch.nn.functional as F
-
-MODEL_ID = "doctolib-lab/doctobert-fr-base"
-
-tokenizer = AutoTokenizer.from_pretrained(MODEL_ID)
-model = AutoModelForMaskedLM.from_pretrained(MODEL_ID).to("cuda")
-model.eval()
-
-MASK_TOKEN_DISPLAY = "<mask>"
-MASK_TOKEN_ID = tokenizer.mask_token_id
+# ZeroGPU: `spaces` must be imported before torch. Present on any HF Spaces
+# Gradio SDK image regardless of hardware tier; absent locally.
+try:
+    import spaces
 
+    _HAS_SPACES = True
+except ImportError:
+    _HAS_SPACES = False
 
-@spaces.GPU(duration=15)
-def predict_mask(text: str, top_k: int = 5):
-    """Predict the most likely tokens to fill the <mask> in a French medical sentence.
+import html
 
-    Args:
-        text: A French medical sentence containing one <mask> token.
-        top_k: Number of top predictions to return.
-    """
-    if MASK_TOKEN_DISPLAY not in text:
-        error_msg = f"Le texte doit contenir le jeton {MASK_TOKEN_DISPLAY}."
-        return {"error": error_msg}, error_msg
-
-    # Split text around the <mask> display token and tokenize each part,
-    # then insert the actual mask token id between them.
-    parts = text.split(MASK_TOKEN_DISPLAY)
-    if len(parts) != 2:
-        error_msg = f"Le texte doit contenir exactement un jeton {MASK_TOKEN_DISPLAY}."
-        return {"error": error_msg}, error_msg
-
-    left_ids = tokenizer(parts[0], add_special_tokens=False)["input_ids"]
-    right_ids = tokenizer(parts[1], add_special_tokens=False)["input_ids"]
-
-    # Build: [CLS] left... [MASK] right... [SEP]
-    cls_id = tokenizer.cls_token_id
-    sep_id = tokenizer.sep_token_id
-    input_ids = [cls_id] + left_ids + [MASK_TOKEN_ID] + right_ids + [sep_id]
-    input_ids = torch.tensor([input_ids], device="cuda")
-    attention_mask = torch.ones_like(input_ids)
+import gradio as gr
+import torch
+import torch.nn.functional as F
+from transformers import AutoModelForMaskedLM, AutoTokenizer
+
+MODEL_ID = "doctolib-lab/doctomodernbert-fr-base"
+
+# torch.cuda.is_available() is the single source of truth for whether a CUDA
+# device is actually usable here -- true on real GPU hardware, true on a
+# properly-attached ZeroGPU Space (its emulation mode makes this report
+# correctly outside @spaces.GPU functions too), false on cpu-basic. No need to
+# infer it from whether `spaces` merely imported, which says nothing about
+# the Space's actual hardware tier.
+if torch.cuda.is_available():
+    DEVICE = "cuda"
+elif torch.backends.mps.is_available():
+    DEVICE = "mps"
+else:
+    DEVICE = "cpu"
+print(f"[doctobert-fill-mask] device={DEVICE} has_spaces={_HAS_SPACES}", flush=True)
 
-    mask_pos = len(left_ids) + 1  # +1 for [CLS]
+tokenizer = AutoTokenizer.from_pretrained(MODEL_ID)
+model = AutoModelForMaskedLM.from_pretrained(MODEL_ID).to(DEVICE).eval()
+MASK = tokenizer.mask_token  # ModernBERT -> "[MASK]"
+
+
+def gpu(fn):
+    """@spaces.GPU is documented as effect-free outside real ZeroGPU hardware,
+    so it's always safe to apply once `spaces` is importable; no-op locally,
+    where the package isn't installed at all."""
+    return spaces.GPU(duration=15)(fn) if _HAS_SPACES else fn
+
+
+def _card(inner: str) -> str:
+    return f'<div class="fm">{inner}</div>'
+
+
+def _err(msg: str) -> str:
+    return _card(f'<div class="fm-err">{html.escape(msg)}</div>')
+
+
+def _render(text: str, toks: list[str], scores: list[float]) -> str:
+    left, right = text.split(MASK, 1)
+    top1 = toks[0] if toks else ""
+    sentence = (
+        f'<span class="fm-txt">{html.escape(left)}</span>'
+        + f'<span class="fm-slot">{html.escape(top1) or "·"}</span>'
+        + f'<span class="fm-txt">{html.escape(right)}</span>'
+    )
+    maxp = max(scores) if scores else 1.0
+    rows = []
+    for r, (t, s) in enumerate(zip(toks, scores), 1):
+        w = (s / maxp * 100) if maxp else 0
+        rows.append(
+            f'<div class="fm-row"><span class="fm-rank">{r}</span>'
+            f'<span class="fm-tok">{html.escape(t) or "∅"}</span>'
+            f'<div class="fm-bar"><i style="--w:{w:.1f}%"></i></div>'
+            f'<span class="fm-pct">{s * 100:.1f}%</span></div>'
+        )
+    return _card(
+        f'<div class="fm-sentence">{sentence}</div>'
+        '<div class="fm-cap">Top predictions</div>' + "".join(rows)
+    )
+
+
+@gpu
+def predict(text: str, top_k: int = 8) -> str:
+    text = (text or "").strip()
+    if not text:
+        return _err("Type a sentence first.")
+    if MASK not in text:
+        return _err(f"Your sentence must contain a {MASK} token.")
+    if text.count(MASK) != 1:
+        return _err(f"Use exactly one {MASK} token.")
+
+    enc = tokenizer(text, return_tensors="pt").to(DEVICE)
+    ids = enc["input_ids"][0]
+    pos = (ids == tokenizer.mask_token_id).nonzero(as_tuple=True)[0]
+    if pos.numel() == 0:
+        return _err("Could not locate the mask after tokenization.")
+    mp = int(pos[0])
 
     with torch.no_grad():
-        logits = model(input_ids=input_ids, attention_mask=attention_mask).logits
-
-    mask_logits = logits[0, mask_pos, :]
-    probs = F.softmax(mask_logits, dim=-1)
-    topk_probs, topk_indices = torch.topk(probs, min(int(top_k), probs.shape[-1]))
-
-    results = []
-    for i in range(topk_probs.shape[0]):
-        token_id = topk_indices[i].item()
-        token = tokenizer.convert_ids_to_tokens(token_id)
-        if isinstance(token, str):
-            token = token.replace("▁", "").strip()
-        prob = topk_probs[i].item()
-        # Build the full sequence with this token substituted
-        filled_ids = input_ids[0].clone()
-        filled_ids[mask_pos] = token_id
-        sequence = tokenizer.decode(filled_ids, skip_special_tokens=True).strip()
-        results.append({
-            "token": token,
-            "score": round(prob, 4),
-            "sequence": sequence,
-        })
-
-    best = results[0]["sequence"] if results else ""
-
-    return {"predictions": results}, best
-
-
+        logits = model(**enc).logits
+    probs = F.softmax(logits[0, mp].float(), dim=-1)
+    k = min(int(top_k), probs.shape[-1])
+    tp, ti = torch.topk(probs, k)
+    toks = [
+        tokenizer.decode([i], clean_up_tokenization_spaces=False).strip()
+        for i in ti.tolist()
+    ]
+    scores = [float(x) for x in tp.tolist()]
+    return _render(text, toks, scores)
+
+
+# Palette matches the DoctoBERT Diffusion demo (dark card, accent gradient bars).
 CSS = """
-#col-container { max-width: 900px; margin: 0 auto; }
-.dark .gradio-container { color: var(--body-text-color); }
+#col { max-width: 900px; margin: 0 auto; }
+.fm {
+  --fg:#e6edf6; --dim:#7d8aa0; --line:#1e2636; --accent:#7aa2ff; --accent2:#b78bff;
+  background: linear-gradient(180deg,#0e1420,#0b0f17); border:1px solid var(--line);
+  border-radius:14px; padding:18px 20px; color:var(--fg);
+  font-family: ui-monospace,"JetBrains Mono",Menlo,Consolas,monospace;
+}
+.fm-sentence { font-size:18px; line-height:1.9; margin-bottom:14px; }
+/* Gradio's own text color otherwise wins on the bare text nodes; force ours.
+   Blue accent matches the DoctoBERT Diffusion demo's prompt text. */
+.fm-txt { color: var(--accent) !important; }
+.fm-slot { background:linear-gradient(90deg,var(--accent),var(--accent2)); color:#0b0f17;
+  font-weight:600; padding:1px 9px; border-radius:6px; }
+.fm-cap { color:var(--dim) !important; font-size:12px; text-transform:uppercase;
+  letter-spacing:.08em; margin:6px 0 10px; }
+.fm-row { display:flex; align-items:center; gap:12px; margin:7px 0; font-size:14px; }
+.fm-rank { color:var(--dim) !important; width:1.5em; text-align:right; }
+/* Token + probability are the actual content, not secondary metadata -- keep
+   them fully bright (and !important, like .fm-txt) regardless of bar size. */
+.fm-tok { min-width:130px; color:var(--fg) !important; overflow:hidden;
+  text-overflow:ellipsis; white-space:nowrap; }
+.fm-bar { flex:1; height:10px; background:#182034; border-radius:6px; overflow:hidden; }
+.fm-bar>i { display:block; height:100%; width:var(--w);
+  background:linear-gradient(90deg,var(--accent),var(--accent2)); border-radius:6px;
+  animation: fmgrow .6s cubic-bezier(.2,.8,.2,1); }
+@keyframes fmgrow { from { width:0 } }
+.fm-pct { color:var(--fg) !important; width:4.5em; text-align:right; font-weight:600; }
+.fm-err { color:#ff6b6b !important; }
 """
 
-with gr.Blocks(theme=gr.themes.Citrus(), css=CSS) as demo:
-    with gr.Column(elem_id="col-container"):
+EXAMPLES = [
+    f"Le patient présente une douleur {MASK} aiguë.",
+    f"Le médecin a prescrit un {MASK} pour traiter l'infection.",
+    f"Une douleur thoracique irradiant vers le bras gauche évoque un {MASK} du myocarde.",
+    f"Les analyses sanguines ont révélé un taux élevé de {MASK}.",
+]
+
+with gr.Blocks(title="DoctoBERT Fill-Mask") as demo:
+    with gr.Column(elem_id="col"):
         gr.Markdown(
-            """
-            # 🏥 DoctoBERT-fr-base
-            **DoctoBERT-fr-base** est un encodeur médical français (RoBERTa, 111M de paramètres)
-            pré-entraîné sur des données biomédicales et cliniques françaises.
-            Démo de **fill-mask** : donnez une phrase médicale contenant `<mask>`
-            et le modèle prédit les jetons les plus probables.
-
-            🔗 [Modèle sur le Hub](https://huggingface.co/doctolib-lab/doctobert-fr-base)
-            """
+            "# 🩺 DoctoBERT Fill-Mask\n"
+            "Fill-mask demo of "
+            f"[doctolib-lab/doctomodernbert-fr-base](https://huggingface.co/{MODEL_ID}), "
+            "a French medical ModernBERT encoder.\n\n"
+            f"Write a sentence with a `{MASK}` token and the model predicts the most likely "
+            "words, ranked by probability."
         )
-
         with gr.Row():
-            text_input = gr.Textbox(
-                label="Phrase médicale (avec <mask>)",
-                placeholder=f"Le patient souffre d'une {MASK_TOKEN_DISPLAY} aiguë.",
+            inp = gr.Textbox(
+                label=f"Sentence (with {MASK})",
+                value=EXAMPLES[0],
+                placeholder=f"Le patient souffre d'une {MASK} aiguë.",
                 scale=4,
             )
-            run_btn = gr.Button("Prédire", variant="primary", scale=1)
-
-        with gr.Accordion("Paramètres avancés", open=False):
-            top_k_input = gr.Slider(
-                minimum=1, maximum=20, value=5, step=1,
-                label="Nombre de prédictions (top-k)",
-            )
-
-        best_pred_output = gr.Textbox(label="Meilleure prédiction", interactive=False)
-        predictions_output = gr.JSON(label="Top-k prédictions")
-
-        run_btn.click(
-            fn=predict_mask,
-            inputs=[text_input, top_k_input],
-            outputs=[predictions_output, best_pred_output],
-            api_name="predict",
-        )
-
-        gr.Examples(
-            examples=[
-                [f"Le patient souffre d'une {MASK_TOKEN_DISPLAY} aiguë.", 5],
-                [f"Le médecin prescrit un traitement pour l'{MASK_TOKEN_DISPLAY}.", 5],
-                [f"La dose recommandée est de {MASK_TOKEN_DISPLAY} mg par jour.", 5],
-                [f"Le patient présente des symptômes de {MASK_TOKEN_DISPLAY} chronique.", 5],
-                [f"L'examen clinique révèle une {MASK_TOKEN_DISPLAY} au niveau du thorax.", 5],
-            ],
-            inputs=[text_input, top_k_input],
-            outputs=[predictions_output, best_pred_output],
-            fn=predict_mask,
-            cache_examples=True,
-            cache_mode="lazy",
-        )
+            btn = gr.Button("Predict", variant="primary", scale=1)
+        with gr.Accordion("Options", open=False):
+            topk = gr.Slider(1, 20, value=8, step=1, label="Top-k predictions")
+        out = gr.HTML()
+        gr.Examples(EXAMPLES, inputs=inp)
+        btn.click(predict, [inp, topk], out)
+        inp.submit(predict, [inp, topk], out)
 
 if __name__ == "__main__":
-    demo.launch(mcp_server=True)
+    demo.launch(theme=gr.themes.Soft(), css=CSS, show_error=True)