Blame for runtime/litert_boundary_runtime.py

"""gliner2.5-base-v1 on LiteRT with every trained head: gliner2's own schema handling and decoders over seven graphs.

Ported from FluidInference/gliner2-5-base-coreml@ca8ddf6c (Apache-2.0): extraction_runtime.py (the Core ML runtime),
preprocessing.py (prepare_extraction) and extraction_pool.py (select_candidates). Changes: LiteRT CompiledModel graphs in
place of Core ML packages; the features stage is two graphs (the encoder over the host-mapped float16 token table, then
the head); the bucket sizes come from the graphs.

    from litert_boundary_runtime import LiteRTBoundaryExtractor
    from gliner2 import Schema
    model = LiteRTBoundaryExtractor("path/to/bundle")
    model.extract("Alice founded Acme in Toronto.", Schema().entities(["person", "organization"]).relations(["founded"]))
    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. The graphs have fixed sizes (128 subword tokens, 64 text words,
8 queries and 8 classification choices per request); a larger request is refused with a ValueError, never truncated.
For long text use the *_long calls with a chunk_size that fits, e.g. extract_entities_long(text, labels, chunk_size=40,
chunk_overlap=8).

Needs gliner2==2.0.0 (schema, collator and decoders; no PyTorch weights are loaded), torch (the library's tensor
types), numpy and ai-edge-litert. The bundle folder holds the .tflite graphs, the float16 table, tokenizer/ and
source/config.json (the checkpoint's configuration).
"""
from __future__ import annotations

import json
import math
from contextvars import ContextVar
from pathlib import Path

import numpy as np
import torch
from gliner2.configuration import BoundaryHeadSettings
from gliner2.models.boundary.constants import MASK_LOGIT
from gliner2.models.boundary.engine import BoundaryExtractor
from gliner2.models.boundary.indexing import gather_rows
from gliner2.models.boundary.pool import PooledCandidates, _deduplicate_pool
from gliner2.models.boundary.proposal import select_top_boundaries
from gliner2.models.boundary.records import RecordGroupOutput
from gliner2.models.boundary.relations import RelationProposalSettings, RelationTypeSpec, TypedRelationPairGenerator
from gliner2.models.base import load_extractor_tokenizer
from gliner2.models.outputs import CandidateTensorBatch, ExtractorOutput
from gliner2.processor import SchemaTransformer
from gliner2.training.trainer import ExtractorCollator

FEATURE_NAMES = ("text_states", "query_states", "boundary_states", "boundary_mask", "start_logits", "end_logits", "inside_prefix",
                 "inside_prefix_mean", "pool_start_projection", "pool_end_projection", "null_logits", "count_log_rates",
                 "classification_logits")
HEAD_INPUTS = ("hidden_states", "text_indices", "text_mask", "query_indices", "query_mask", "cls_indices", "cls_mask")
SCORER_INPUTS = ("text_states", "text_mask", "query_states", "query_mask", "boundary_states", "start_logits", "end_logits",
                 "inside_prefix", "inside_prefix_mean", "candidate_indices", "candidate_mask", "candidate_compatibility")
TABLE = "gliner2_base_embeddings_f16.bin"
FILES = {
    "head": "gliner2_base_head_fp32.tflite", "scorer": "gliner2_base_scorer_fp32.tflite",
    "explicit": "gliner2_base_explicit_fp32.tflite", "relation": "gliner2_base_relation_fp32.tflite",
    "record_assignment": "gliner2_base_record_assignment_fp32.tflite", "record_anchorless": "gliner2_base_record_anchorless_fp32.tflite",
}


