app.py

import os
import torch
import pandas as pd
import gradio as gr
from transformers import AutoTokenizer, AutoModelForSequenceClassification

# ---------------------------------------------------------------------------
# Models in the family — all built on the from-scratch `mist-encoder-base-ng`.
#   base  : olaverse/mist-encoder-base-ng   (encoder, not demoed directly)
#   LID   : olaverse/lid-neural-5.1         (4-way Nigerian language ID)
#   embed : olaverse/naija-embed-base       (cross-lingual sentence embeddings)
# Repos must be public, or set an HF_TOKEN secret in the Space with read access.
# requirements.txt: gradio, transformers, torch, sentence-transformers, pandas
# ---------------------------------------------------------------------------
LID_ID   = "olaverse/lid-neural-5.1"
EMBED_ID = "olaverse/naija-embed-base"
BASE_ID  = "olaverse/mist-encoder-base-ng"
TOKEN = os.environ.get("HF_TOKEN")

# --- Language ID model (eager load; it's the primary feature) ---------------
tok = AutoTokenizer.from_pretrained(LID_ID, token=TOKEN)
model = AutoModelForSequenceClassification.from_pretrained(
    LID_ID, attn_implementation="eager", token=TOKEN
)
model.eval()
id2label = model.config.id2label

# --- Embedder (lazy: Space still boots if sentence-transformers/weights miss)
_embed = {}
def _get_embedder():
    if "m" not in _embed:
        try:
            from sentence_transformers import SentenceTransformer
            _embed["m"] = SentenceTransformer(EMBED_ID, token=TOKEN)
        except Exception as e:
            _embed["m"] = None
            print(f"embedder unavailable: {e}")
    return _embed["m"]

def _embed_encode(m, texts):
    """Manual encode that is robust to the tokenizer emitting `token_type_ids`
    (ModernBERT's forward() rejects it). Returns L2-normalized embeddings.
    Accepts a single string or a list; always returns a 2D tensor."""
    single = isinstance(texts, str)
    batch = [texts] if single else list(texts)
    feats = m.tokenize(batch)                       # ST builds input_ids/attention_mask(/token_type_ids)
    feats.pop("token_type_ids", None)               # <- the fix: ModernBERT doesn't accept it
    dev = next(m.parameters()).device
    feats = {k: (v.to(dev) if hasattr(v, "to") else v) for k, v in feats.items()}
    emb = m.forward(feats)["sentence_embedding"]
    emb = torch.nn.functional.normalize(emb, p=2, dim=1)
    return emb

# --- Optional diacritizers (character-level BiLSTMs via the olaverse SDK) ----
_diac = {}
def _get_diacritizer(name):
    if name not in _diac:
        try:
            from olaverse.nlp.diacritizer import Diacritizer
            _diac[name] = Diacritizer(model=name)
        except Exception as e:
            _diac[name] = None
            print(f"diacritizer {name} unavailable: {e}")
    return _diac[name]


# ===========================================================================
# Tab 1 — Language identification
# ===========================================================================
@torch.no_grad()
def _classify(text):
    enc = tok(text, return_tensors="pt", truncation=True, max_length=128)
    enc.pop("token_type_ids", None)                 # base is ModernBERT too — be safe
    probs = model(**enc).logits.softmax(-1)[0].tolist()
    return {id2label[i]: float(p) for i, p in enumerate(probs)}

def identify(text, restore):
    text = (text or "").strip()
    if not text:
        return {}, ""
    shown = text
    if restore and restore != "None":
        name = {"Yor\u00f9b\u00e1": "diacnet-yor", "Igbo": "diacnet-ig"}[restore]
        d = _get_diacritizer(name)
        if d is not None:
            try:
                shown = d.restore(text)
            except Exception as e:
                shown = text + f"   (diacritizer error: {e})"
        else:
            shown = text  # SDK unavailable on this Space — classify as-is
    note = "" if shown == text else f"**Restored:** {shown}"
    return _classify(shown), note


