"""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)]