gradio app for masked modeling task

README.md

@@ -1,14 +1,68 @@
----
-title: Masked Language Modeling
-emoji: 🐢
-colorFrom: yellow
-colorTo: purple
-sdk: gradio
-sdk_version: 6.18.0
-python_version: '3.13'
-app_file: app.py
-pinned: false
-short_description: 'AfriBERT (Kenya adapted) Masked Language Modeling '
----
-
-Check out the configuration reference at https://huggingface.co/docs/hub/spaces-config-reference
+# AfriBERT Kenya Masked LM Gradio App
+
+Gradio demo for comparing masked-language-modeling predictions from:
+
+- Base model: `castorini/afriberta_large`
+- Adapted model: `Rogendo/afribert-kenya-adapted`
+
+The app uses the same tokenizer, `castorini/afriberta_large`, for both models so the MLM predictions are directly comparable.
+
+The app supports Swahili, Sheng, Kenyan institutional text, M-PESA language, and English-Swahili code-switching examples.
+
+## Run locally
+
+PyTorch does not currently install on Python 3.14. Use Python 3.10 for this app.
+
+```bash
+cd /Users/bitzsupport/Desktop/Portfoliio/afribert-kenya-mlm-gradio
+python3.10 -m venv venv
+source venv/bin/activate
+python -m pip install --upgrade pip
+pip install -r requirements.txt
+export HF_TOKEN="your_huggingface_read_token"
+python app.py
+```
+
+If `python3.10` is not installed on macOS:
+
+```bash
+brew install python@3.10
+```
+
+If the model is public, `HF_TOKEN` is optional. If it is private, the token must have read access.
+
+Optional overrides:
+
+```bash
+export MODEL_ID="Rogendo/afribert-kenya-adapted"
+export ADAPTED_MODEL_ID="Rogendo/afribert-kenya-adapted"
+export BASE_MODEL_ID="castorini/afriberta_large"
+export TOKENIZER_ID="castorini/afriberta_large"
+```
+
+## Hugging Face Space
+
+Create a Gradio Space and upload:
+
+- `app.py`
+- `requirements.txt`
+- `README.md`
+- `runtime.txt`
+
+Then add a Space secret named `HF_TOKEN` with a Hugging Face token that can read the model.
+
+## Usage
+
+Use the tokenizer mask token shown in the app: `<mask>`. `[MASK]` is also accepted and automatically converted.
+
+Examples:
+
+```text
+Tulifanya meeting jana na manager akasema <mask> itakuwa ready wiki ijayo.
+```
+
+```text
+Msee alikuwa poa sana, akanisaidia kupata <mask> ya ofisi.
+```
+
+The first output table compares the base and adapted model rank-by-rank. The second table shows each model's completed sentence for every prediction.

app.py

