feat: demo app

app.py

@@ -1,32 +1,46 @@
 """
 ModernBERTić Large - HF Space demo
-Three tabs: fill-mask, side-by-side vs BERTić, long-context fill-mask.
+Three tabs: fill-mask, side-by-side vs XLM-R, long-context fill-mask.
 """
 
+import os
+
 import gradio as gr
+import spaces
 import torch
 import torch.nn.functional as F
 from transformers import AutoTokenizer, AutoModelForMaskedLM
 
 MODEL_NAME = "permitt/galton-modernbertic-large"
-BASELINE_NAME = "classla/bcms-bertic"
+BASELINE_NAME = "FacebookAI/xlm-roberta-large"
 
-device = "cuda" if torch.cuda.is_available() else "cpu"
-dtype = torch.bfloat16 if device == "cuda" else torch.float32
+HF_TOKEN = os.environ.get("HF_TOKEN")
+if HF_TOKEN is None:
+    raise RuntimeError(
+        "HF_TOKEN secret not set. Add it under Space Settings -> Variables and secrets."
+    )
 
-tokenizer = AutoTokenizer.from_pretrained(MODEL_NAME)
-model = AutoModelForMaskedLM.from_pretrained(MODEL_NAME, torch_dtype=dtype).to(device).eval()
+# Load on CPU. ZeroGPU allocates GPU only inside @spaces.GPU functions.
+tokenizer = AutoTokenizer.from_pretrained(MODEL_NAME, token=HF_TOKEN)
+model = AutoModelForMaskedLM.from_pretrained(
+    MODEL_NAME, dtype=torch.bfloat16, token=HF_TOKEN
+).eval()
 
 baseline_tokenizer = AutoTokenizer.from_pretrained(BASELINE_NAME)
-baseline_model = AutoModelForMaskedLM.from_pretrained(BASELINE_NAME).to(device).eval()
+baseline_model = AutoModelForMaskedLM.from_pretrained(
+    BASELINE_NAME, dtype=torch.bfloat16
+).eval()
 
 OUR_MASK = tokenizer.mask_token
 THEIR_MASK = baseline_tokenizer.mask_token
 
 
 @torch.inference_mode()
-def fill_mask(text: str, tok, mdl, top_k: int = 5):
-    inputs = tok(text, return_tensors="pt", truncation=True, max_length=8192).to(device)
+def _run(text, tok, mdl, top_k=5):
+    mdl = mdl.to("cuda")
+    inputs = tok(
+        text, return_tensors="pt", truncation=True, max_length=8192
+    ).to("cuda")
     mask_id = tok.mask_token_id
     pos = (inputs.input_ids == mask_id).nonzero(as_tuple=True)
     if len(pos[1]) == 0:
@@ -35,25 +49,30 @@ def fill_mask(text: str, tok, mdl, top_k: int = 5):
     mask_logits = logits[0, pos[1][0]]
     probs = F.softmax(mask_logits.float(), dim=-1)
     top_probs, top_ids = probs.topk(top_k)
-    return [(tok.decode([tid]).strip(), float(p)) for tid, p in zip(top_ids, top_probs)]
+    return [
+        (tok.decode([tid]).strip(), float(p))
+        for tid, p in zip(top_ids, top_probs)
+    ]
 
 
 def fmt(preds):
     return "\n".join(f"{w:<20}  {p:.3f}" for w, p in preds)
 
 
-def predict_ours(text):
-    return fmt(fill_mask(text, tokenizer, model))
+@spaces.GPU
+def predict_ours(text: str) -> str:
+    return fmt(_run(text, tokenizer, model))
 
 
-def predict_compare(text):
-    ours = fill_mask(text, tokenizer, model)
-    bertic_text = text.replace(OUR_MASK, THEIR_MASK)
-    theirs = fill_mask(bertic_text, baseline_tokenizer, baseline_model)
+@spaces.GPU
+def predict_compare(text: str):
+    ours = _run(text, tokenizer, model)
+    their_text = text.replace(OUR_MASK, THEIR_MASK)
+    theirs = _run(their_text, baseline_tokenizer, baseline_model)
     return fmt(ours), fmt(theirs)
 
 
-with gr.Blocks(title="ModernBERTić Large", theme=gr.themes.Soft()) as demo:
+with gr.Blocks(title="ModernBERTić Large") as demo:
     gr.Markdown(
         f"""
         # ModernBERTić Large
@@ -75,14 +94,17 @@ with gr.Blocks(title="ModernBERTić Large", theme=gr.themes.Soft()) as demo:
                 f"Glavni grad Srbije je {OUR_MASK}.",
                 f"Najveći grad u Hrvatskoj je {OUR_MASK}.",
                 f"Pisac romana 'Na Drini ćuprija' je {OUR_MASK} Andrić.",
-                f"Главни град Србије је {OUR_MASK}.",  # cyrillic
+                f"Главни град Србије је {OUR_MASK}.",
             ],
             inputs=inp,
         )
         btn.click(predict_ours, inp, out)
 
-    with gr.Tab("vs BERTić"):
-        gr.Markdown("Same input, both models. ModernBERTić-large vs `classla/bcms-bertic`.")
+    with gr.Tab("vs XLM-R large"):
+        gr.Markdown(
+            "Same input, both models. ModernBERTić-large vs `xlm-roberta-large` "
+            "(the standard multilingual MLM baseline for BCMS)."
+        )
         inp2 = gr.Textbox(
             label="Input",
             value=f"Najveće jezero u Crnoj Gori je {OUR_MASK} jezero.",
@@ -91,13 +113,13 @@ with gr.Blocks(title="ModernBERTić Large", theme=gr.themes.Soft()) as demo:
         btn2 = gr.Button("Compare", variant="primary")
         with gr.Row():
             out_ours = gr.Textbox(label="ModernBERTić-large (ours)", lines=6)
-            out_theirs = gr.Textbox(label="BERTić (Ljubešić et al.)", lines=6)
+            out_theirs = gr.Textbox(label="XLM-R large", lines=6)
         btn2.click(predict_compare, inp2, [out_ours, out_theirs])
 
     with gr.Tab("Long context (8192)"):
         gr.Markdown(
             "Paste a long passage with one mask token deep in the text. "
-            "BERTić truncates at 512 tokens. ModernBERTić handles up to 8192."
+            "Most BCMS encoders truncate at 512 tokens. ModernBERTić handles up to 8192."
         )
         inp3 = gr.Textbox(
             label="Long input",
@@ -117,4 +139,4 @@ with gr.Blocks(title="ModernBERTić Large", theme=gr.themes.Soft()) as demo:
 
 
 if __name__ == "__main__":
-    demo.launch()
+    demo.launch(theme=gr.themes.Soft())