class Graph:
    """One signature of a LiteRT CompiledModel: run(feed) -> {output name: float32 array}."""

    def __init__(self, path: Path, signature: str | None = None, 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)
        if signature is None:
            signature = self.model.get_signature_list()[0] if hasattr(self.model, "get_signature_list") else None
        self.index = self.model.get_signature_index(signature) if signature else 0
        sig = self.model.get_signature_by_index(self.index)
        self.signature = signature or sig.get("key")
        self.inputs, self.outputs = sig["inputs"], sig["outputs"]
        self.ins, self.outs = self.model.create_input_buffers(self.index), self.model.create_output_buffers(self.index)
        details = self.model.get_input_tensor_details(self.signature)
        self.dtypes = {k: (np.int32 if "int" in str(v.get("dtype", "")).lower() else np.float32) for k, v in details.items()}
        self.input_shapes = {k: tuple(int(x) for x in v["shape"]) for k, v in details.items()}
        self.shapes = {k: tuple(int(x) for x in v["shape"]) for k, v in self.model.get_output_tensor_details(self.signature).items()}

    def __call__(self, feed: dict) -> dict:
        for buffer, name in zip(self.ins, self.inputs):
            buffer.write(np.ascontiguousarray(np.asarray(feed[name]), self.dtypes[name]))
        self.model.run_by_index(self.index, self.ins, self.outs)
        return {name: buffer.read(int(np.prod(self.shapes[name])), np.float32).reshape(self.shapes[name])
                for buffer, name in zip(self.outs, self.outputs)}


def signature_of(path: Path) -> str:
    from ai_edge_litert.compiled_model import CompiledModel
    model = CompiledModel.from_file(str(path))
    key = model.get_signature_by_index(0)["key"]
    model.close()
    return key


def to_tensor(value):
    return torch.from_numpy(np.asarray(value).copy())


def padded(value, size):
    if value.shape[0] > size:
        raise ValueError(f"Request exceeds the bucket capacity {size}")
    result = value.new_zeros((size, *value.shape[1:]))
    result[: value.shape[0]] = value
    return result


def prepare_extraction(processor, text: str, schema, length: int, max_words: int, max_queries: int, max_choices: int):
    """gliner2's collator on one request, padded to the bucket; refuses requests that do not fit (never truncates)."""
    built = schema.build() if hasattr(schema, "build") else schema
    batch = ExtractorCollator(processor, is_training=False, max_len=None, architecture="boundary")([(text, built)])
    if batch.input_ids.shape[1] > length:
        raise ValueError(f"Schema and text require {batch.input_ids.shape[1]} subwords; bucket holds {length}")
    if batch.text_word_indices.shape[1] > max_words:
        raise ValueError(f"Text requires {batch.text_word_indices.shape[1]} words; bucket holds {max_words}")
    if batch.query_marker_indices.shape[1] > max_queries:
        raise ValueError(f"Schema requires {batch.query_marker_indices.shape[1]} queries; bucket holds {max_queries}")
    if batch.cls_marker_indices.shape[1] > max_choices:
        raise ValueError(f"Schema requires {batch.cls_marker_indices.shape[1]} choices; bucket holds {max_choices}")
    if batch.query_marker_indices.shape[1] == 0 and batch.cls_marker_indices.shape[1] == 0:
        raise ValueError("Schema has no extraction or classification queries")

    def pad(values, width, fill=0):
        array = values.numpy()
        return np.pad(array, ((0, 0), (0, width - array.shape[1])), constant_values=fill)

    arrays = {
        "input_ids": pad(batch.input_ids, length, processor.tokenizer.pad_token_id).astype(np.int32),
        "attention_mask": pad(batch.attention_mask, length).astype(np.int32),
        "text_indices": pad(batch.text_word_indices, max_words).astype(np.int32),
        "text_mask": pad(batch.text_word_mask, max_words).astype(np.float32),
        "query_indices": pad(batch.query_marker_indices, max_queries).astype(np.int32),
        "query_mask": pad(batch.query_marker_mask, max_queries).astype(np.float32),
        "cls_indices": pad(batch.cls_marker_indices, max_choices).astype(np.int32),
        "cls_mask": pad(batch.cls_marker_mask, max_choices).astype(np.float32),
    }
    return arrays, batch