@@ -0,0 +1,192 @@
+import os
+from typing import Any
+
+import gradio as gr
+import torch
+from transformers import AutoModelForMaskedLM, AutoTokenizer
+
+
+ADAPTED_MODEL_ID = os.getenv("ADAPTED_MODEL_ID", os.getenv("MODEL_ID", "Rogendo/afribert-kenya-adapted"))
+BASE_MODEL_ID = os.getenv("BASE_MODEL_ID", "castorini/afriberta_large")
+TOKENIZER_ID = os.getenv("TOKENIZER_ID", "castorini/afriberta_large")
+HF_TOKEN = os.getenv("HF_TOKEN") or os.getenv("HUGGINGFACE_HUB_TOKEN")
+
+
+def load_models() -> tuple[Any, Any, Any, torch.device]:
+    device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
+    tokenizer = AutoTokenizer.from_pretrained(TOKENIZER_ID, token=HF_TOKEN, use_fast=False)
+    base_model = AutoModelForMaskedLM.from_pretrained(BASE_MODEL_ID, use_safetensors=True)
+    adapted_model = AutoModelForMaskedLM.from_pretrained(
+        ADAPTED_MODEL_ID,
+        token=HF_TOKEN,
+        use_safetensors=True,
+    )
+
+    base_model.to(device)
+    adapted_model.to(device)
+    base_model.eval()
+    adapted_model.eval()
+    return tokenizer, base_model, adapted_model, device
+
+
+tokenizer, base_model, adapted_model, device = load_models()
+MASK_TOKEN = tokenizer.mask_token or "[MASK]"
+
+
+EXAMPLES = [
+    f"Oya, twendeni zetu, kuna {MASK_TOKEN} flani ameniudhi.",
+    f"Tuma {MASK_TOKEN} kwa kutumia nambari ya simu kupitia huduma ya M-PESA.",
+    f"Mtoto aliripotiwa kwa ofisi ya {MASK_TOKEN} wa jamii baada ya kudhulumiwa nyumbani.",
+    f"Tulifanya meeting jana na manager akasema {MASK_TOKEN} itakuwa ready wiki ijayo.",
+    f"Msee alikuwa poa sana, akanisaidia kupata {MASK_TOKEN} ya ofisi.",
+]
+
+
+def normalize_input(text: str) -> str:
+    text = (text or "").strip()
+    if "[MASK]" in text and MASK_TOKEN != "[MASK]":
+        text = text.replace("[MASK]", MASK_TOKEN)
+    return text
+
+
+def model_predictions(model, inputs, mask_positions, top_k: int, model_label: str) -> list[list[Any]]:
+    with torch.no_grad():
+        outputs = model(**inputs)
+        logits = outputs.logits[0]
+
+    rows = []
+    for mask_index, position in enumerate(mask_positions.tolist(), start=1):
+        probabilities = torch.softmax(logits[position], dim=-1)
+        scores, token_ids = torch.topk(probabilities, k=int(top_k))
+
+        for rank, (score, token_id) in enumerate(zip(scores, token_ids), start=1):
+            token = tokenizer.decode([token_id.item()]).strip()
+            completed = inputs["input_ids"][0].clone()
+            completed[position] = token_id
+            sequence = tokenizer.decode(completed, skip_special_tokens=True)
+
+            rows.append([
+                model_label,
+                mask_index,
+                rank,
+                token,
+                round(float(score.item()), 4),
+                sequence,
+            ])
+
+    return rows
+
+
+def predict_masks(text: str, top_k: int) -> tuple[str, list[list[Any]], list[list[Any]]]:
+    text = normalize_input(text)
+
+    if not text:
+        return "Enter a sentence with a mask token.", [], []
+
+    if MASK_TOKEN not in text:
+        return f"Add at least one mask token: `{MASK_TOKEN}`", [], []
+
+    inputs = tokenizer(text, return_tensors="pt").to(device)
+    mask_positions = (inputs["input_ids"][0] == tokenizer.mask_token_id).nonzero(as_tuple=True)[0]
+
+    if len(mask_positions) == 0:
+        return f"No valid mask token found. Use `{MASK_TOKEN}`.", [], []
+
+    base_rows = model_predictions(base_model, inputs, mask_positions, top_k, "Base AfriBERT")
+    adapted_rows = model_predictions(adapted_model, inputs, mask_positions, top_k, "Adapted AfriBERT Kenya")
+    comparison_rows = []
+
+    for base_row, adapted_row in zip(base_rows, adapted_rows):
+        comparison_rows.append([
+            base_row[1],
+            base_row[2],
+            base_row[3],
+            base_row[4],
+            adapted_row[3],
+            adapted_row[4],
+        ])
+
+    summary = (
+        f"Base model: `{BASE_MODEL_ID}`\n\n"
+        f"Adapted model: `{ADAPTED_MODEL_ID}`\n\n"
+        f"Tokenizer: `{TOKENIZER_ID}`\n\n"
+        f"Mask token: `{MASK_TOKEN}`\n\n"
+        f"Found {len(mask_positions)} mask position{'s' if len(mask_positions) != 1 else ''}."
+    )
+    return summary, comparison_rows, base_rows + adapted_rows
+
+
+with gr.Blocks(title="AfriBERT Kenya Masked LM") as demo:
+    gr.Markdown(
+        """
+        # AfriBERT Kenya Masked Language Modeling
+
+        Compare base AfriBERT against the Kenya-adapted model on Swahili, Sheng,
+        Kenyan institutional text, M-PESA language, and English-Swahili code-switching.
+        """
+    )
+
+    with gr.Row():
+        with gr.Column(scale=2):
+            text_input = gr.Textbox(
+                label="Input text",
+                value=EXAMPLES[0],
+                lines=4,
+                placeholder=f"Type a sentence containing {MASK_TOKEN}",
+            )
+            top_k = gr.Slider(
+                label="Top predictions",
+                minimum=1,
+                maximum=10,
+                value=5,
+                step=1,
+            )
+            predict_button = gr.Button("Compare masked-token predictions", variant="primary")
+
+        with gr.Column(scale=1):
+            gr.Markdown(
+                f"""
+                **How to use**
+
+                Add `{MASK_TOKEN}` where you want the model to predict a token.
+                `[MASK]` is also accepted and converted automatically.
+
+                For private models, set `HF_TOKEN` before launching the app.
+                The same base AfriBERT tokenizer is used for both models.
+                """
+            )
+
+    summary_output = gr.Markdown()
+    comparison_output = gr.Dataframe(
+        headers=["Mask", "Rank", "Base prediction", "Base score", "Adapted prediction", "Adapted score"],
+        datatype=["number", "number", "str", "number", "str", "number"],
+        label="Side-by-side comparison",
+        wrap=True,
+    )
+    details_output = gr.Dataframe(
+        headers=["Model", "Mask", "Rank", "Prediction", "Score", "Completed sentence"],
+        datatype=["str", "number", "number", "str", "number", "str"],
+        label="Detailed predictions",
+        wrap=True,
+    )
+
+    gr.Examples(
+        examples=EXAMPLES,
+        inputs=text_input,
+    )
+
+    predict_button.click(
+        fn=predict_masks,
+        inputs=[text_input, top_k],
+        outputs=[summary_output, comparison_output, details_output],
+    )
+
+    text_input.submit(
+        fn=predict_masks,
+        inputs=[text_input, top_k],
+        outputs=[summary_output, comparison_output, details_output],
+    )
+
+
+if __name__ == "__main__":
+    demo.launch()

requirements.txt

@@ -0,0 +1,6 @@
+gradio>=4.44.0,<6
+numpy<2
+torch>=2.2.0
+transformers>=4.44.0,<5
+sentencepiece==0.1.99
+protobuf==3.20.3

runtime.txt

@@ -0,0 +1 @@
+python-3.10.13