LFM2.5-Encoder-350M-Spellchecker

LFM2.5-Encoder-350M-Spellchecker — LiteRT

LiquidAI/LFM2.5-Encoder-350M-Spellchecker converted to LiteRT (.tflite) for on-device inference. A GECToR-style two-head tagger that corrects misspellings and grammar token by token, fully offline (demo Space).

Model description

One signature, gec_128 (S = 128, batch 1, right-padded; the base model's own decode also uses max_len 128). input_ids int32 [1, 128] — the tokenizer prepends <|startoftext|>, which the model uses as the sentence anchor — and attention_mask int32 [1, 128]. Two outputs, both zeroed at padded positions:

The tag space is 0 = $KEEP, 1 = $DELETE, 2 … 2+V = $REPLACE_<piece>, 2+V … = $APPEND_<piece>, with V = 64400 BPE pieces. For a $REPLACE/$APPEND tag the piece id is the tag minus its base — i.e. tag 7393 means "replace with vocabulary id 7391", which is Ġgoes.

Decoding is the base repo's algorithm: argmax the tags, gate them by softmax(detect)[1] >= min_error_prob, apply the surviving edits, and repeat (at most 3 passes) until the text stops changing. The base repo also bundles an optional PyTorch reranker for its published maximum-precision operating point; that stays host-side on desktop. This artifact covers the tagger, which is a fully supported mode of the base model's .correct().

How to use

1. Install dependencies

pip install ai-edge-litert numpy tokenizers huggingface_hub

2. Save the script below as spellcheck.py:

#!/usr/bin/env python3
"""Correct text with litert-community/LFM2.5-Encoder-350M-Spellchecker."""
import argparse

import numpy as np
from ai_edge_litert.interpreter import Interpreter
from huggingface_hub import hf_hub_download
from tokenizers import Tokenizer

REPO = "litert-community/LFM2.5-Encoder-350M-Spellchecker"
SEQ_LEN = 128
VOCAB = 64400  # $REPLACE_<piece> occupies tags 2..2+VOCAB, $APPEND_<piece> the rest


def correct_once(ids, runner, min_error_prob):
    """One tagging pass. Returns the edited id list and whether anything changed."""
    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

    out = runner(input_ids=input_ids, attention_mask=attention_mask)
    label_logits, detect_logits = out["output_0"][0], out["output_1"][0]

    edits = []
    for t in range(len(ids)):
        scores = detect_logits[t]
        error_prob = np.exp(scores[1] - scores.max()) / np.exp(scores - scores.max()).sum()
        if error_prob < min_error_prob:
            continue
        tag = int(label_logits[t].argmax())
        if tag == 0:  # $KEEP
            continue
        edits.append((t, tag))

    edited = list(ids)
    for t, tag in reversed(edits):  # right-to-left keeps earlier indices valid
        if tag == 1:  # $DELETE
            del edited[t]
        elif tag < 2 + VOCAB:  # $REPLACE_<piece>
            edited[t] = tag - 2
        else:  # $APPEND_<piece>
            edited.insert(t + 1, tag - 2 - VOCAB)
    return edited, bool(edits)


def main():
    parser = argparse.ArgumentParser()
    parser.add_argument("--text", required=True, help="Text to correct.")
    parser.add_argument("--min-error-prob", type=float, default=0.5)
    parser.add_argument("--max-passes", type=int, default=3)
    args = parser.parse_args()

    model_path = hf_hub_download(REPO, "LFM2.5-Encoder-350M-Spellchecker_wi8fc.tflite")
    tokenizer = Tokenizer.from_file(hf_hub_download(REPO, "tokenizer.json"))

    ids = tokenizer.encode(args.text).ids  # the tokenizer prepends the BOS anchor
    if len(ids) > SEQ_LEN:
        raise SystemExit(f"{len(ids)} tokens exceed the {SEQ_LEN}-token window")

    interpreter = Interpreter(model_path=model_path)
    runner = interpreter.get_signature_runner("gec_128")

    for _ in range(args.max_passes):
        ids, changed = correct_once(ids, runner, args.min_error_prob)
        if not changed:
            break

    print(tokenizer.decode(ids).strip())


if __name__ == "__main__":
    main()

3. Run it

python spellcheck.py --text "I has recieved you're mesage yesterday and will responde soon."

I have received your message yesterday and will respond soon.

A sentence with nothing to fix comes back unchanged. On Android/iOS use the LiteRT runtime's SignatureRunner APIs with the same signature name; the tokenizer is the standard Hugging Face tokenizer.json.

Performance

One gec_128 pass with the int8 (wi8fc) file, CPU only.

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 per output head — both heads measured 64 ms — 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 299 ms against a 49.8 ms steady state. Model load itself was 0.28 s on the iPhone, with a peak footprint of 746 MiB.

Correction is iterative, so a three-pass correction runs the graph three times. The signature is fixed-shape, so input language or content does not change the per-pass time.

Accuracy note

Task-level parity against the PyTorch reference on "She go to school every day ." — a single $REPLACE on "go": fp32, fp16 and int8 all produce the identical edit, at the same position, with the same replacement piece and an agreeing detect head. That is a single-sentence spot check, not a benchmark over a labelled corpus.

On the iPhone 17 Pro the int8 file reproduces the desktop outputs bit-exactly on both heads — including the full [1, 128, 128802] label tensor — at cosine 1.000000, max absolute difference 0.0.

Android (Pixel 8a)

Android figures use the standard TFLite benchmark_model on a Pixel 8a (Tensor G3, Android 16) — 5 warm-up runs then 20 timed runs, the signature selected explicitly with --signature_to_run_for, CPU at 4 threads.

GPU status (2026-08-13 re-export): still CPU on mobile. The re-export respells the one idiom mobile GPU delegates refuse — transformers' rank-5 repeat_kv expand — into an equivalent rank-4 matmul (outputs bitwise-identical on CPU) — removing the family-wide GPU blocker — but this model's int8 file still does not compile on mobile GPU delegates: its 128802-token tied-vocabulary head hits a runtime kernel limit ("failed to initialize kernel"), consistently on Metal and OpenCL. CPU remains the mobile path for the int8 file (bit-exact on device). The fp16 file does run under the GPU delegates with the head falling back to CPU and matches the fp32 reference (cosine 1.000000, desktop-verified). It is also the only file here that reached a mobile accelerator: AOT-compiled for the Hexagon it runs on a Galaxy S26 NPU at 82.96 ms (see Snapdragon NPU (Hexagon) below), while the int8 file produced no usable row on either S26 accelerator.

Snapdragon NPU (Hexagon)

• LFM2.5-Encoder-350M-Spellchecker_fp16.tflite — the NPU runs it at 82.96 ms. The GPU does not — LiteRtException: Failed to compile model.
• LFM2.5-Encoder-350M-Spellchecker_wi8fc.tflite — neither accelerator produced a usable row on the S26. NPU: the graph compiles and then fails to run (LiteRtException: Failed to invoke the compiled model). GPU: LiteRtException: Failed to compile model.

Measured on a Samsung Galaxy S26 (Snapdragon 8 Elite Gen 5 / SM8850, Hexagon v81, Android 16) with LiteRT CompiledModel 2.2.0, one accelerator per process, 5 warm-up runs then N=50 timed runs, median reported. The run held thermal status NONE throughout. Headroom 0.76–0.77, where 1.0 is the throttling threshold.

The NPU row marked AOT ran an artifact compiled ahead of time for SM8850 (ai-edge-litert 2.2.0 + QAIRT 2.47.0), not the published file. That artifact is not distributed here; the compile is one command in the NPU guide.

GPU wiring: GPU guide.

License

LFM Open License v1.0 (see LICENSE, unchanged from the base model). Note the license's commercial-use threshold (Section 5). This repository redistributes converted Derivative Works of LiquidAI/LFM2.5-Encoder-350M-Spellchecker with modification notices per Section 4; all credit for the model to Liquid AI.