def select_candidates(start_projection, end_projection, boundary_mask, query_mask, start_logits, end_logits, *, boundary_top_k: int,
                      pool_size: int, min_pool_per_query: int) -> PooledCandidates:
    """gliner2 2.0.0's shared candidate pool (stable ranking and deduplication) from the head graph's projections."""
    batch, n_boundaries, dim = start_projection.shape
    n_queries = query_mask.shape[1]
    floor = torch.full_like(start_logits, MASK_LOGIT)
    valid_boundary = boundary_mask.unsqueeze(1) & query_mask.unsqueeze(-1)
    union_start = torch.where(valid_boundary, start_logits, floor).amax(1)
    union_end = torch.where(valid_boundary, end_logits, floor).amax(1)
    union_valid = boundary_mask & query_mask.any(-1, keepdim=True)
    _, starts, starts_valid = select_top_boundaries(union_start.unsqueeze(1), union_valid.unsqueeze(1), boundary_top_k)
    _, ends, ends_valid = select_top_boundaries(union_end.unsqueeze(1), union_valid.unsqueeze(1), boundary_top_k)
    starts, ends = starts[:, 0], ends[:, 0]
    starts_valid, ends_valid = starts_valid[:, 0], ends_valid[:, 0]
    n_starts, n_ends = starts.shape[1], ends.shape[1]
    pair_start = starts.unsqueeze(-1).expand(batch, n_starts, n_ends).reshape(batch, -1)
    pair_end = ends.unsqueeze(1).expand(batch, n_starts, n_ends).reshape(batch, -1)
    pair_valid = (starts_valid.unsqueeze(-1) & ends_valid.unsqueeze(1) & (ends.unsqueeze(1) > starts.unsqueeze(-1))).reshape(batch, -1)
    compatibility = (gather_rows(start_projection, pair_start) * gather_rows(end_projection, pair_end)).sum(-1) / math.sqrt(dim)
    union_pair_score = (compatibility + union_start.gather(1, pair_start.clamp(0, n_boundaries - 1))
                        + union_end.gather(1, pair_end.clamp(0, n_boundaries - 1)))
    quota = min(min_pool_per_query, pair_start.shape[-1])
    quota_keys = pair_start.new_zeros((batch, 0))
    quota_scores = union_pair_score.new_zeros((batch, 0))
    quota_valid = pair_valid.new_zeros((batch, 0))
    if quota:
        start_idx = pair_start.clamp(0, start_logits.shape[2] - 1).unsqueeze(1).expand(batch, n_queries, -1)
        end_idx = pair_end.clamp(0, end_logits.shape[2] - 1).unsqueeze(1).expand(batch, n_queries, -1)
        per_query = start_logits.gather(2, start_idx) + end_logits.gather(2, end_idx) + compatibility.unsqueeze(1)
        per_query_valid = pair_valid.unsqueeze(1) & query_mask.unsqueeze(-1)
        ranked = torch.argsort(per_query.masked_fill(~per_query_valid, MASK_LOGIT), dim=-1, descending=True, stable=True)[..., :quota]
        quota_start, quota_end = start_idx.gather(-1, ranked), end_idx.gather(-1, ranked)
        quota_valid = per_query_valid.gather(-1, ranked).reshape(batch, -1)
        quota_keys = (quota_start * n_boundaries + quota_end).reshape(batch, -1)
        rank_bonus = torch.arange(quota, 0, -1, device=start_projection.device, dtype=union_pair_score.dtype)
        quota_scores = (union_pair_score.new_full((batch, n_queries, quota), -MASK_LOGIT * 0.5) + rank_bonus.view(1, 1, quota)).reshape(batch, -1)
    all_keys = torch.cat((quota_keys, pair_start * n_boundaries + pair_end), -1)
    all_scores = torch.cat((quota_scores, union_pair_score.detach()), -1)
    all_valid = torch.cat((quota_valid, pair_valid), -1)
    selected_keys, selected_valid = _deduplicate_pool(all_keys, all_scores, all_valid, pool_size, n_boundaries)
    selected_keys = torch.where(selected_valid, selected_keys, torch.zeros_like(selected_keys))
    selected_s = torch.div(selected_keys, n_boundaries, rounding_mode="floor")
    selected_e = selected_keys - selected_s * n_boundaries
    indices = torch.where(selected_valid.unsqueeze(-1), torch.stack((selected_s, selected_e), -1), torch.zeros_like(torch.stack((selected_s, selected_e), -1)))
    selected_compat = (gather_rows(start_projection, selected_s) * gather_rows(end_projection, selected_e)).sum(-1) / math.sqrt(dim)
    selected_score = (selected_compat + union_start.gather(1, selected_s.clamp(0, n_boundaries - 1))
                      + union_end.gather(1, selected_e.clamp(0, n_boundaries - 1))).masked_fill(~selected_valid, MASK_LOGIT)
    selected_compat = torch.where(selected_valid, selected_compat, torch.zeros_like(selected_compat))
    return PooledCandidates(indices=indices, mask=selected_valid, proposal_logits=selected_score, gold_mask=None,
                            compat_logits=selected_compat, stats=None)


