extraction_export.py

"""Core ML graph wrappers for GLiNER2.5's trained boundary extraction path."""

import math
from contextlib import contextmanager

import torch
from coremltools.converters.mil import Builder as mb
from coremltools.converters.mil.frontend.torch.ops import _get_inputs
from coremltools.converters.mil.frontend.torch.torch_op_registry import register_torch_op
from coremltools.converters.mil.mil import types
from gliner2.models.boundary import encoding, heads
from gliner2.models.boundary.pool import PooledCandidates
from gliner2.models.boundary.proposal import BoundaryProposals
from transformers.models.deberta_v2 import modeling_deberta_v2

from export_model import coreml_safe_attention_forward


@register_torch_op(override=True)
def clamp_min(context, node):
    """Preserve the tensor dtype when TorchScript supplied a Python scalar."""
    x, y = _get_inputs(context, node, expected=2)
    if x.dtype != y.dtype:
        y = mb.cast(x=y, dtype=types.builtin_to_string(x.dtype))
    context.add(mb.maximum(x=x, y=y, name=node.name))


@register_torch_op(torch_alias=["clip"], override=True)
def clamp(context, node):
    """Avoid promoting integer span indices to float for an absent bound."""
    inputs = _get_inputs(context, node, expected=[1, 2, 3])
    x = inputs[0]
    lower = inputs[1] if len(inputs) > 1 and inputs[1] is not None else None
    upper = inputs[2] if len(inputs) > 2 and inputs[2] is not None else None
    result = x
    for bound, op in ((upper, mb.minimum), (lower, mb.maximum)):
        if bound is None:
            continue
        if bound.dtype != x.dtype:
            bound = mb.cast(x=bound, dtype=types.builtin_to_string(x.dtype))
        result = op(x=result, y=bound)
    context.add(mb.identity(x=result, name=node.name))


def shift_left(text_states, bos_state):
    """Functional equivalent of the upstream in-place BOS placement."""
    bos = bos_state.to(text_states.dtype).view(1, 1, -1)
    return torch.cat((bos.expand(text_states.shape[0], 1, -1), text_states), 1)


def shift_right(text_states, text_lengths, eos_state):
    """Functional equivalent of the upstream in-place EOS placement."""
    batch, length, hidden = text_states.shape
    eos = eos_state.to(text_states.dtype).view(1, 1, hidden)
    right = torch.cat((text_states, eos.expand(batch, 1, hidden)), 1)
    positions = torch.arange(length + 1, device=text_states.device).view(1, length + 1, 1)
    return torch.where(
        positions == text_lengths.view(batch, 1, 1),
        eos.expand(batch, length + 1, hidden),
        right,
    )


def safe_boundary_attention(self, states, mask):
    """Explicit scaled attention with the same finite masked result as upstream."""
    batch, length, dim = states.shape
    qkv = self.qkv_projection(self.norm(states)).view(batch, length, 3, self.num_heads, self.head_dim)
    query, key, value = qkv.permute(2, 0, 3, 1, 4)
    allowed = mask.view(batch, 1, 1, length).expand(batch, 1, length, length)
    if self.window > 0:
        positions = torch.arange(length, device=states.device)
        local = (positions.view(length, 1) - positions.view(1, length)).abs() <= self.window
        allowed = allowed & local.view(1, 1, length, length)
    diagonal = torch.eye(length, dtype=torch.bool, device=states.device).view(1, 1, length, length)
    allowed = (allowed.float() + diagonal.float()) > 0.5
    scores = torch.matmul(query, key.transpose(-1, -2)) / math.sqrt(self.head_dim)
    scores = scores.masked_fill(~allowed, -1e4)
    attended = torch.matmul(torch.softmax(scores, dim=-1), value)
    attended = attended.transpose(1, 2).reshape(batch, length, dim)
    return (states + self.dropout(self.output_projection(attended))) * mask.unsqueeze(-1).to(states.dtype)


