"""Native GLiNER2 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 GLiNER2ClassificationExport(nn.Module):
    def __init__(self, native: nn.Module):
        super().__init__()
        self.encoder = native.encoder
        self.classifier = native.classifier
        self.temperature = float(native.boundary_settings.classification_temperature)

    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
        indices = marker_indices.long().unsqueeze(-1).expand(-1, -1, hidden.shape[-1])
        states = hidden.gather(1, indices)
        logits = self.classifier(states).squeeze(-1) / self.temperature
        logits = torch.where(marker_mask > 0.5, logits, torch.full_like(logits, -1e4))
        return logits, torch.softmax(logits, dim=-1)