# ===========================================================================
# Tab 2 — Semantic search with naija-embed-base (cross-lingual retrieval)
# ===========================================================================
@torch.no_grad()
def semantic_search(query, candidates_text):
    m = _get_embedder()
    if m is None:
        return pd.DataFrame({"candidate": ["embedder unavailable — check Space logs / requirements.txt"],
                             "cosine": [None]})
    query = (query or "").strip()
    cands = [c.strip() for c in (candidates_text or "").splitlines() if c.strip()]
    if not query or not cands:
        return pd.DataFrame({"candidate": [], "cosine": []})
    try:
        q = _embed_encode(m, query)                 # (1, d)
        C = _embed_encode(m, cands)                 # (n, d)
        sims = (C @ q.T).squeeze(1).tolist()
    except Exception as e:
        return pd.DataFrame({"candidate": [f"encode error: {str(e)[:80]}"], "cosine": [None]})
    ranked = sorted(zip(cands, sims), key=lambda x: -x[1])
    return pd.DataFrame({"candidate": [c for c, _ in ranked],
                         "cosine":    [round(float(s), 3) for _, s in ranked]})


# ===========================================================================
# UI
# ===========================================================================
lid_examples = [
    ["Ojo lo si oja lana", "Yor\u00f9b\u00e1"],                       # un-accented Yoruba -> restore
    ["B\u00e1wo ni? \u1e62\u00e9 d\u00e1ad\u00e1a ni?", "None"],       # Yoruba
    ["Ina kwana? Yaya gida?", "None"],                                 # Hausa
    ["Kedu ka \u1ecb mere?", "None"],                                  # Igbo
    ["How you dey na? I dey kampe o.", "None"],                        # Pidgin
    ["Sannu da zuwa, yaya aiki?", "None"],                             # Hausa
    ["\u1eb8 k\u00fa \u00e0\u00e1r\u1ecd, \u1e63\u00e9 \u00e0l\u00e1\u00e0f\u00eda ni?", "None"],  # Yoruba (greeting)
    ["Daal\u1ee5, k\u00e8d\u1ee5 ihe \u1ecb na-eme?", "None"],         # Igbo
    ["Wetin dey happen for here?", "None"],                            # Pidgin
    ["Oga abeg make we waka go market", "None"],                       # Pidgin
]

# Cross-lingual examples: each query (one language) against mixed-language candidates.
embed_examples = [
    # Hausa query: "I like reading books" -> Yoruba/Igbo matches vs unrelated Hausa
    [
        "Ina son karatun littattafai",
        "Mo n\u00edf\u1eb9\u0301\u1eb9\u0301 k\u00edk\u00e0 \u00ecw\u00e9\n"
        "\u1ecc na-amas\u1ecb m \u1ecbg\u1ee5 akw\u1ee5kw\u1ecd\n"
        "Mota ya tsufa sosai\n"
        "Yau akwai ruwan sama",
    ],
    # Yoruba query: "food is very tasty" -> Hausa/Igbo food matches vs unrelated
    [
        "O\u0301unjE\u0323 yi\u00ed dun gan",
        "Abinci ya yi da\u0257i sosai\n"
        "Nri a t\u1ecdr\u1ecd \u1ee5t\u1ecd nke ukwuu\n"
        "Yara suna wasa a filin\n"
        "Akw\u1ee5kw\u1ecd a d\u1ecb \u1ecdh\u1ee5r\u1ee5",
    ],
    # Igbo query: "the weather is cold today" -> Hausa/Yoruba weather vs unrelated
    [
        "Ihu igwe j\u1ee5r\u1ee5 oyi taa",
        "Yau akwai sanyi sosai\n"
        "Oj\u00f9-\u1ecd\u0300run t\u00fa\u0300tu\u0300 l\u00f3n\u00ec\n"
        "Mo f\u1eb9\u0301r\u00e0n orin yi\u00ed\n"
        "Ah\u1ecba na-ad\u1ecb oke \u1ecdn\u1ee5",
    ],
    # Hausa query: "I am going to the market" -> Pidgin/Yoruba/Igbo "market"
    [
        "Ina zuwa kasuwa yanzu",
        "I dey go market now\n"
        "Mo n\u00ed l\u1ecd s\u00ed \u1ecdj\u00e0\n"
        "Ana m na-aga ah\u1ecba\n"
        "Ruwan sama yana sauka",
    ],
    # Within-language (Hausa) sanity check: paraphrase vs unrelated
    [
        "Yara suna karatu a makaranta",
        "Da\u0301lib\u00e9\u0301 suna karatu a aji\n"
        "Malam yana koyar da \u0257alibai\n"
        "Mota ta \u0253aci a hanya\n"
        "Kifi yana iyo a ruwa",
    ],
]