def safe_query_head(self, boundary, boundary_mask, text, text_mask, query, query_mask):
    """Upstream marginals with a dtype-safe count clamp for Core ML."""
    scale = 1.0 / math.sqrt(self.boundary_dim)
    start = (
        torch.einsum(
            "bld,bqd->bql",
            self.dropout(self.start_boundary_projection(boundary)),
            self.start_query_projection(query),
        )
        * scale
    )
    end = (
        torch.einsum(
            "bld,bqd->bql",
            self.dropout(self.end_boundary_projection(boundary)),
            self.end_query_projection(query),
        )
        * scale
    )
    inside = (
        torch.einsum(
            "bld,bqd->bql",
            self.dropout(self.inside_text_projection(text)),
            self.inside_query_projection(query),
        )
        * scale
    )
    boundary_keep = boundary_mask.unsqueeze(1) & query_mask.unsqueeze(-1)
    text_keep = text_mask.unsqueeze(1) & query_mask.unsqueeze(-1)
    start = heads._masked_fill_min(start, boundary_keep)
    end = heads._masked_fill_min(end, boundary_keep)
    inside = heads._masked_fill_min(inside, text_keep)
    inside_for_prefix = inside.masked_fill(~text_keep, 0.0).float()
    count = torch.clamp(text_keep.sum(-1, keepdim=True).float(), min=1.0)
    mean = (inside_for_prefix.sum(-1, keepdim=True) / count).detach()
    centered = (inside_for_prefix - mean) * text_keep.to(inside_for_prefix.dtype)
    zeros = torch.zeros(centered.shape[0], centered.shape[1], 1, dtype=torch.float32, device=text.device)
    prefix = torch.cat((zeros, centered.cumsum(dim=-1)), dim=-1)
    return heads.BoundaryMarginals(start, end, inside, prefix, mean)


@contextmanager
def coreml_trace_patches():
    """Apply and restore mathematically equivalent trace-safe operations."""
    saved = (
        modeling_deberta_v2.scaled_size_sqrt,
        modeling_deberta_v2.build_rpos,
        modeling_deberta_v2.DisentangledSelfAttention.forward,
        encoding.shift_left_with_bos,
        encoding.shift_right_with_eos,
        encoding.BoundaryAttentionBlock.forward,
        heads.BoundaryQueryHead.forward,
    )

    def static_scale(query_layer, scale_factor):
        value = math.sqrt(float(query_layer.shape[-1] * scale_factor))
        return torch.tensor(value, dtype=torch.float32, device=query_layer.device)

    modeling_deberta_v2.scaled_size_sqrt = static_scale
    modeling_deberta_v2.build_rpos = lambda query, key, relative_pos, buckets, max_pos: relative_pos
    modeling_deberta_v2.DisentangledSelfAttention.forward = coreml_safe_attention_forward
    encoding.shift_left_with_bos = shift_left
    encoding.shift_right_with_eos = shift_right
    encoding.BoundaryAttentionBlock.forward = safe_boundary_attention
    heads.BoundaryQueryHead.forward = safe_query_head
    try:
        yield
    finally:
        (
            modeling_deberta_v2.scaled_size_sqrt,
            modeling_deberta_v2.build_rpos,
            modeling_deberta_v2.DisentangledSelfAttention.forward,
            encoding.shift_left_with_bos,
            encoding.shift_right_with_eos,
            encoding.BoundaryAttentionBlock.forward,
            heads.BoundaryQueryHead.forward,
        ) = saved


