GLiNER2.5-Decide LiteRT: fp16-weight graphs s128/s256/s512 (+fp32 references), fp16/fp32 embedding tables, host contract, Python runtime, conversion scripts, Android sample, card
d9af00e verified | """GLiNER2.5-Decide classification graph (DESIGN.md contract), one fixed shape per instance. | |
| Inputs: inputs_embeds [1,N,1024] f32 (raw word-embedding rows, pad id 0 included), | |
| attention_mask [1,N] f32, label_routing [1,L,N] f32 (one-hot rows at the [L] markers). | |
| Output: logits [1,1,1,L] f32 = classifier(hidden[pos_j]); rows beyond the label count are | |
| classifier(0), ignored by the host. | |
| Rewritten body = GLiNER2.5-Small's `ShapedDebertaLayer` (vendored verbatim in shaped_layer.py) | |
| over the 24 DeBERTa-v3-large layers. `native=True` runs the unmodified transformers | |
| `DebertaV2Model` on the same inputs, then the same routing matmul and classifier. | |
| """ | |
| from dataclasses import dataclass | |
| import torch | |
| from torch import nn | |
| from torch.nn import functional as F | |
| from common import HIDDEN, LABEL_SLOTS | |
| from shaped_layer import ShapedDebertaLayer, rank4_matmul | |
| INPUT_NAMES = ["inputs_embeds", "attention_mask", "label_routing"] | |
| OUTPUT_NAME = "logits" | |
| class Shape: | |
| seq: int | |
| labels: int = LABEL_SLOTS | |
| def contract(self): | |
| shapes = [(1, self.seq, HIDDEN), (1, self.seq), (1, self.labels, self.seq)] | |
| return dict(seq=self.seq, label_slots=self.labels, | |
| inputs=[dict(name=n, shape=list(s), dtype="float32") for n, s in zip(INPUT_NAMES, shapes)], | |
| output=dict(name=OUTPUT_NAME, shape=[1, 1, 1, self.labels], dtype="float32")) | |
| def label_positions(batch): | |
| """[L] marker positions in request order: task order, label order, each task's [P] dropped.""" | |
| return [p for positions in batch.schema_special_indices[0] for p in positions[1:]] | |
| def fixed_inputs(model, captured, shape): | |
| batch = captured["batch"] | |
| n = batch.input_ids.shape[1] | |
| positions = label_positions(batch) | |
| assert n <= shape.seq and len(positions) <= shape.labels, (n, len(positions), shape) | |
| ids = F.pad(batch.input_ids, (0, shape.seq - n), value=model.processor.tokenizer.pad_token_id) | |
| attention = F.pad(batch.attention_mask, (0, shape.seq - n)).float() | |
| routing = torch.zeros(1, shape.labels, shape.seq) | |
| for j, p in enumerate(positions): | |
| routing[0, j, p] = 1.0 | |
| with torch.no_grad(): | |
| embeds = model.encoder.embeddings.word_embeddings(ids) | |
| return tuple(x.detach().contiguous() for x in (embeds, attention, routing)) | |
| class DecideClassifier(nn.Module): | |
| def __init__(self, model, shape, native=False): | |
| super().__init__() | |
| self.shape = shape | |
| self.native = native | |
| self.classifier = model.classifier | |
| assert [type(m).__name__ for m in self.classifier] == ["Linear", "ReLU", "Linear"] | |
| emb = model.encoder.embeddings | |
| assert emb.position_embeddings is None and emb.token_type_embeddings is None and emb.embed_proj is None | |
| assert model.encoder.encoder.conv is None | |
| if native: | |
| self.native_encoder = model.encoder | |
| else: | |
| self.embedding_norm = emb.LayerNorm | |
| with torch.no_grad(): | |
| rel = model.encoder.encoder.get_rel_embedding().detach() | |
| self.layers = nn.ModuleList(ShapedDebertaLayer(x, rel, shape.seq) for x in model.encoder.encoder.layer) | |
| def forward(self, inputs_embeds, attention_mask, label_routing): | |
| if self.native: | |
| hidden = self.native_encoder(inputs_embeds=inputs_embeds, attention_mask=attention_mask).last_hidden_state | |
| else: | |
| hidden = self.embedding_norm(inputs_embeds) * attention_mask.unsqueeze(-1) | |
| mask = attention_mask[:, None, None, :] * attention_mask[:, None, :, None] | |
| for layer in self.layers: | |
| hidden = layer(hidden, mask) | |
| rows = rank4_matmul(label_routing.unsqueeze(1), hidden.unsqueeze(1)) | |
| return self.classifier(rows).reshape(1, 1, 1, self.shape.labels) | |