app/fp/parser3_model.py

"""Lossless character inputs and complete-span heads for the v3 controls.

The encoder still sees one position per Unicode scalar, so UTF-8 fallback does
not consume a different positional budget or split an output character. Existing
vocabulary IDs and embeddings are preserved. Unknown scalars have a reversible
integer ID and four ordered UTF-8 byte slots, never the shared UNK ID alone.
"""

import numpy as np
import torch
from torch import nn

from fp import charclass
from fp.parser2 import Extractor, KINDS, LABELS, predict_spans, tag_paths, legality_matrix

TYPES = ["O", *KINDS]


def utf8_slots(codepoints):
    """UTF-8 bytes + 1, padded with zero to four slots (tensorized, ordered)."""
    cp = codepoints.clamp_min(0)
    two, three, four = cp >= 0x80, cp >= 0x800, cp >= 0x10000
    first = torch.where(four, 0xF0 | (cp >> 18), torch.where(
        three, 0xE0 | (cp >> 12), torch.where(two, 0xC0 | (cp >> 6), cp)))
    second = torch.where(four, 0x80 | ((cp >> 12) & 63), torch.where(
        three, 0x80 | ((cp >> 6) & 63), 0x80 | (cp & 63)))
    third = torch.where(four, 0x80 | ((cp >> 6) & 63), 0x80 | (cp & 63))
    fourth = 0x80 | (cp & 63)
    return torch.stack((first + 1, torch.where(two, second + 1, 0),
                        torch.where(three, third + 1, 0),
                        torch.where(four, fourth + 1, 0)), dim=-1)