class LiteRTBoundaryHead(torch.nn.Module):
    """Gives gliner2's decoder the tensors the boundary head would have produced, from the LiteRT stages."""

    def __init__(self, context: ContextVar, explicit: Graph):
        super().__init__()
        self.context, self.explicit = context, explicit
        self.max_queries, self.max_spans = explicit.input_shapes["span_indices"][1], explicit.input_shapes["span_indices"][2]

    def forward(self, text, text_mask, query, query_mask, return_candidates=True):
        state = self.context.get()
        features, full, count = state["features"], state["candidates"], query.shape[1]
        selected = CandidateTensorBatch(indices=full.indices[:, :count], proposal_logits=full.proposal_logits[:, :count],
                                        pair_logits=full.pair_logits[:, :count], valid_mask=full.valid_mask[:, :count],
                                        query_mask=full.query_mask[:, :count], candidate_states=full.candidate_states[:, :count])
        return ExtractorOutput(candidates=selected if return_candidates else None, start_logits=features["start_logits"][:, :count],
                               end_logits=features["end_logits"][:, :count], null_logits=features["null_logits"][:, :count],
                               count_log_rates=features["count_log_rates"][:, :count], batch_size=1)

    def score_explicit_spans(self, text, text_mask, query, query_mask, indices, valid_mask=None):
        state = self.context.get()
        features = state["features"]
        full_queries = features["query_states"][0]
        active = int(state["arrays"]["query_mask"].sum())
        query_count, span_count = query.shape[1], indices.shape[2]
        if query_count > self.max_queries or span_count > self.max_spans:
            raise ValueError("Explicit-span request exceeds the bucket capacity")
        selected = []
        for row in query[0]:
            matches = torch.isclose(full_queries[:active], row, atol=1e-6, rtol=0).all(-1).nonzero(as_tuple=False).flatten()
            if matches.numel() != 1:
                raise ValueError("Explicit-span query cannot be mapped to the encoded schema")
            selected.append(int(matches[0]))
        selection = torch.tensor(selected, dtype=torch.long)

        def feature(name):
            return padded(features[name][0].index_select(0, selection), self.max_queries).unsqueeze(0)

        span_indices = torch.zeros(1, self.max_queries, self.max_spans, 2, dtype=torch.int32)
        span_mask = torch.zeros(1, self.max_queries, self.max_spans, dtype=torch.float32)
        span_indices[:, :query_count, :span_count] = indices.int()
        span_mask[:, :query_count, :span_count] = valid_mask.float() if valid_mask is not None else 1.0
        feed = {"text_states": text, "text_mask": text_mask.float(), "query_states": padded(query[0], self.max_queries).unsqueeze(0),
                "query_mask": padded(query_mask[0].float(), self.max_queries).unsqueeze(0), "boundary_states": features["boundary_states"],
                "start_logits": feature("start_logits"), "end_logits": feature("end_logits"), "inside_prefix": feature("inside_prefix"),
                "inside_prefix_mean": feature("inside_prefix_mean"), "span_indices": span_indices, "span_mask": span_mask}
        output = self.explicit({k: v.numpy() for k, v in feed.items()})["span_logits"]
        return to_tensor(output)[:, :query_count, :span_count]


