source/src/strata/modeling/compose/natural_compiler.py

"""Factorized natural predicate, operator, and argument-span compiler."""

from __future__ import annotations

from dataclasses import dataclass

import torch
from torch import nn


@dataclass(slots=True)
class NaturalCompilerOutput:
    predicate_logits: torch.Tensor
    operator_logits: torch.Tensor
    span_start_logits: torch.Tensor
    span_end_logits: torch.Tensor
    argument_head_logits: torch.Tensor
    role_presence_logits: torch.Tensor


class NaturalAtomicCompiler(nn.Module):
    """Compile atomic natural queries without an answer or graph-state bypass."""

    def __init__(
        self,
        pretrained_word_vectors: torch.Tensor,
        *,
        hidden_dim: int = 192,
        query_dim: int = 128,
        role_count: int = 5,
    ) -> None:
        super().__init__()
        if pretrained_word_vectors.ndim != 2:
            raise ValueError("pretrained_word_vectors must be [vocabulary, width]")
        self.role_count = int(role_count)
        self.word_vectors = nn.Embedding.from_pretrained(
            pretrained_word_vectors.to(torch.float32), freeze=True, padding_idx=0
        )
        self.word_projection = nn.Sequential(
            nn.LayerNorm(pretrained_word_vectors.shape[1]),
            nn.Linear(pretrained_word_vectors.shape[1], hidden_dim),
            nn.GELU(),
        )
        self.anchor_projection = nn.Sequential(
            nn.LayerNorm(pretrained_word_vectors.shape[1]),
            nn.Linear(pretrained_word_vectors.shape[1], hidden_dim),
            nn.GELU(),
        )
        self.anchor_scale = nn.Parameter(torch.tensor(10.0))
        self.context_encoder = nn.GRU(
            hidden_dim,
            hidden_dim,
            num_layers=2,
            batch_first=True,
            bidirectional=True,
            dropout=0.1,
        )
        context_dim = 2 * hidden_dim
        self.predicate_head = nn.Linear(context_dim, 1)

        self.query_projection = nn.Sequential(
            nn.LayerNorm(pretrained_word_vectors.shape[1]),
            nn.Linear(pretrained_word_vectors.shape[1], query_dim),
            nn.GELU(),
        )
        self.query_encoder = nn.GRU(
            query_dim,
            query_dim,
            batch_first=True,
            bidirectional=True,
        )
        self.operator_head = nn.Sequential(
            nn.LayerNorm(2 * query_dim),
            nn.Linear(2 * query_dim, context_dim),
            nn.GELU(),
            nn.Linear(context_dim, role_count),
        )

        self.role_embedding = nn.Embedding(role_count, hidden_dim)
        self.word_span_projection = nn.Linear(context_dim, hidden_dim, bias=False)
        self.predicate_span_projection = nn.Linear(context_dim, hidden_dim, bias=False)
        self.start_head = nn.Linear(hidden_dim, 1, bias=False)
        self.end_head = nn.Linear(hidden_dim, 1, bias=False)
        self.argument_head = nn.Linear(hidden_dim, 1, bias=False)
        self.role_presence = nn.Sequential(
            nn.Linear(context_dim + hidden_dim, hidden_dim),
            nn.GELU(),
            nn.Linear(hidden_dim, 1),
        )

    def forward(
        self,
        word_ids: torch.Tensor,
        word_mask: torch.Tensor,
        query_word_ids: torch.Tensor,
        query_mask: torch.Tensor,
        *,
        anchor_word_ids: torch.Tensor | None = None,
        gold_predicate_index: torch.Tensor | None = None,
    ) -> NaturalCompilerOutput:
        vectors = self.word_projection(self.word_vectors(word_ids))
        lengths = word_mask.sum(dim=-1).to(torch.long).cpu()
        packed = nn.utils.rnn.pack_padded_sequence(
            vectors, lengths, batch_first=True, enforce_sorted=False
        )
        encoded, _ = self.context_encoder(packed)
        words, _ = nn.utils.rnn.pad_packed_sequence(
            encoded, batch_first=True, total_length=word_ids.shape[1]
        )
        predicate_logits = self.predicate_head(words).squeeze(-1)
        if anchor_word_ids is not None:
            anchor = self.anchor_projection(self.word_vectors(anchor_word_ids))
            lexical = torch.nn.functional.normalize(vectors, dim=-1)
            anchor = torch.nn.functional.normalize(anchor, dim=-1)
            predicate_logits = predicate_logits + self.anchor_scale.clamp(0.0, 20.0) * torch.einsum(
                "bsd,bd->bs", lexical, anchor
            )
        predicate_logits = predicate_logits.masked_fill(~word_mask, -torch.inf)
        if gold_predicate_index is None:
            predicate_weights = torch.softmax(predicate_logits, dim=-1)
            predicate = torch.einsum("bs,bsd->bd", predicate_weights, words)
        else:
            predicate = words.gather(
                1,
                gold_predicate_index[:, None, None].expand(-1, 1, words.shape[-1]),
            ).squeeze(1)

        query_lengths = query_mask.sum(dim=-1).to(torch.long).cpu()
        query_packed = nn.utils.rnn.pack_padded_sequence(
            self.query_projection(self.word_vectors(query_word_ids)),
            query_lengths,
            batch_first=True,
            enforce_sorted=False,
        )
        _query_words, query_hidden = self.query_encoder(query_packed)
        query = torch.cat([query_hidden[-2], query_hidden[-1]], dim=-1)
        operator_logits = self.operator_head(query)

        role_ids = torch.arange(self.role_count, device=word_ids.device)
        role = self.role_embedding(role_ids)
        span_hidden = torch.tanh(
            self.word_span_projection(words)[:, None, :, :]
            + self.predicate_span_projection(predicate)[:, None, None, :]
            + role[None, :, None, :]
        )
        start = self.start_head(span_hidden).squeeze(-1).masked_fill(~word_mask[:, None, :], -torch.inf)
        end = self.end_head(span_hidden).squeeze(-1).masked_fill(~word_mask[:, None, :], -torch.inf)
        argument_head = self.argument_head(span_hidden).squeeze(-1).masked_fill(
            ~word_mask[:, None, :], -torch.inf
        )
        presence_input = torch.cat(
            [
                predicate[:, None, :].expand(-1, self.role_count, -1),
                role[None, :, :].expand(word_ids.shape[0], -1, -1),
            ],
            dim=-1,
        )
        presence = self.role_presence(presence_input).squeeze(-1)
        return NaturalCompilerOutput(
            predicate_logits, operator_logits, start, end, argument_head, presence
        )