class Parser(Extractor):
    def __init__(self, representation="unicode", head="bio", *, encoder=None, tokenizer=None,
                 char_classes=False):
        if representation not in ("legacy", "unicode", "subword") or head not in ("bio", "span"):
            raise ValueError("unknown representation or head")
        if encoder is None:
            super().__init__("pretrained", len(LABELS))
        else:
            # Tiny injected encoders make alignment/export tests independent of downloads.
            nn.Module.__init__(self)
            self.encoder, self.tokenizer, self.mode = encoder, tokenizer, "pretrained"
            self.tags = nn.Linear(encoder.config.hidden_size, len(LABELS))
        self.representation, self.head = representation, head
        self.base_vocab_size = self.encoder.config.vocab_size
        hidden = self.encoder.config.hidden_size
        if representation == "unicode":
            if hidden % 4:
                raise ValueError("hidden size must be divisible by four")
            self.byte_slots = nn.ModuleList([nn.Embedding(257, hidden // 4, padding_idx=0)
                                             for _ in range(4)])
            for embedding in self.byte_slots:
                nn.init.normal_(embedding.weight, std=0.02)
                with torch.no_grad():
                    embedding.weight[0].zero_()
            self.byte_projection = nn.Linear(hidden, hidden, bias=False)
            nn.init.eye_(self.byte_projection.weight)
        if head == "span":
            self.tags = nn.Linear(hidden, len(TYPES))
            self.starts = nn.Linear(hidden, len(TYPES))
            self.ends = nn.Linear(hidden, len(TYPES))
        self.char_classes = char_classes
        # (structural, other, fallback) identity-dropout rates; set by the trainer.
        self.identity_dropout = (0.0, 0.0, 0.0)
        if char_classes:
            if representation != "unicode":
                raise ValueError("character classes need the unicode representation")
            # Zero-initialised: a model built on an existing encoder starts out
            # computing exactly what it computed without class features.
            self.category_embedding = nn.Embedding(charclass.N_CATEGORIES, hidden)
            self.script_embedding = nn.Embedding(charclass.N_SCRIPTS, hidden)
            nn.init.zeros_(self.category_embedding.weight)
            nn.init.zeros_(self.script_embedding.weight)
            # Stands in for a dropped character identity during training only.
            self.anonymous = nn.Parameter(torch.randn(hidden) * 0.02)
            category, script = charclass.tables()
            id_cp = torch.full((self.base_vocab_size,), -1, dtype=torch.long)
            for token, index in self.tokenizer.get_vocab().items():
                if len(token) == 1 and index < self.base_vocab_size:
                    id_cp[index] = ord(token)
            # Lookup tables are rebuilt from code, never saved: the exported Go
            # artifact carries the same assignment as ranges in its JSON sidecar.
            self.register_buffer("id_codepoint", id_cp, persistent=False)
            self.register_buffer("category_table", torch.from_numpy(category.astype("int64")),
                                 persistent=False)
            self.register_buffer("script_table", torch.from_numpy(script.astype("int64")),
                                 persistent=False)
            structural = torch.zeros(charclass.N_CATEGORIES, dtype=torch.bool)
            structural[list(charclass.STRUCTURAL)] = True
            self.register_buffer("structural_category", structural, persistent=False)

    def fits(self, row):
        """Encoder-window check: characters for character encoders, tokens for subword.

        The frozen character encoder sees one position per code point plus the
        two specials; a subword encoder sees the tokenizer's pieces. No row is
        ever truncated, so the caller filters with this before encoding.
        """
        if self.representation == "subword":
            tokens = self.tokenizer(row["query"], add_special_tokens=True)["input_ids"]
            return len(tokens) <= self.encoder.config.max_position_embeddings
        return len(row["query"]) + 2 <= self.encoder.config.max_position_embeddings

    def encode(self, row):
        if self.representation == "subword":
            return self._encode_subword(row)
        query = row["query"]
        # encode() rejects surrogate code points as well as invalid UTF-8 text.
        query.encode("utf-8", errors="strict")
        tokens = self.tokenizer.tokenize(query)
        if tokens != list(query):
            raise ValueError("v3 requires a character tokenizer with exact original offsets")
        ids = self.tokenizer.convert_tokens_to_ids(tokens)
        if self.representation == "unicode":
            ids = [self.base_vocab_size + ord(char) if token == self.tokenizer.unk_token_id
                   else token for char, token in zip(query, ids)]
        ids = [self.tokenizer.cls_token_id, *ids, self.tokenizer.sep_token_id]
        if len(ids) > self.encoder.config.max_position_embeddings:
            raise ValueError(f"{len(ids)} positions exceed encoder limit; no truncation")
        labels = self._labels(query, row["spans"], row.get("partial", False))
        return ids, list(range(1, len(query) + 1)), list(range(2, len(query) + 2)), labels, row

    def _labels(self, query, spans, partial=False):
        """BIO labels for one row, refusing invalid or misquoted spans.

        Partial rows mark uncovered characters -100 so the loss ignores them:
        the row asserts only what its spans cover, not that nothing else exists.
        """
        labels = [-100 if partial else 0] * len(query)
        previous = 0
        for span in spans:
            a, b, kind = span["start"], span["end"], span["kind"]
            if not 0 <= previous <= a < b <= len(query):
                raise ValueError("invalid or overlapping span offsets")
            previous = b
            if span.get("text", query[a:b]) != query[a:b]:
                raise ValueError("span text does not match original offsets")
            labels[a] = LABELS.index("B-" + kind)
            labels[a + 1:b] = [LABELS.index("I-" + kind)] * (b - a - 1)
        return labels

    def _encode_subword(self, row):
        """Subword tokens with one character window per code point.

        Each character shares the window of the token(s) covering it, so the
        cumulative-mean pooling in forward() still yields exactly one feature per
        character and the span head keeps character resolution. A character the
        tokenizer never covers is an error, never a silent guess.
        """
        query = row["query"]
        query.encode("utf-8", errors="strict")
        encoding = self.tokenizer(query, add_special_tokens=True,
                                  return_offsets_mapping=True)
        ids, offsets = encoding["input_ids"], encoding["offset_mapping"]
        if len(ids) > self.encoder.config.max_position_embeddings:
            raise ValueError(f"{len(ids)} tokens exceed encoder limit; no truncation")
        starts, ends = [], []
        for index in range(len(query)):
            first = last = None
            for position, (start, stop) in enumerate(offsets):
                if start == stop or not start <= index < stop:
                    continue
                if first is None:
                    first = position
                last = position
            if first is not None:
                starts.append((first, last + 1))
            else:
                starts.append(None)
        if query and all(window is None for window in starts):
            raise ValueError(f"tokenizer covers nothing in {query!r}")
        # Word-piece tokenizers normalize some characters away (sentencepiece
        # collapses repeated whitespace). Those characters inherit the nearest
        # covered character's window so the character grid stays intact; the
        # count is recorded on the model and reported in the run config.
        uncovered = 0
        for index in range(len(query)):
            if starts[index] is not None:
                continue
            uncovered += 1
            previous = next((starts[position] for position in range(index - 1, -1, -1)
                             if starts[position] is not None), None)
            following = next((starts[position] for position in range(index + 1, len(query))
                              if starts[position] is not None), None)
            starts[index] = previous or following
        self.uncovered_chars = getattr(self, "uncovered_chars", 0) + uncovered
        ends = [window[1] for window in starts]
        starts = [window[0] for window in starts]
        return (ids, starts, ends,
                self._labels(query, row["spans"], row.get("partial", False)), row)

    def character_embeddings(self, ids):
        fallback = ids >= self.base_vocab_size
        known = ids.masked_fill(fallback, self.tokenizer.unk_token_id)
        embedded = self.encoder.embeddings.word_embeddings(known)
        if self.representation == "legacy":
            return embedded
        slots = utf8_slots(ids - self.base_vocab_size)
        composed = torch.cat([embedding(slots[..., i])
                              for i, embedding in enumerate(self.byte_slots)], dim=-1)
        composed = self.byte_projection(composed.reshape(-1, composed.shape[-1])).reshape_as(embedded)
        embedded = torch.where(fallback.unsqueeze(-1), composed, embedded)
        if not self.char_classes:
            return embedded
        codepoint = torch.where(fallback, ids - self.base_vocab_size,
                                self.id_codepoint[ids.clamp(0, self.base_vocab_size - 1)])
        is_char = codepoint >= 0
        safe = codepoint.clamp(0, charclass.LIMIT - 1)
        category = torch.where(is_char, self.category_table[safe], 0)
        script = torch.where(is_char, self.script_table[safe], 0)
        structural, other, unseen = self.identity_dropout
        if self.training and (structural or other or unseen):
            rate = torch.where(self.structural_category[category], structural, other)
            rate = torch.where(fallback, torch.clamp(rate, min=unseen), rate)
            dropped = is_char & (torch.rand(rate.shape, device=rate.device) < rate)
            embedded = torch.where(dropped.unsqueeze(-1), self.anonymous.to(embedded.dtype),
                                   embedded)
        return embedded + self.category_embedding(category) + self.script_embedding(script)

    def forward(self, batch):
        ids, attention, starts, ends = batch
        if self.representation in ("legacy", "subword"):
            hidden = self.encoder(input_ids=ids, attention_mask=attention).last_hidden_state
        else:
            hidden = self.encoder(inputs_embeds=self.character_embeddings(ids),
                                  attention_mask=attention).last_hidden_state
        cumulative = torch.cat((torch.zeros_like(hidden[:, :1]), hidden.cumsum(1)), 1)
        a = starts.unsqueeze(-1).expand(-1, -1, hidden.shape[-1])
        b = ends.unsqueeze(-1).expand_as(a)
        chars = (cumulative.gather(1, b) - cumulative.gather(1, a)) / (
            ends - starts).clamp_min(1).unsqueeze(-1)
        logits = self.tags(chars.reshape(-1, chars.shape[-1]))
        return logits.reshape(*chars.shape[:2], -1), chars

    def boundary_logits(self, chars):
        flat = chars.reshape(-1, chars.shape[-1])
        shape = (*chars.shape[:2], len(TYPES))
        return self.starts(flat).reshape(shape), self.ends(flat).reshape(shape)


def span_targets(labels):
    """Type/start/end class targets, respecting adjacent same-type gold spans."""
    valid = labels >= 0
    types = torch.where(valid, (labels + 1) // 2, -100)
    starts = torch.where(valid, torch.where(labels % 2 == 1, types, 0), -100)
    following = torch.cat((labels[:, 1:], torch.zeros_like(labels[:, :1])), dim=1)
    continues = (following == 2 * types) & (types > 0)
    ends = torch.where(valid, torch.where((types > 0) & ~continues, types, 0), -100)
    return types, starts, ends


def class_weights(targets, size):
    """Smoothed inverse-square-root weights; unseen classes never have zero loss."""
    valid = targets[targets >= 0]
    counts = torch.bincount(valid, minlength=size).float().add(1)
    weights = counts.rsqrt()
    return weights / (weights * counts).sum() * counts.sum()


def loss(model, logits, chars, labels, weights, *, semicrf=0.0, lengths=None,
         teacher=None, kd_weight=0.0, kd_temperature=1.0):
    """Factorised span/BIO cross-entropy, plus the optional structured and KD terms.

    `weights` entries may be None (uniform classes). `semicrf` adds the globally
    normalised interval objective for fully labelled rows; `teacher` holds the
    teacher's (type, start, end) logits for distillation on labelled characters.
    """
    if model.head == "bio":
        return nn.functional.cross_entropy(logits.reshape(-1, len(LABELS)),
                                           labels.reshape(-1), weight=weights[0])
    start, end = model.boundary_logits(chars)
    targets = span_targets(labels)
    terms = [nn.functional.cross_entropy(pred.reshape(-1, len(TYPES)), target.reshape(-1),
                                        weight=weight)
             for pred, target, weight in zip((logits, start, end), targets, weights)]
    total = sum(terms) / len(terms)
    if semicrf:
        total = total + semicrf * semicrf_loss(logits, start, end, labels, lengths)
    if teacher is not None and kd_weight:
        total = total + kd_weight * distillation_loss((logits, start, end), teacher, labels,
                                                      kd_temperature)
    return total


def interval_potentials(type_logits, start_logits, end_logits):
    """Per-kind log-odds against the null class: the scores `interval_decode` sums."""
    return [x[..., 1:] - x[..., :1] for x in (type_logits, start_logits, end_logits)]


def semicrf_log_partition(body, start, end, lengths):
    """log sum over every segmentation of exp(score), batched.

    A segmentation is background characters (score 0) and non-overlapping typed
    spans scoring interior body + start + end, exactly the space `interval_decode`
    maximises over. O(L^2 K) per row.
    """
    batch, length, kinds = body.shape
    prefix = torch.cat((body.new_zeros(batch, 1, kinds), body.cumsum(1)), dim=1)
    alpha = [body.new_zeros(batch)]
    for stop in range(1, length + 1):
        previous = torch.stack(alpha, dim=1)                       # [B, stop]
        spans = (previous.unsqueeze(-1) + prefix[:, stop:stop + 1] - prefix[:, :stop]
                 + start[:, :stop] + end[:, stop - 1:stop])        # [B, stop, K]
        alpha.append(torch.logaddexp(alpha[-1], spans.reshape(batch, -1).logsumexp(-1)))
    alpha = torch.stack(alpha, dim=1)
    return alpha.gather(1, lengths.unsqueeze(1)).squeeze(1)


def semicrf_loss(type_logits, start_logits, end_logits, labels, lengths):
    """Per-character negative log-likelihood of the gold segmentation.

    Rows carrying unlabelled characters (partial rows) have no single gold
    segmentation and are left to the factorised terms.
    """
    positions = torch.arange(labels.shape[1], device=labels.device)
    inside = positions.unsqueeze(0) < lengths.unsqueeze(1)
    full = ~((labels < 0) & inside).any(1) & (lengths > 0)
    if not full.any():
        return type_logits.sum() * 0
    body, start, end = interval_potentials(type_logits[full], start_logits[full],
                                           end_logits[full])
    labels, lengths = labels[full], lengths[full]
    types, starts, ends = span_targets(labels)

    def gathered(scores, target):
        picked = scores.gather(2, (target.clamp_min(1) - 1).unsqueeze(-1)).squeeze(-1)
        return (picked * (target > 0)).sum(1)

    gold = gathered(body, types) + gathered(start, starts) + gathered(end, ends)
    partition = semicrf_log_partition(body, start, end, lengths)
    return ((partition - gold) / lengths.clamp_min(1)).mean()


def distillation_loss(student, teacher, labels, temperature):
    """Temperature-scaled KL to the teacher, on labelled characters only.

    Characters a row leaves unlabelled are excluded, so no teacher prediction ever
    stands in for a missing gold label.
    """
    mask = (labels >= 0).reshape(-1)
    if not mask.any():
        return student[0].sum() * 0
    total = 0
    for s, t in zip(student, teacher):
        s = s.reshape(-1, s.shape[-1])[mask] / temperature
        t = t.reshape(-1, t.shape[-1])[mask] / temperature
        total = total + nn.functional.kl_div(s.log_softmax(-1), t.log_softmax(-1),
                                             log_target=True, reduction="batchmean")
    return total * temperature ** 2 / len(student)


def record_confidence(type_logits, start_logits, end_logits, length):
    """exp(score(best segmentation) - log Z): the semi-CRF probability of the decode.

    A score for ranking and thresholds, not a calibrated probability unless the
    model was trained with the structured term and checked for calibration.
    """
    body, start, end = interval_potentials(type_logits[:length].unsqueeze(0),
                                           start_logits[:length].unsqueeze(0),
                                           end_logits[:length].unsqueeze(0))
    spans = interval_decode(type_logits[:length], start_logits[:length], end_logits[:length])
    best = 0.0
    for span in spans:
        k = KINDS.index(span["kind"])
        best += float(body[0, span["start"]:span["end"], k].sum()
                      + start[0, span["start"], k] + end[0, span["end"] - 1, k])
    partition = semicrf_log_partition(body, start, end,
                                      torch.tensor([length], device=body.device))
    return float(torch.exp(torch.tensor(best) - partition.cpu()[0]))


def interval_decode(type_logits, start_logits, end_logits):
    """Best non-overlapping complete spans, with an explicit all-background path.

    A candidate's score is interior type-vs-background log odds plus type-vs-null
    odds at its start and end. No vocabulary, delimiter, or position rule enters
    decoding. Adjacent spans are legal, including two of the same type.
    """
    arrays = [x.detach().cpu().numpy() if isinstance(x, torch.Tensor) else np.asarray(x)
              for x in (type_logits, start_logits, end_logits)]
    if any(a.ndim != 2 or a.shape != arrays[0].shape for a in arrays):
        raise ValueError("span logits must be aligned matrices")
    if not all(np.isfinite(a).all() for a in arrays):
        raise ValueError("non-finite span logits")
    body, start, end = [a[:, 1:] - a[:, :1] for a in arrays]
    length, kinds = body.shape
    prefix = np.concatenate((np.zeros((1, kinds)), body.cumsum(0)), axis=0)
    best, back = np.zeros(length + 1), [None] * (length + 1)
    for stop in range(1, length + 1):
        scores = (best[:stop, None] + prefix[stop] - prefix[:stop]
                  + start[:stop] + end[stop - 1])
        index = int(scores.argmax())
        begin, kind = divmod(index, kinds)
        if scores[begin, kind] > best[stop - 1]:
            best[stop], back[stop] = scores[begin, kind], (begin, kind)
        else:
            best[stop] = best[stop - 1]
    spans, stop = [], length
    while stop:
        if back[stop] is None:
            stop -= 1
        else:
            begin, kind = back[stop]
            spans.append({"kind": KINDS[kind], "start": begin, "end": stop})
            stop = begin
    return list(reversed(spans))


def decode_batch(model, logits, chars, items, decoder):
    lengths = [len(item[-1]["query"]) for item in items]
    if model.head == "bio":
        paths = tag_paths(logits.detach().cpu(), lengths, decoder, legality_matrix())
        return [predict_spans(item[-1]["query"], path) for item, path in zip(items, paths)]
    if decoder != "span":
        raise ValueError("span head requires span decoding")
    starts, ends = model.boundary_logits(chars)
    return [interval_decode(logits[i, :n], starts[i, :n], ends[i, :n])
            for i, n in enumerate(lengths)]