class LiteRTRelationScorer(torch.nn.Module):
    """Runs the relation graph after gliner2's own typed pair selection."""

    def __init__(self, context: ContextVar, graph: Graph):
        super().__init__()
        self.context, self.graph = context, graph
        self.max_words = graph.input_shapes["text_states"][1]
        self.max_relations = graph.input_shapes["relation_states"][1]
        self.pair_cap = graph.input_shapes["pair_mask"][0]

    def forward(self, text, relation, candidates, pairs):
        count = len(pairs)
        if text.shape[1] != self.max_words or relation.shape[1] > self.max_relations or count > self.pair_cap:
            raise ValueError("Relation request exceeds the bucket capacity")
        relation_states = torch.zeros(1, self.max_relations, relation.shape[-1], dtype=relation.dtype)
        relation_states[:, : relation.shape[1]] = relation
        text_length = int(self.context.get()["arrays"]["text_mask"].sum())
        feed = {"text_states": text, "text_length": torch.tensor([text_length], dtype=torch.int32), "relation_states": relation_states,
                "batch_index": padded(pairs.batch_index.int(), self.pair_cap), "relation_index": padded(pairs.relation_index.int(), self.pair_cap),
                "head_start": padded(pairs.head_start.int(), self.pair_cap), "head_end": padded(pairs.head_end.int(), self.pair_cap),
                "tail_start": padded(pairs.tail_start.int(), self.pair_cap), "tail_end": padded(pairs.tail_end.int(), self.pair_cap),
                "pair_mask": padded(pairs.pair_mask.float(), self.pair_cap)}
        return to_tensor(self.graph({k: v.numpy() for k, v in feed.items()})["relation_logits"])[:count]


class LiteRTRecordHead(torch.nn.Module):
    """Keeps gliner2's instance ordering while the record decoder's learned layers run on LiteRT."""

    def __init__(self, assignment: Graph, anchorless: Graph):
        super().__init__()
        self.assignment, self.anchorless = assignment, anchorless
        self.max_instances = assignment.input_shapes["instance_states"][0]
        self.max_fields, self.max_candidates = assignment.input_shapes["field_candidate_states"][:2]

    def forward_group(self, spec, query_states, candidates, sample_index):
        field_specs = list(spec.fields)
        field_query_ids = [field.query_id for field in field_specs]
        if len(field_specs) > self.max_fields:
            raise ValueError("Record has more fields than the bucket holds")
        if any(q < 0 or q >= query_states.shape[0] for q in field_query_ids):
            raise ValueError("Record field query is outside the encoded schema")
        field_states, field_spans, field_logits, field_masks = [], [], [], []
        for query_id in field_query_ids:
            mask = candidates.valid_mask[sample_index, query_id]
            if int(mask.sum()) > self.max_candidates:
                raise ValueError("Record candidate count exceeds the bucket")
            field_states.append(candidates.candidate_states[sample_index, query_id][mask])
            field_spans.append(candidates.indices[sample_index, query_id][mask])
            field_logits.append(candidates.pair_logits[sample_index, query_id][mask])
            field_masks.append(torch.ones(int(mask.sum()), dtype=torch.bool))
        instance_seed, instance_spans = [], []
        if spec.mode == "natural":
            anchor = field_query_ids.index(spec.anchor_query_id)
            instances = field_states[anchor]
            for index, span in enumerate(field_spans[anchor]):
                instance_seed.append((anchor, index))
                instance_spans.append((int(span[0]), int(span[1])))
        elif spec.mode == "latent":
            instances = torch.cat(field_states, 0) if field_states else query_states.new_zeros((0, query_states.shape[-1]))
            for field_index, spans in enumerate(field_spans):
                for index, span in enumerate(spans):
                    instance_seed.append((field_index, index))
                    instance_spans.append((int(span[0]), int(span[1])))
        else:
            context = torch.cat(field_states, 0) if field_states else query_states.new_zeros((0, query_states.shape[-1]))
            capacity = self.anchorless.input_shapes["context_mask"][0]
            context_mask = torch.zeros(capacity, dtype=torch.float32)
            context_mask[: context.shape[0]] = 1.0
            out = self.anchorless({"context_states": padded(context, capacity).numpy(), "context_mask": context_mask.numpy()})
            instances = to_tensor(out["instance_states"])
            instance_seed = [None] * instances.shape[0]
            instance_spans = [None] * instances.shape[0]
        count = instances.shape[0]
        if count > self.max_instances:
            raise ValueError("Record instance count exceeds the bucket")
        candidate_states = torch.zeros(self.max_fields, self.max_candidates, query_states.shape[-1])
        for index, states in enumerate(field_states):
            candidate_states[index, : states.shape[0]] = states
        predicted = self.assignment({"instance_states": padded(instances, self.max_instances).numpy(),
                                     "field_queries": padded(query_states[field_query_ids], self.max_fields).numpy(),
                                     "field_candidate_states": candidate_states.numpy()})
        assignment = to_tensor(predicted["assignment_logits"])
        assignment_by_field = [assignment[:count, i, : 1 + states.shape[0]] for i, states in enumerate(field_states)]
        if spec.mode == "natural":
            object_logits = field_logits[field_query_ids.index(spec.anchor_query_id)]
        elif spec.mode == "latent":
            object_logits = to_tensor(predicted["latent_seed_logits"])[:count]
        else:
            object_logits = to_tensor(predicted["object_logits"])[:count]
        return RecordGroupOutput(spec=spec, object_logits=object_logits, assign_logits=assignment_by_field, field_query_ids=field_query_ids,
                                 field_specs=field_specs, field_spans=field_spans, field_cand_mask=field_masks, field_cand_logits=field_logits,
                                 instance_seed=instance_seed, instance_spans=instance_spans)


