export_model.py

"""Native GLiNER2.5-Decide multi-head classification path with explicit marker routing."""
import torch
from torch import nn
from transformers.models.deberta_v2 import modeling_deberta_v2


def coreml_safe_attention_forward(
    self, hidden_states, attention_mask, output_attentions=False,
    query_states=None, relative_pos=None, rel_embeddings=None,
):
    """Native DeBERTa attention with a finite mask sentinel for FP16 Core ML."""
    if query_states is None:
        query_states = hidden_states
    query = self.transpose_for_scores(self.query_proj(query_states), self.num_attention_heads)
    key = self.transpose_for_scores(self.key_proj(hidden_states), self.num_attention_heads)
    value = self.transpose_for_scores(self.value_proj(hidden_states), self.num_attention_heads)
    factor = 1 + int("c2p" in self.pos_att_type) + int("p2c" in self.pos_att_type)
    scale = modeling_deberta_v2.scaled_size_sqrt(query, factor)
    scores = torch.bmm(query, key.transpose(-1, -2) / scale.to(dtype=query.dtype))
    if self.relative_attention:
        relative = self.disentangled_attention_bias(
            query, key, relative_pos, self.pos_dropout(rel_embeddings), factor
        )
        scores = scores + relative
    scores = scores.view(-1, self.num_attention_heads, scores.size(-2), scores.size(-1))
    scores = scores.masked_fill(~attention_mask.bool(), -1e4)
    probabilities = self.dropout(torch.softmax(scores, dim=-1))
    context = torch.bmm(probabilities.view(-1, probabilities.size(-2), probabilities.size(-1)), value)
    context = context.view(-1, self.num_attention_heads, context.size(-2), context.size(-1))
    context = context.permute(0, 2, 1, 3).contiguous()
    context = context.view(context.size()[:-2] + (-1,))
    return (context, probabilities) if output_attentions else (context, None)


class GLiNER2DecideExport(nn.Module):
    """Encoder plus the shared label classifier, scored for every (head, label) marker at once.

    Logits are returned per head so the host can apply softmax (single-label) or sigmoid
    (multi-label) exactly as the native runtime does. Probabilities are the per-head softmax.
    """

    def __init__(self, native: nn.Module):
        super().__init__()
        self.encoder = native.encoder
        self.classifier = native.classifier

    def forward(self, input_ids, attention_mask, marker_indices, marker_mask):
        hidden = self.encoder(input_ids=input_ids.long(), attention_mask=attention_mask.long()).last_hidden_state
        heads, options = marker_indices.shape[1], marker_indices.shape[2]
        flat = marker_indices.long().reshape(1, heads * options, 1).expand(-1, -1, hidden.shape[-1])
        states = hidden.gather(1, flat)
        logits = self.classifier(states).reshape(1, heads, options)
        logits = torch.where(marker_mask > 0.5, logits, torch.full_like(logits, -1e4))
        return logits, torch.softmax(logits, dim=-1)