Card: measured Performance table (M4 Max + iPhone 17 Pro), verified runnable usage script, accuracy note

README.md

@@ -11,8 +11,8 @@ tags:
 - on-device
 - edge
 - encoder
-- zero-shot
 - routing
+- zero-shot
 - liquid
 - lfm2
 - lfm2.5
@@ -22,52 +22,153 @@ tags:
 
 [LiquidAI/LFM2.5-Encoder-350M-Prompt-Router](https://huggingface.co/LiquidAI/LFM2.5-Encoder-350M-Prompt-Router) converted to **LiteRT** (`.tflite`) for on-device inference. Zero-shot prompt routing: define your routing lanes as free text and the model scores the whole prompt against every lane in one CPU pass ([demo Space](https://huggingface.co/spaces/LiquidAI/prompt-routing)).
 
-| File | Recipe | Size | |
+## Model description
+
+| File | Recipe | Size | Target |
 |---|---|---|---|
-| `LFM2.5-Encoder-350M-Prompt-Router_wi8fc.tflite` | int8 dynamic-range (linears + embedding, convs float) | 365 MB | mobile + desktop (iPhone-verified bit-exact, 145 ms) |
+| `LFM2.5-Encoder-350M-Prompt-Router_wi8fc.tflite` | int8 dynamic-range (linears + embedding, convs float) | 365 MB | mobile + desktop |
 | `LFM2.5-Encoder-350M-Prompt-Router_fp16.tflite` | fp16 weights, float compute | 713 MB | desktop — phone memory limits (XNNPACK per-signature fp32 unpacking) |
 
-## Signatures
+Two signatures, `route_128` and `route_512` (S = 128 / 512, batch 1, right-padded, up to **8 lane slots**):
 
-`route_128` / `route_512` (S = 128 / 512, batch 1, right-padded, up to **8 lane slots**):
-
-| Input | Shape | |
+| Tensor | Shape | Meaning |
 |---|---|---|
-| `input_ids` | int32 `[1, S]` | prompt tokens: `Categories:\n- <lane 1>\n- <lane 2>…\n\nText:\n<prompt>` |
-| `attention_mask` | int32 `[1, S]` | 1 = token, 0 = pad |
-| `text_pool` | float32 `[1, 1, S]` | mean-pool weights over the prompt's text tokens (`1/n` each) |
-| `category_pool` | float32 `[1, 8, S]` | row r = mean-pool weights over lane r's tokens; unused lane rows all-zero |
+| `input_ids` | int32 `[1, S]` | `Categories:\n- <lane 1>\n- <lane 2>…\n\nText:\n<prompt>` |
+| `attention_mask` | int32 `[1, S]` | 1 = real token, 0 = pad |
+| `text_pool` | float32 `[1, 1, S]` | mean-pool weights over the prompt's own tokens (`1/n` each) |
+| `category_pool` | float32 `[1, 8, S]` | row *r* = mean-pool weights over lane *r*'s tokens; unused rows all-zero |
+| output | float32 `[1, 8]` | one logit per lane slot |
+
+Softmax over the first N (real) lanes only — an all-zero pool row produces a constant bias logit that must be ignored.
 
-Output: `logits` float32 `[1, 8]`. Softmax over the first N (real) lanes only — all-zero pool rows produce a constant bias logit that must be ignored.
+## How to use
 
-The pool matrices are built host-side from tokenizer character offsets, exactly like the base repo's `route()` helper:
+**1. Install dependencies**
+
+```bash
+pip install ai-edge-litert numpy tokenizers huggingface_hub
+```
+
+**2. Save the script** below as `route_prompt.py`:
 
 ```python
+#!/usr/bin/env python3
+"""Route a prompt to one of your lanes with litert-community/LFM2.5-Encoder-350M-Prompt-Router."""
+import argparse
+
 import numpy as np
+from ai_edge_litert.interpreter import Interpreter
+from huggingface_hub import hf_hub_download
 from tokenizers import Tokenizer
 
-def build_inputs(text, lanes, tok, S=512):
-    body = "\n".join(f"- {r}" for r in lanes)
+REPO = "litert-community/LFM2.5-Encoder-350M-Prompt-Router"
+MAX_LANES = 8
+
+
+def build_inputs(text, lanes, tokenizer, seq_len):
+    """Builds input_ids/attention_mask plus the two mean-pool matrices."""
+    body = "\n".join(f"- {lane}" for lane in lanes)
     prefix = f"Categories:\n{body}\n\nText:\n"
-    enc = tok.encode(prefix + text)
-    ids, offs = enc.ids, enc.offsets
-    x = np.zeros((1, S), np.int32); m = np.zeros((1, S), np.int32)
-    x[0, :len(ids)] = ids; m[0, :len(ids)] = 1
-    tp = np.zeros((1, 1, S), np.float32)
-    ti = [i for i, (a, b) in enumerate(offs) if b > len(prefix) and a != b]
-    tp[0, 0, ti] = 1 / len(ti)
-    cp = np.zeros((1, 8, S), np.float32)
+    encoding = tokenizer.encode(prefix + text)
+    ids, offsets = encoding.ids, encoding.offsets
+    if len(ids) > seq_len:
+        raise SystemExit(f"{len(ids)} tokens exceed --seq-len {seq_len}")
+
+    input_ids = np.zeros((1, seq_len), np.int32)
+    attention_mask = np.zeros((1, seq_len), np.int32)
+    input_ids[0, : len(ids)] = ids
+    attention_mask[0, : len(ids)] = 1
+
+    # Mean-pool over the document's own tokens.
+    text_pool = np.zeros((1, 1, seq_len), np.float32)
+    text_idx = [i for i, (a, b) in enumerate(offsets) if b > len(prefix) and a != b]
+    text_pool[0, 0, text_idx] = 1 / len(text_idx)
+
+    # Mean-pool over each lane's tokens; unused lane rows stay all-zero.
+    category_pool = np.zeros((1, MAX_LANES, seq_len), np.float32)
     pos = len("Categories:\n")
     for r, lane in enumerate(lanes):
-        a, b = pos + 2, pos + 2 + len(lane); pos = b + 1
-        idx = [i for i, (ta, tb) in enumerate(offs) if ta < b and tb > a and ta != tb]
-        cp[0, r, idx] = 1 / len(idx)
-    return {"input_ids": x, "attention_mask": m, "text_pool": tp, "category_pool": cp}
+        start, end = pos + 2, pos + 2 + len(lane)
+        pos = end + 1
+        idx = [i for i, (a, b) in enumerate(offsets) if a < end and b > start and a != b]
+        category_pool[0, r, idx] = 1 / len(idx)
+
+    return {
+        "input_ids": input_ids,
+        "attention_mask": attention_mask,
+        "text_pool": text_pool,
+        "category_pool": category_pool,
+    }
+
+
+def main():
+    parser = argparse.ArgumentParser()
+    parser.add_argument("--text", required=True, help="The prompt to route.")
+    parser.add_argument("--lane", action="append", required=True,
+                        help="A routing lane, repeatable (up to 8).")
+    parser.add_argument("--seq-len", type=int, default=512, choices=[128, 512])
+    args = parser.parse_args()
+    if len(args.lane) > MAX_LANES:
+        raise SystemExit(f"at most {MAX_LANES} lanes")
+
+    model_path = hf_hub_download(REPO, "LFM2.5-Encoder-350M-Prompt-Router_wi8fc.tflite")
+    tokenizer = Tokenizer.from_file(hf_hub_download(REPO, "tokenizer.json"))
+
+    feed = build_inputs(args.text, args.lane, tokenizer, args.seq_len)
+    interpreter = Interpreter(model_path=model_path)
+    runner = interpreter.get_signature_runner(f"route_{args.seq_len}")
+    logits = list(runner(**feed).values())[0][0]
+
+    # Softmax over the real lanes only — unused rows carry a constant bias logit.
+    real = logits[: len(args.lane)]
+    probs = np.exp(real - real.max())
+    probs /= probs.sum()
+    for lane, p in sorted(zip(args.lane, probs), key=lambda x: -x[1]):
+        print(f"{p:6.3f}  {lane}")
+
+
+if __name__ == "__main__":
+    main()
+```
+
+**3. Run it**
+
+```bash
+python route_prompt.py \
+  --text "My Python script throws a KeyError on a dict lookup, how do I fix it?" \
+  --lane "coding question" --lane "travel planning" \
+  --lane "medical advice" --lane "small talk"
+```
+
+```
+ 0.838  coding question
+ 0.054  small talk
+ 0.054  travel planning
+ 0.054  medical advice
 ```
 
-## Verification
+On Android/iOS use the LiteRT runtime's SignatureRunner APIs with the same signature names; the tokenizer is the standard Hugging Face `tokenizer.json`.
+
+## Performance
+
+One pass over a padded sequence with the int8 (`wi8fc`) file, CPU only.
+
+| Device | Threads | `route_128` | `route_512` |
+|---|---|---|---|
+| Apple M4 Max (macOS) | 8 | 34.5 ms | 112.3 ms |
+| iPhone 17 Pro | 6 | not measured | 145 ms |
+
+Mac figures are the median of 20 warm runs (`ai-edge-litert` 2.1.6, XNNPACK, otherwise idle machine). The iPhone figure comes from the on-device gate (TFLite C API + SignatureRunner + XNNPACK) and is a single run, not a median.
+
+**Budget for one slow first call.** The first inference after loading pays a one-time graph preparation: on the Mac it took 372 ms against a 34.5 ms steady state. Later signatures on the same loaded model do not pay it again — `route_512` measured 110 ms cold against 112 ms warm. Model load itself was 0.38 s on the iPhone, with a peak footprint of 649 MiB.
+
+One pass scores the prompt against all eight lane slots at once, so the cost does not grow with the number of lanes. The signatures are fixed-shape, so input language or content does not change the time.
+
+### Accuracy note
+
+Task-level parity against the PyTorch reference on the demo prompt with four lanes: fp32, fp16 **and int8 all reproduce the reference lane probabilities to four decimal places** — 0.838 for "coding question". That is a single-prompt spot check, not a benchmark over a labelled corpus.
 
-Task-level parity vs the PyTorch reference (demo prompt, 4 lanes): fp32, fp16 **and int8 all reproduce the reference lane probabilities to 4 decimal places** ("coding question" 0.838). On an iPhone 17 Pro the int8 file reproduces the desktop outputs **bit-exactly** (cosine 1.000000, max diff 0.0) at 145 ms per `route_512` pass (6 threads, XNNPACK).
+On the iPhone 17 Pro the int8 file reproduces the desktop outputs **bit-exactly** — cosine 1.000000, max absolute difference 0.0.
 
 ## License