Every trained head: explicit spans, relations and records; reference runtime; phone results
ff95cf8 verified | """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) | |