with gr.Blocks(title="Naija NLP \u2014 models on mist-encoder-base-ng") as demo:
    gr.Markdown(
        "# \U0001f1f3\U0001f1ec Naija NLP \u2014 the `mist-encoder-base-ng` family\n"
        "A small suite of models built on the from-scratch **`olaverse/mist-encoder-base-ng`** "
        "encoder (ModernBERT architecture):\n\n"
        f"- **[`lid-neural-5.1`](https://huggingface.co/{LID_ID})** \u2014 language ID "
        "(Hausa, Yoruba, Igbo, Nigerian Pidgin).\n"
        f"- **[`naija-embed-base`](https://huggingface.co/{EMBED_ID})** \u2014 cross-lingual "
        "sentence embeddings for semantic search, clustering, and RAG (Hausa, Yoruba, Igbo).\n"
        f"- **[`mist-encoder-base-ng`](https://huggingface.co/{BASE_ID})** \u2014 the shared base "
        "encoder everything is fine-tuned from.\n"
    )

    with gr.Tabs():
        # ---- Tab 1: Language ID --------------------------------------------
        with gr.Tab("Language ID"):
            gr.Markdown(
                "Detects **Hausa, Yoruba, Igbo, or Nigerian Pidgin** from short text. "
                "For **un-accented Yor\u00f9b\u00e1/Igbo**, turn on diacritic restoration "
                "(`diacnet-*`) \u2014 it puts tone marks and dots back before classifying.\n\n"
                "\u26a0\ufe0f No English/other class \u2014 non-Nigerian text is (confidently) "
                "mislabeled, usually as Pidgin."
            )
            with gr.Row():
                with gr.Column():
                    txt = gr.Textbox(lines=3, label="Text",
                                     placeholder="Type Hausa, Yoruba, Igbo, or Nigerian Pidgin\u2026")
                    restore = gr.Radio(["None", "Yor\u00f9b\u00e1", "Igbo"], value="None",
                                       label="Restore diacritics first (for un-accented input)")
                    btn = gr.Button("Identify", variant="primary")
                with gr.Column():
                    out = gr.Label(num_top_classes=4, label="Predicted language")
                    restored = gr.Markdown()
            btn.click(identify, [txt, restore], [out, restored])
            txt.submit(identify, [txt, restore], [out, restored])
            gr.Examples(lid_examples, [txt, restore])

        # ---- Tab 2: Semantic search ----------------------------------------
        with gr.Tab("Semantic search"):
            gr.Markdown(
                "Rank candidate sentences by meaning against a query, using "
                "`naija-embed-base`. Works **within a language** and **across** Hausa, "
                "Yoruba, and Igbo \u2014 e.g. a Hausa query can retrieve a Yoruba sentence "
                "about the same thing.\n\n"
                "Put one candidate per line. Scores are cosine similarity (higher = closer).\n\n"
                "\u2139\ufe0f Cross-lingual alignment is real but modest (FLORES Hausa\u2192Yoruba "
                "\u22480.67); within-language is stronger. Nigerian Pidgin isn't a trained "
                "language for this embedder."
            )
            with gr.Row():
                with gr.Column():
                    q = gr.Textbox(lines=2, label="Query",
                                   placeholder="A sentence to search with\u2026")
                    cands = gr.Textbox(lines=6, label="Candidates (one per line)",
                                       placeholder="Sentence 1\nSentence 2\nSentence 3\u2026")
                    sbtn = gr.Button("Search", variant="primary")
                with gr.Column():
                    results = gr.Dataframe(label="Ranked by similarity", wrap=True)
            sbtn.click(semantic_search, [q, cands], results)
            gr.Examples(embed_examples, [q, cands])


if __name__ == "__main__":
    demo.launch()