"""GLiNER2.5-Decide on LiteRT with every trained head: gliner2's own schema handling and decoders over LiteRT graphs. gliner2's span architecture runs five trained modules: the encoder, `span_rep`, `count_pred`, `count_embed` and the classifier. This runtime is gliner2's SpanExtractor with those five replaced by LiteRT graphs; everything else (schema building, the collator, span scoring, decoding, overlap rules, chunking for long text) is the library's own code. from litert_span_runtime import LiteRTSpanExtractor from gliner2 import Schema model = LiteRTSpanExtractor("path/to/bundle") model.classify_text("Please cancel my order.", {"intent": ["cancel", "refund", "track"]}) model.extract_entities("Alice founded Acme in Toronto.", ["person", "organization", "city"]) model.extract_relations("Alice founded Acme in Toronto.", ["founded", "located_in"]) Every public call of gliner2's extractor works (extract, batch_extract, extract_entities, classify_text, extract_json, extract_relations and their batch_* and *_long forms). A request over 512 tokens is refused with a ValueError, never truncated: use the *_long calls with a chunk_size that fits, e.g. extract_entities_long(text, labels, chunk_size=300). Graphs: the encoder file's `e{N}` signatures (hidden states for N = 128, 160, 256, 512 tokens; requests over 512 tokens are refused, use the library's *_long calls to chunk), and the heads file's `w{W}` (span representations), `count`, `c{C}_m{M}` (count embeddings) and `k{K}` (classifier rows). The span file (`span=`: the same `w{W}`, `count` and `k{K}`, built so LiteRT's GPU delegate compiles them) and the count file (`count=`: `c{C}_m{M}`) stand in for the heads file; with both, the heads file is not loaded: model = LiteRTSpanExtractor("path/to/bundle", span="gliner25_decide_span_wfp16.tflite", count="gliner25_decide_count_wfp16.tflite") Needs gliner2==2.0.0 (no PyTorch weights are loaded), torch (the library's tensor types), numpy and ai-edge-litert. The bundle folder holds the graphs, the float16 token table, the tokenizer files and source/config.json (the checkpoint's configuration). """ from __future__ import annotations import json import re from pathlib import Path from types import SimpleNamespace import numpy as np import torch from gliner2.inference.engine import SpanExtractor from gliner2.models.base import load_extractor_tokenizer from gliner2.processor import SchemaTransformer TABLE = "gliner25_decide_embeddings_f16.bin" class Model: """A LiteRT CompiledModel with every signature callable by name: model(signature, feed) -> {output: float32 array}.""" def __init__(self, path: Path, threads: int = 4, accelerator: str = "cpu", cache: str = ""): from ai_edge_litert.compiled_model import CompiledModel, CpuOptions, GpuOptions, HardwareAccelerator, Options if accelerator == "cpu": options = Options(HardwareAccelerator.CPU, cpu_options=CpuOptions(num_threads=threads, xnnpack_weight_cache_path=cache)) else: # "gpu" (FP32) or "gpu-fp16" options = Options(HardwareAccelerator.GPU, gpu_options=GpuOptions(enforce_f32=accelerator == "gpu", constant_tensor_sharing=True)) self.model = CompiledModel.from_file(str(path), options=options) self.signatures, self.buffers = {}, {} for index in range(len(self.model.get_signature_list())): sig = self.model.get_signature_by_index(index) key = sig["key"] details = self.model.get_input_tensor_details(key) self.signatures[key] = { "index": index, "inputs": sig["inputs"], "outputs": sig["outputs"], "in_shapes": {k: tuple(int(x) for x in v["shape"]) for k, v in details.items()}, "out_shapes": {k: tuple(int(x) for x in v["shape"]) for k, v in self.model.get_output_tensor_details(key).items()}, } def keys(self, pattern: str) -> list[tuple[tuple[int, ...], str]]: """Signatures whose name matches `pattern`, with the pattern's integer groups, sorted.""" found = [] for key in self.signatures: m = re.fullmatch(pattern, key) if m: found.append((tuple(int(g) for g in m.groups()), key)) return sorted(found) def __call__(self, key: str, feed: dict) -> dict: sig = self.signatures[key] if key not in self.buffers: self.buffers[key] = (self.model.create_input_buffers(sig["index"]), self.model.create_output_buffers(sig["index"])) ins, outs = self.buffers[key] for buffer, name in zip(ins, sig["inputs"]): buffer.write(np.ascontiguousarray(np.asarray(feed[name], np.float32))) self.model.run_by_index(sig["index"], ins, outs) return {name: buffer.read(int(np.prod(sig["out_shapes"][name])), np.float32).reshape(sig["out_shapes"][name]) for buffer, name in zip(outs, sig["outputs"])} def smallest(options: list[tuple[tuple[int, ...], str]], need: int, what: str) -> tuple[int, str]: for (size, *_), key in options: if size >= need: return size, key raise ValueError(f"{what} {need} exceeds the largest graph bucket {options[-1][0][0]}") class Encoder(torch.nn.Module): """The encoder: token ids -> hidden states, through the host-mapped float16 table and the `e{N}` signatures.""" def __init__(self, model: Model, table: np.ndarray): super().__init__() self.model, self.table = model, table self.buckets = model.keys(r"e(\d+)") def forward(self, input_ids, attention_mask, **_): batch, length = input_ids.shape out = torch.zeros(batch, length, self.table.shape[1]) for b in range(batch): n = int(attention_mask[b].sum()) size, key = smallest(self.buckets, n, "Request tokens") ids = np.zeros(size, np.int64) ids[:n] = input_ids[b, :n].numpy() mask = np.zeros((1, size), np.float32) mask[0, :n] = 1.0 hidden = self.model(key, {"inputs_embeds": self.table[ids].astype(np.float32)[None], "attention_mask": mask})["hidden_states"] out[b, :n] = torch.from_numpy(hidden[0, :n]) return SimpleNamespace(last_hidden_state=out) class SpanRep(torch.nn.Module): """span_rep: word states [B, L, D] -> span representations [B, L, max_width, D] (the library's span layout).""" def __init__(self, model: Model): super().__init__() self.model = model self.buckets = model.keys(r"w(\d+)") def forward(self, h, span_idx=None): batch, length, d = h.shape rows = [] for b in range(batch): size, key = smallest(self.buckets, length, "Text words") words = np.zeros((1, size, d), np.float32) words[0, :length] = h[b].numpy() rows.append(torch.from_numpy(self.model(key, {"words": words})["span_rep"][0, :length])) return torch.stack(rows, 0) class CountPred(torch.nn.Module): def __init__(self, model: Model): super().__init__() self.model = model def forward(self, prompt): return torch.from_numpy(self.model("count", {"prompt": prompt.numpy()})["count_logits"]) class CountEmbed(torch.nn.Module): """count_embed(fields [M, D], count) -> [count, M, D]: the smallest `c{C}_m{M}` that holds both, sliced.""" max_count = 20 def __init__(self, model: Model): super().__init__() self.model = model self.buckets = model.keys(r"c(\d+)_m(\d+)") def forward(self, fields, gold_count_val: int): count, m = min(int(gold_count_val), self.max_count), fields.shape[0] for (steps, width), key in sorted(self.buckets, key=lambda x: (x[0][0] * x[0][1], x[0])): if steps >= count and width >= m: padded = np.zeros((width, fields.shape[1]), np.float32) padded[:m] = fields.numpy() return torch.from_numpy(self.model(key, {"fields": padded})["struct_proj"][:count, :m]) raise ValueError(f"{m} fields with count {count} exceed the largest count graph") class Classifier(torch.nn.Module): """classifier(states [n, D]) -> [n, 1], in chunks of the largest `k{K}` signature (the classifier is row-wise).""" def __init__(self, model: Model): super().__init__() self.model = model self.buckets = model.keys(r"k(\d+)") def forward(self, states): n, d = states.shape out, start = [], 0 while start < n: rows = min(n - start, self.buckets[-1][0][0]) size, key = smallest(self.buckets, rows, "Classifier rows") padded = np.zeros((size, d), np.float32) padded[:rows] = states[start:start + rows].numpy() out.append(torch.from_numpy(self.model(key, {"states": padded})["logits"][:rows])) start += rows return torch.cat(out).reshape(n, 1) class LiteRTSpanExtractor(SpanExtractor): """gliner2's SpanExtractor (schemas, collator, scoring and decoders) with every trained module on LiteRT.""" def __init__(self, bundle: str, encoder: str = "gliner25_decide_wfp16.tflite", heads: str = "gliner25_decide_heads_wfp16.tflite", accelerator: str = "cpu", threads: int = 4, cache: str = "", span: str | None = None, count: str | None = None): torch.nn.Module.__init__(self) folder = Path(bundle) config = json.loads((folder / "source" / "config.json").read_text()) self.config = SimpleNamespace(**config) self.max_width = config["max_width"] self.processor = SchemaTransformer(tokenizer=load_extractor_tokenizer(str(folder)), token_pooling=config.get("token_pooling", "first")) self.strict_extraction = True self._device_anchor = torch.nn.Parameter(torch.zeros(1), requires_grad=False) # the runtime reads device and dtype from parameters() encoder_model = Model(folder / encoder, threads, accelerator, cache) # The span file (span_rep, count prediction, classifier) and the count file (count_embed) stand in for the heads file. heads_model = None if span and count else Model(folder / heads, threads) span_model = Model(folder / span, threads) if span else heads_model count_model = Model(folder / count, threads) if count else heads_model hidden = encoder_model.signatures[encoder_model.keys(r"e(\d+)")[0][1]]["in_shapes"]["inputs_embeds"][2] rows = (folder / TABLE).stat().st_size // (2 * hidden) table = np.memmap(folder / TABLE, dtype="