class LiteRTBoundaryExtractor(BoundaryExtractor):
    """gliner2's BoundaryExtractor (schemas, candidate selection, decoders) with every trained module on LiteRT."""

    def __init__(self, bundle: str, encoder: str = "gliner2_base_encoder_wfp16.tflite", accelerator: str = "cpu", threads: int = 4,
                 cache: str = ""):
        torch.nn.Module.__init__(self)
        folder = Path(bundle)
        config = json.loads((folder / "source" / "config.json").read_text())
        self.boundary_settings = BoundaryHeadSettings(**config["boundary_head"])
        self.processor = SchemaTransformer(tokenizer=load_extractor_tokenizer(str(folder / "tokenizer")), token_pooling="first")
        self.enable_records = self.boundary_settings.enable_records
        self.enable_relations = self.boundary_settings.enable_relations
        self.strict_extraction = True
        self._device_anchor = torch.nn.Parameter(torch.zeros(1), requires_grad=False)  # the runtime reads device and dtype from parameters()
        self._context = ContextVar("gliner2_litert_extraction_context")
        self.encoder_graph = Graph(folder / encoder, signature_of(folder / encoder), threads, accelerator, cache)
        self.head_graph = Graph(folder / FILES["head"], signature_of(folder / FILES["head"]), 1)
        self.scorer_graph = Graph(folder / FILES["scorer"], signature_of(folder / FILES["scorer"]), 1)
        graphs = {k: Graph(folder / FILES[k], signature_of(folder / FILES[k]), 1)
                  for k in ("explicit", "relation", "record_assignment", "record_anchorless")}
        self.length = self.encoder_graph.input_shapes["attention_mask"][1]
        self.max_words = self.head_graph.input_shapes["text_indices"][1]
        self.max_queries = self.head_graph.input_shapes["query_indices"][1]
        self.max_choices = self.head_graph.input_shapes["cls_indices"][1]
        hidden = self.encoder_graph.input_shapes["inputs_embeds"][2]
        rows = (folder / TABLE).stat().st_size // (2 * hidden)
        self.table = np.memmap(folder / TABLE, dtype="<f2", mode="r", shape=(rows, hidden))
        self.boundary_head = LiteRTBoundaryHead(self._context, graphs["explicit"])
        self.relation_scorer = LiteRTRelationScorer(self._context, graphs["relation"])
        self.record_decoder = LiteRTRecordHead(graphs["record_assignment"], graphs["record_anchorless"])
        self.relation_pair_generator = TypedRelationPairGenerator(RelationProposalSettings(
            heads_per_relation=self.boundary_settings.relation_heads_per_type, tails_per_relation=self.boundary_settings.relation_tails_per_type,
            pair_cap=self.boundary_settings.relation_pair_cap, argument_threshold=self.boundary_settings.relation_argument_proposal_threshold))

    def features(self, arrays: dict) -> dict:
        ids = arrays["input_ids"].reshape(-1).astype(np.int64)
        hidden = self.encoder_graph({"inputs_embeds": self.table[ids].astype(np.float32)[None],
                                     "attention_mask": arrays["attention_mask"].astype(np.float32).reshape(1, -1)})["hidden_states"]
        feed = {k: arrays[k] for k in HEAD_INPUTS if k != "hidden_states"}
        feed["hidden_states"] = hidden
        return {name: to_tensor(value) for name, value in self.head_graph(feed).items()}

    def _encode_core(self, batch):
        state = self._context.get()
        features = state["features"]
        query_mask = to_tensor(state["arrays"]["query_mask"]).bool()
        query_count = int(query_mask.sum())
        query_states = features["query_states"][:, :query_count]
        query_mask = query_mask[:, :query_count]
        text_mask = to_tensor(state["arrays"]["text_mask"]).bool()
        ext_specs, cls_specs, rel_specs, word_offsets = [], [], [], []
        for sample in range(len(batch)):
            specs = [{"group_index": q.task_index, "field_index": q.role_index, "task_type": q.task_type, "task_name": q.task_name,
                      "field_name": q.role_name} for q in batch.query_layouts[sample].queries]
            ext_specs.append(specs)
            classifications, choice_offset = [], 0
            for group in range(batch.schema_counts[sample]):
                if batch.task_types[sample][group] != "classifications":
                    continue
                count = max(len(batch.schema_special_indices[sample][group]) - 1, 0)
                tokens = batch.schema_tokens_list[sample][group]
                if count:
                    classifications.append({"group_index": group, "task_name": tokens[2], "schema_tokens": tokens,
                                            "group_embs": features["classification_logits"][sample, choice_offset: choice_offset + count]})
                choice_offset += count
            cls_specs.append(classifications)
            word_offsets.append(max(int(batch.text_word_counts[sample]) - len(batch.start_mappings[sample]), 0))
            groups = {}
            for query_id, spec in enumerate(specs):
                groups.setdefault(spec["group_index"], []).append(query_id)
            relations = []
            for group, role_ids in groups.items():
                if batch.task_types[sample][group] != "relations" or len(role_ids) < 2:
                    continue
                head_id, tail_id = role_ids[:2]
                roles = query_states[sample, [head_id, tail_id]]
                state_vec = torch.cat((roles[0], roles[1]), -1) if self.boundary_settings.directional_relation_states else roles.mean(0)
                relations.append({"group_index": group, "relation_type": specs[head_id]["task_name"],
                                  "spec": RelationTypeSpec(specs[head_id]["task_name"], head_query_ids=(head_id,), tail_query_ids=(tail_id,)),
                                  "query_state": state_vec})
            rel_specs.append(relations)
        return {"text_states": features["text_states"], "text_mask": text_mask, "text_lengths": text_mask.sum(-1).long(),
                "query_states": query_states, "query_mask": query_mask, "ext_specs": ext_specs, "cls_specs": cls_specs,
                "rel_specs": rel_specs, "word_offsets": word_offsets}

    def _extract_classification_result(self, results, schema_name, schema, embs, schema_tokens, temperature=1.0):
        cls_config = self._resolve_classification_config(schema_tokens[2], schema.get("classifications", []))
        if cls_config is None:
            return
        if temperature <= 0:
            raise ValueError("Classification temperature must be positive")
        logits = embs / temperature
        activation, multi_label = cls_config.get("class_act", "auto"), cls_config.get("multi_label", False)
        probabilities = torch.sigmoid(logits) if activation == "sigmoid" or (activation != "softmax" and multi_label) else torch.softmax(logits, dim=-1)
        labels = cls_config["labels"]
        if multi_label:
            threshold = cls_config.get("cls_threshold", 0.5)
            chosen = [(labels[i], float(probabilities[i])) for i in range(len(labels)) if float(probabilities[i]) >= threshold]
            if not chosen:
                best = int(probabilities.argmax())
                chosen = [(labels[best], float(probabilities[best]))]
            results[cls_config["task"]] = chosen
            return
        best = int(probabilities.argmax())
        results[cls_config["task"]] = (labels[best], float(probabilities[best]))

    def batch_extract(self, texts, schemas, batch_size: int = 8, threshold: float = 0.5, num_workers: int = 0, format_results: bool = True,
                      include_confidence: bool = False, include_spans: bool = False, max_len=None, overlap_policy=None):
        """One request at a time through the fixed-shape graphs (gliner2's batch API, so the *_long calls work too)."""
        if isinstance(schemas, list):
            if len(schemas) != len(texts):
                raise ValueError(f"Schema count ({len(schemas)}) != text count ({len(texts)})")
            schema_list = schemas
        else:
            schema_list = [schemas] * len(texts)
        return [self.extract(text, schema, threshold=threshold, format_results=format_results, include_confidence=include_confidence,
                             include_spans=include_spans, overlap_policy=overlap_policy) for text, schema in zip(texts, schema_list)]

    def extract(self, text: str, schema, threshold: float = 0.5, format_results: bool = True, include_confidence: bool = False,
                include_spans: bool = False, max_len=None, overlap_policy=None):
        schema_dicts, metadata_list = self._build_schema_dicts_and_metadata([schema])
        if overlap_policy is not None:
            metadata_list[0]["_overlap_policy"] = self._resolved_overlap_policy(overlap_policy)
        arrays, batch = prepare_extraction(self.processor, text, schema_dicts[0], self.length, self.max_words, self.max_queries, self.max_choices)
        features = self.features(arrays)
        query_mask = to_tensor(arrays["query_mask"]).bool()
        candidates = None
        if bool(query_mask.any()):
            head = self.boundary_settings
            pool = select_candidates(features["pool_start_projection"], features["pool_end_projection"], features["boundary_mask"].bool(),
                                     query_mask, features["start_logits"], features["end_logits"], boundary_top_k=head.pool_boundary_top_k,
                                     pool_size=head.pool_size, min_pool_per_query=head.min_pool_per_query)
            feed = {"text_states": features["text_states"], "text_mask": to_tensor(arrays["text_mask"]), "query_states": features["query_states"],
                    "query_mask": to_tensor(arrays["query_mask"]), "boundary_states": features["boundary_states"],
                    "start_logits": features["start_logits"], "end_logits": features["end_logits"], "inside_prefix": features["inside_prefix"],
                    "inside_prefix_mean": features["inside_prefix_mean"], "candidate_indices": pool.indices.int(),
                    "candidate_mask": pool.mask.float(), "candidate_compatibility": pool.compat_logits}
            scored = self.scorer_graph({k: v.numpy() for k, v in feed.items()})
            q = self.max_queries
            candidates = CandidateTensorBatch(indices=pool.indices.unsqueeze(1).expand(1, q, -1, 2),
                                              proposal_logits=pool.proposal_logits.unsqueeze(1).expand(1, q, -1),
                                              pair_logits=to_tensor(scored["pair_logits"]), valid_mask=pool.mask.unsqueeze(1).expand(1, q, -1),
                                              query_mask=query_mask,
                                              candidate_states=to_tensor(scored["candidate_states"]).unsqueeze(1).expand(1, q, -1, -1))
        token = self._context.set({"arrays": arrays, "features": features, "candidates": candidates})
        try:
            raw = self._extract_from_batch(batch, threshold, metadata_list, include_confidence, include_spans)[0]
            if format_results:
                return self.format_results(raw, include_confidence, metadata_list[0].get("relation_order", []),
                                           metadata_list[0].get("classification_tasks", []))
            return raw
        finally:
            self._context.reset(token)