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()