Update app.py

app.py

@@ -1,29 +1,48 @@
 import gradio as gr
 from transformers import AutoTokenizer, AutoModelForSequenceClassification, pipeline
+from huggingface_hub import hf_hub_download
+import fasttext
 
 # ---------------------------------------------------------------------------
 # Model registry — all LID models published under olaverse.
-# Lazy-loaded on first use so the Space doesn't need to hold all four in
-# memory (or pay startup time for all four) before anyone's clicked anything.
+# Two different underlying formats:
+#   - "transformers": standard AutoTokenizer + AutoModelForSequenceClassification
+#   - "fasttext": a .bin file loaded directly with the fasttext library
+#     (lid-lite-25 is NOT a transformers model — it's a character-n-gram
+#     linear classifier, hence no tokenizer/config to load via AutoTokenizer)
+# Lazy-loaded on first use so the Space doesn't need to hold all of these in
+# memory (or pay startup time for all of them) before anyone's clicked anything.
 # ---------------------------------------------------------------------------
 MODELS = {
     "lid-neural-5 (Nigerian, 4 langs)": {
+        "type": "transformers",
         "repo": "olaverse/lid-neural-5",
         "note": "Yoruba, Hausa, Igbo, Nigerian Pidgin",
     },
     "lid-neural-5.1 (Nigerian, 4 langs, sentence-level)": {
+        "type": "transformers",
         "repo": "olaverse/lid-neural-5.1",
         "note": "Hausa, Yoruba, Igbo, Nigerian Pidgin — built on mist-encoder-base-ng, tuned for short/sentence-level text (97.6% acc)",
     },
-    "lid-lite-25 (fastText, 25 langs)": {
+    "lid-lite-25 (fastText, passages)": {
+        "type": "fasttext",
         "repo": "olaverse/lid-lite-25",
-        "note": "CPU-only, sub-ms inference — trained on both long passages and short queries",
+        "filename": "passages.bin",
+        "note": "25 languages, fastText character n-gram model, tuned for long-form passages",
+    },
+    "lid-lite-25 (fastText, short queries)": {
+        "type": "fasttext",
+        "repo": "olaverse/lid-lite-25",
+        "filename": "questions.bin",
+        "note": "25 languages, fastText character n-gram model, tuned for short questions/queries",
     },
     "lid-neural-25.1 (XLM-R, passages)": {
+        "type": "transformers",
         "repo": "olaverse/lid-neural-25.1",
         "note": "25 languages, tuned for long-form text",
     },
     "lid-neural-25.2 (XLM-R, short queries)": {
+        "type": "transformers",
         "repo": "olaverse/lid-neural-25.2",
         "note": "25 languages, tuned for short questions/queries",
     },
@@ -41,10 +60,29 @@ LANGUAGE_NAMES = {
 
 import gc
 
+
+class FastTextWrapper:
+    """Mimics the shape of a transformers text-classification pipeline output
+    ([[{'label': ..., 'score': ...}, ...]]) so downstream code doesn't need
+    to branch on model type."""
+
+    def __init__(self, ft_model):
+        self._model = ft_model
+
+    def __call__(self, text: str, top_k: int = 5):
+        clean = text.replace("\n", " ").strip()
+        labels, probs = self._model.predict(clean, k=top_k)
+        results = [
+            {"label": label.replace("__label__", ""), "score": float(prob)}
+            for label, prob in zip(labels, probs)
+        ]
+        return [results]
+
+
 # Only ONE model is kept in memory at a time. Loading a new one evicts
 # whatever was previously cached. This trades re-download/reload time on
 # every model switch for a flat, predictable memory footprint — important
-# on a free/quota-limited CPU Space running four separate checkpoints.
+# on a free/quota-limited CPU Space running several separate checkpoints.
 _CACHED_LABEL = None
 _CACHED_PIPELINE = None
 
@@ -67,10 +105,17 @@ def get_pipeline(model_label: str):
     # Evict whatever's currently loaded before loading the new one.
     _unload_current()
 
-    repo = MODELS[model_label]["repo"]
-    tok = AutoTokenizer.from_pretrained(repo)
-    model = AutoModelForSequenceClassification.from_pretrained(repo)
-    pipe = pipeline("text-classification", model=model, tokenizer=tok, top_k=5)
+    entry = MODELS[model_label]
+
+    if entry["type"] == "fasttext":
+        local_path = hf_hub_download(repo_id=entry["repo"], filename=entry["filename"])
+        ft_model = fasttext.load_model(local_path)
+        pipe = FastTextWrapper(ft_model)
+    else:
+        repo = entry["repo"]
+        tok = AutoTokenizer.from_pretrained(repo)
+        model = AutoModelForSequenceClassification.from_pretrained(repo)
+        pipe = pipeline("text-classification", model=model, tokenizer=tok, top_k=5)
 
     _CACHED_LABEL = model_label
     _CACHED_PIPELINE = pipe
@@ -173,9 +218,9 @@ with gr.Blocks() as demo:
 
     with gr.Tab("Compare all models"):
         gr.Markdown(
-            "Runs the same text through all five models and shows top-3 predictions from each. "
-            "Only one model is kept in memory at a time, so this reloads each checkpoint in turn — "
-            "expect it to take a few seconds longer than the single-model tab."
+            "Runs the same text through all six model/checkpoint combinations and shows top-3 "
+            "predictions from each. Only one model is kept in memory at a time, so this reloads "
+            "each checkpoint in turn — expect it to take a bit longer than the single-model tab."
         )
         compare_input = gr.Textbox(
             label="Enter text",
@@ -193,9 +238,9 @@ with gr.Blocks() as demo:
     **Model notes:**
     - `lid-neural-5` — Nigerian-focused, 4 languages (Yoruba, Hausa, Igbo, Pidgin)
     - `lid-neural-5.1` — Nigerian-focused, same 4 languages, sentence-level tuned on `mist-encoder-base-ng` (97.6% acc; most residual error involves Pidgin, which shares vocabulary with the others)
-    - `lid-lite-25` — fastText, CPU-only, 25 languages
+    - `lid-lite-25` — fastText (character n-gram, CPU-only), 25 languages, two checkpoints (passages / short queries)
     - `lid-neural-25.1` / `.2` — XLM-R fine-tunes, 25 languages, tuned for passages vs. short queries respectively.
-      Known limitation across both: Zulu/Xhosa confusion on short text (see model cards).
+      Known limitation across both `-25` families: Zulu/Xhosa confusion on short text (see model cards).
     """)
 
 if __name__ == "__main__":