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

File	Recipe	Size	Target
LFM2.5-Encoder-350M-Spellchecker_wi8fc.tflite	int8 dynamic-range (linears + embedding + tied vocab heads, convs float)	429 MB	mobile + desktop
LFM2.5-Encoder-350M-Spellchecker_fp16.tflite	fp16 weights, float compute	847 MB	desktop — phone memory limits (XNNPACK per-signature fp32 unpacking)

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:

Output	Shape	Meaning
output_0 (label_logits)	float32 [1, 128, 128802]	per-token edit tag
output_1 (detect_logits)	float32 [1, 128, 2]	P(token is part of an error) gate

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.

Device	Threads	gec_128
Apple M4 Max (macOS)	8	49.8 ms
iPhone 17 Pro	6	64 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 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.

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.