class ExtractionFeaturesExport(torch.nn.Module):
    """Trained encoder, boundary marginals, pool projections and null/count heads."""

    def __init__(self, native):
        super().__init__()
        self.encoder = native.encoder
        self.head = native.boundary_head
        self.classifier = native.classifier

    def forward(
        self,
        input_ids,
        attention_mask,
        text_indices,
        text_mask,
        query_indices,
        query_mask,
        cls_indices,
        cls_mask,
    ):
        hidden = self.encoder(input_ids=input_ids.long(), attention_mask=attention_mask.long()).last_hidden_state
        text_idx = text_indices.long().unsqueeze(-1).expand(-1, -1, hidden.shape[-1])
        query_idx = query_indices.long().unsqueeze(-1).expand(-1, -1, hidden.shape[-1])
        text = hidden.gather(1, text_idx) * text_mask.unsqueeze(-1)
        query = hidden.gather(1, query_idx) * query_mask.unsqueeze(-1)
        cls_idx = cls_indices.long().unsqueeze(-1).expand(-1, -1, hidden.shape[-1])
        classification_states = hidden.gather(1, cls_idx)
        classification_logits = self.classifier(classification_states).squeeze(-1)
        classification_logits = torch.where(
            cls_mask > 0.5, classification_logits, torch.full_like(classification_logits, -1e4)
        )
        tm, qm = text_mask > 0.5, query_mask > 0.5
        encoded = self.head.boundary_encoder(text, tm)
        marginal = self.head.boundary_query_head(encoded.states, encoded.mask, text, tm, query, qm)
        return (
            text,
            query,
            encoded.states,
            encoded.mask.float(),
            marginal.start_logits,
            marginal.end_logits,
            marginal.inside_prefix,
            marginal.inside_prefix_mean,
            self.head.shared_pool_builder.start_projection(encoded.states),
            self.head.shared_pool_builder.end_projection(encoded.states),
            self.head.null_projection(query).squeeze(-1),
            self.head.count_head(query).squeeze(-1),
            classification_logits,
        )


class ExtractionScoreExport(torch.nn.Module):
    """Trained shared-pool reranker and record candidate state projection."""

    def __init__(self, native):
        super().__init__()
        self.scorer = native.boundary_head.shared_pool_scorer
        self.candidate_encoder = native.boundary_head.candidate_encoder

    def forward(
        self,
        text,
        text_mask,
        query,
        query_mask,
        boundary,
        starts,
        ends,
        inside,
        inside_mean,
        indices,
        pool_mask,
        compatibility,
    ):
        tm, qm = text_mask > 0.5, query_mask > 0.5
        pooled = PooledCandidates(indices.long(), pool_mask > 0.5, None, None, compatibility)
        score, _ = self.scorer(
            boundary,
            query,
            qm,
            pooled,
            starts,
            ends,
            inside,
            tm.sum(-1).long(),
            text,
            tm,
            inside_prefix_mean=inside_mean,
        )
        index = indices.long()
        start_states = boundary.gather(1, index[..., 0].unsqueeze(-1).expand(-1, -1, boundary.shape[-1]))
        end_states = boundary.gather(1, index[..., 1].unsqueeze(-1).expand(-1, -1, boundary.shape[-1]))
        candidate_states = self.candidate_encoder(torch.cat((start_states, end_states), -1))
        candidate_states = candidate_states * pool_mask.unsqueeze(-1)
        return score.transpose(1, 2), candidate_states


class ExtractionRelationExport(torch.nn.Module):
    """The trained sparse relation scorer with tensor-only pair routing."""

    def __init__(self, native):
        super().__init__()
        self.scorer = native.relation_scorer

    def forward(
        self,
        text,
        text_length,
        relation,
        batch_index,
        relation_index,
        head_start,
        head_end,
        tail_start,
        tail_end,
        pair_mask,
    ):
        scorer = self.scorer
        length = text.shape[1]
        batch_valid = (batch_index >= 0) & (batch_index < text.shape[0])
        relation_valid = (relation_index >= 0) & (relation_index < relation.shape[1])
        valid = batch_valid & relation_valid & (pair_mask > 0.5)
        batch = batch_index.long().clamp(0, text.shape[0] - 1)
        rel_index = relation_index.long().clamp(0, relation.shape[1] - 1)

        def gather(position):
            return text[batch, position.long().clamp(0, length - 1)]

        h_start = gather(head_start)
        h_end = gather(head_end - 1)
        t_start = gather(tail_start)
        t_end = gather(tail_end - 1)
        rel = relation[batch, rel_index]
        delta = (tail_start - head_start).to(text.dtype)
        order = torch.sign(delta).unsqueeze(-1)
        distance = (delta.abs() / text_length.float().clamp_min(1.0)).unsqueeze(-1)
        features = torch.cat((h_start, h_end, t_start, t_end, rel, order, distance), -1)
        score = scorer.mlp(features).squeeze(-1)
        if scorer.use_biaffine_content:
            prefix = torch.cat(
                (text.new_zeros(text.shape[0], 1, scorer.hidden_size), text.float().cumsum(1).to(text.dtype)),
                dim=1,
            )

            def pool(start, end):
                total = prefix[batch, end.long().clamp(0, length)] - prefix[batch, start.long().clamp(0, length)]
                width = (end - start).clamp_min(1).unsqueeze(-1).to(total.dtype)
                return total / width

            head_content = scorer.head_content_projection(pool(head_start, head_end))
            tail_content = scorer.tail_content_projection(pool(tail_start, tail_end))
            gate = torch.sigmoid(scorer.relation_content_gate(rel))
            biaffine = (head_content * gate * tail_content).sum(-1) / (scorer.hidden_size**0.5)
            linear = scorer.content_linear(torch.cat((head_content, tail_content, rel), -1)).squeeze(-1)
            score = score + biaffine + linear
        return score.masked_fill(~valid, 0.0)