def decode_valid_spans(
    start_logits: torch.Tensor,
    end_logits: torch.Tensor,
    word_mask: torch.Tensor,
    *,
    maximum_width: int = 24,
) -> tuple[torch.Tensor, torch.Tensor]:
    """Return the highest scoring valid start/end pair for every role."""

    scores = start_logits.unsqueeze(-1) + end_logits.unsqueeze(-2)
    length = start_logits.shape[-1]
    positions = torch.arange(length, device=scores.device)
    valid = positions[:, None] <= positions[None, :]
    valid &= positions[None, :] - positions[:, None] < maximum_width
    valid = valid[None, None, :, :] & (
        word_mask[:, None, :, None] & word_mask[:, None, None, :]
    )
    flat = scores.masked_fill(~valid, -torch.inf).flatten(-2)
    selected = flat.argmax(dim=-1)
    return selected // length, selected % length


def gated_compiler_reliability(
    confidence: torch.Tensor,
    *,
    threshold: float,
) -> torch.Tensor:
    """Pass qualified compiler confidence to BRR and zero every unsafe read."""

    if not 0.0 <= threshold <= 1.0:
        raise ValueError("threshold must be in [0, 1]")
    finite = torch.isfinite(confidence)
    qualified = finite & (confidence >= threshold)
    return torch.where(qualified, confidence.clamp(0.0, 1.0), torch.zeros_like(confidence))


__all__ = [
    "NaturalAtomicCompiler",
    "NaturalCompilerOutput",
    "decode_valid_spans",
    "gated_compiler_reliability",
]