"""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=" 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)