class ExtractionRecordAssignmentExport(torch.nn.Module):
    """All trained natural/latent/anchorless object and field assignment layers."""

    def __init__(self, native):
        super().__init__()
        self.head = native.record_decoder

    def forward(self, instances, field_queries, field_candidates):
        head = self.head
        instance_projection = head.inst_proj(instances)
        field_projection = head.field_proj(field_queries)
        query = instance_projection.unsqueeze(1) + field_projection.unsqueeze(0)
        null_scores = torch.einsum("ifd,d->if", query, head.null_embed)
        candidate_scores = torch.einsum("ifd,fcd->ifc", query, head.cand_proj(field_candidates))
        assignment = torch.cat((null_scores.unsqueeze(-1), candidate_scores), -1)
        object_scores = head.object_head(instances).squeeze(-1)
        latent_scores = head.latent_seed_head(instances).squeeze(-1)
        return assignment, object_scores, latent_scores


class ExtractionRecordAnchorlessExport(torch.nn.Module):
    """Trained learned-instance and contextual attention path for records."""

    def __init__(self, native):
        super().__init__()
        self.head = native.record_decoder

    def forward(self, context_states, context_mask):
        head = self.head
        instances = head.instance_embed
        query = head.q_proj(instances)
        key = head.k_proj(context_states)
        value = head.v_proj(context_states)
        attention = torch.matmul(query, key.transpose(-1, -2)) / math.sqrt(head.record_dim)
        attention = attention.masked_fill(context_mask.unsqueeze(0) < 0.5, -1e4)
        pooled = torch.matmul(torch.softmax(attention, -1), value)
        return instances + pooled * (context_mask.sum() > 0).to(pooled.dtype)


class ExtractionExplicitSpanExport(torch.nn.Module):
    """Trained proposal prior and reranker for forced attribute/enum spans."""

    def __init__(self, native):
        super().__init__()
        self.proposer = native.boundary_head.boundary_proposer
        self.scorer = native.boundary_head.pair_scorer

    def forward(
        self,
        text,
        text_mask,
        query,
        query_mask,
        boundary,
        starts,
        ends,
        inside,
        inside_mean,
        indices,
        valid_mask,
    ):
        tm, qm = text_mask > 0.5, query_mask > 0.5
        idx = indices.long()
        legal = (
            (idx[..., 0] >= 0)
            & (idx[..., 1] > idx[..., 0])
            & (idx[..., 1] <= tm.sum(-1).view(-1, 1, 1))
            & qm.unsqueeze(-1)
            & (valid_mask > 0.5)
        )
        compatibility = self.proposer.score_explicit_pairs(boundary, query, idx, legal)
        proposals = BoundaryProposals(
            indices=idx,
            logits=None,
            valid_mask=legal,
            compat_logits=compatibility,
        )
        return self.scorer(
            boundary,
            query,
            proposals,
            starts,
            ends,
            inside,
            tm.sum(-1).long(),
            text,
            tm,
            inside_prefix_mean=inside_mean,
        )