Upload ContinuaRoformer-slav

81e263e verified 6 days ago

19.8 kB

"""
Text processing and damage simulation for ContinuaRoformer-slav.
Text
  tokenize_text    raw text -> char ids (lowercased; `[...]` special tokens kept)
  Continua         char ids -> scriptio continua ids + per-char boundary labels:
                   spaces removed (a space after a char becomes its boundary
                   label 1), every punctuation char (Unicode P*) folded into the
                   divider `·`, runs of dividers collapsed
  ngram_ids        id of the char n-gram STARTING at each position, computed on
                   the model input: an n-gram touching a special token
                   ([-], [#], [UNK], [SOS], ...) or the end gets BLANK, so a
                   masked char never leaks through the n-gram channel
Damage (ContinuaCollator), per window, in this order:
  1. crop          random length crop_min..max_len in train, max_len in eval;
                   in train, with short_crop_prob, a short crop of
                   short_crop_min..short_crop_max chars
  2. [UNK]         synthetic unrecoverable damage in every source: target
                   density d ~ U(0, unk_aug_max) [UNK] per char, real [UNK]s
                   counted; each removes Geom(unk_aug_p) chars, no target
  3. [#]           one gap of unknown length (or, with edge_prob, a torn edge),
                   length Geom(unk_geometric_p) - 1; target for the unk head:
                   1 if it hides more than one char
  4. [-]           lacunae covering a share U(0, mask_rate_max) of maskable
                   chars (eval: U(eval_mask_rate_min, mask_rate_max)), length
                   (1-w)Geom(lacuna_p_short) + w Geom(lacuna_p_long), never
                   touching each other
Eval mode ('valid') draws everything from a document-seeded RNG, so eval and
test_a get identical damage on every call.
"""

import json
import math
import random
import unicodedata
from pathlib import Path
from typing import Any, Dict, List, Optional, Set, Tuple

import numpy as np
import torch

PAD_NGRAM_ID, OOV_NGRAM_ID, BLANK_NGRAM_ID, FIRST_NGRAM_ID = 0, 1, 2, 3
DIVIDER = "·"
# letters, space, word dividers (`·`, `:`) and `+` can be masked and predicted
MASKABLE_CATEGORIES = ("Ll", "Lu", "Lo", "Zs", "Po", "Sm")


# ---- text --------------------------------------------------------------------


def load_char_vocab(path: str) -> Dict[str, int]:
    return json.loads(Path(path).read_text(encoding="utf-8"))


def special_ids(char_vocab: Dict[str, int]) -> Set[int]:
    return {v for k, v in char_vocab.items() if k.startswith("[") and k.endswith("]")}


def tokenize_text(text: str, vocab: Dict[str, int]) -> List[int]:
    unk = vocab.get("[UNK]", 1)
    out, i = [], 0
    while i < len(text):
        if text[i] == "[":
            end = text.find("]", i)
            if end != -1 and text[i:end + 1] in vocab:
                out.append(vocab[text[i:end + 1]])
                i = end + 1
                continue
        out.append(vocab.get(text[i].lower(), unk))
        i += 1
    return out


class Continua:
    """Spaced, punctuated ids -> continua ids (+ labels) + boundary labels
    (1 = a space followed this char, 0 = not, -100 = special token / last)."""

    def __init__(self, char_vocab: Dict[str, int]):
        self.space_id = char_vocab[" "]
        self.divider_id = char_vocab[DIVIDER]
        size = max(char_vocab.values()) + 1
        self.remap = np.arange(size)
        for k, v in char_vocab.items():
            if len(k) == 1 and unicodedata.category(k).startswith("P"):
                self.remap[v] = self.divider_id
        self.is_special = np.zeros(size, dtype=bool)
        self.is_special[list(special_ids(char_vocab))] = True
        # ids that can no longer occur in model input or as a target
        self.removed_ids = {self.space_id} | {int(v) for v in np.nonzero(self.remap != np.arange(size))[0]}

    def __call__(self, tokens: List[int], labels: Optional[List[int]] = None):
        t = self.remap[np.asarray(tokens, dtype=np.int64)]
        lab = None
        if labels is not None:
            lab = np.asarray(labels, dtype=np.int64)
            lab = np.where(lab != -100, self.remap[np.maximum(lab, 0)], -100)
        true = t if lab is None else np.where(lab != -100, lab, t)

        is_space = true == self.space_id
        keep = np.nonzero(~is_space)[0]
        is_div = true[keep] == self.divider_id
        dup = np.zeros(len(keep), dtype=bool)
        dup[1:] = is_div[1:] & is_div[:-1]
        keep = keep[~dup]

        # boundary[k] = a space strictly between keep[k] and keep[k+1]
        cs = np.concatenate([[0], np.cumsum(is_space)])
        boundary = np.full(len(keep), -100, dtype=np.int64)
        if len(keep) > 1:
            boundary[:-1] = (cs[keep[1:]] - cs[keep[:-1] + 1] > 0).astype(np.int64)
        boundary[self.is_special[true[keep]]] = -100
        return t[keep].tolist(), (lab[keep].tolist() if lab is not None else None), boundary.tolist()


def maskable_ids(char_vocab: Dict[str, int], continua: Continua) -> Set[int]:
    """Chars that can be masked in continua text and that the restore head may
    predict: MASKABLE_CATEGORIES minus spaces and folded punctuation."""
    ids = {int(v) for k, v in char_vocab.items() if len(k) == 1 and unicodedata.category(k) in MASKABLE_CATEGORIES}
    return ids - continua.removed_ids


def load_ngram_vocab(path: str) -> Tuple[Dict[Tuple[int, ...], int], Dict]:
    raw = json.loads(Path(path).read_text(encoding="utf-8"))
    return {tuple(int(t) for t in k.split(",")): int(v) for k, v in raw["ngrams"].items()}, raw


def ngram_ids(tokens: List[int], vocab: Dict[Tuple[int, ...], int], n: int, specials: Set[int]) -> List[int]:
    bad = [t in specials for t in tokens]
    out = [BLANK_NGRAM_ID] * len(tokens)
    for i in range(len(tokens) - n + 1):
        if not any(bad[i:i + n]):
            out[i] = vocab.get(tuple(tokens[i:i + n]), OOV_NGRAM_ID)
    return out


def get_source(metadata: Optional[str]) -> str:
    try:
        return json.loads(metadata).get("source_dataset") or "unknown"
    except (TypeError, ValueError, AttributeError):
        return "unknown"


# ---- damage simulation ----------------------------------------------------------


def _geometric(rng, p: float, cap: int) -> int:
    """1, 2, ... with P(L=k) = p(1-p)^(k-1), capped."""
    return min(cap, 1 + int(math.log(1.0 - rng.random()) / math.log(1.0 - p)))


class ContinuaCollator:
    def __init__(
        self,
        char_vocab: Dict[str, int],
        ngram_vocab: Dict[Tuple[int, ...], int],
        ngram_n: int = 3,
        mode: str = "train",
        crop_max_len: Optional[int] = 512,
        crop_min_len: int = 128,
        short_crop_prob: float = 0.2,
        short_crop_min: int = 30,
        short_crop_max: int = 100,
        unk_aug_prob: float = 1.0,
        unk_aug_max: float = 0.04,
        unk_aug_p: float = 0.2,
        unk_geometric_p: float = 0.25,
        edge_prob: float = 0.1,
        mask_rate_max: float = 0.30,
        eval_mask_rate_min: float = 0.02,
        lacuna_p_short: float = 0.6,
        lacuna_p_long: float = 0.2,
        lacuna_long_weight: float = 0.1,
        lacuna_max_len: int = 30,
        date_bins: int = 20,
        pad_to_len: Optional[int] = None,
        eval_seed: int = 42,
    ) -> None:
        self.vocab = char_vocab
        self.pad_id = char_vocab["[PAD]"]
        self.mask_id = char_vocab["[-]"]  # one missing char, restore target
        self.gap_id = char_vocab["[#]"]  # gap of unknown length
        self.unk_id = char_vocab["[UNK]"]  # unrecoverable damage, no target
        self.sos_id = char_vocab["[SOS]"]
        self.special_ids = special_ids(char_vocab)
        self.continua = Continua(char_vocab)
        self.maskable_ids = maskable_ids(char_vocab, self.continua)
        self.ngram_vocab, self.ngram_n = ngram_vocab, ngram_n

        self.mode = mode
        self.crop_max_len, self.crop_min_len = crop_max_len, crop_min_len
        self.short_crop_prob, self.short_crop_min, self.short_crop_max = short_crop_prob, short_crop_min, short_crop_max
        self.unk_aug_prob, self.unk_aug_max, self.unk_aug_p = unk_aug_prob, unk_aug_max, unk_aug_p
        self.unk_geometric_p, self.edge_prob = unk_geometric_p, edge_prob
        self.mask_rate_max, self.eval_mask_rate_min = mask_rate_max, eval_mask_rate_min
        self.lacuna_p_short, self.lacuna_p_long = lacuna_p_short, lacuna_p_long
        self.lacuna_long_weight, self.lacuna_max_len = lacuna_long_weight, lacuna_max_len
        self.date_bins, self.pad_to_len, self.eval_seed = date_bins, pad_to_len, eval_seed

    def ngram_ids(self, tokens: List[int]) -> List[int]:
        return ngram_ids(tokens, self.ngram_vocab, self.ngram_n, self.special_ids)

    def to_continua(self, feature: Dict[str, Any]) -> Tuple[List[int], List[int]]:
        raw = feature["input_ids"] if "input_ids" in feature else tokenize_text(feature["text"], self.vocab)
        tokens, _, boundary = self.continua(list(raw))
        return tokens, boundary

    # 1. crop
    def crop_bounds(self, n: int, rng) -> Tuple[int, int]:
        if self.mode == "train" and n > self.short_crop_max and rng.random() < self.short_crop_prob:
            crop_len = rng.randint(self.short_crop_min, self.short_crop_max)
        elif self.crop_max_len is None or n <= self.crop_max_len:
            return 0, n
        elif self.mode == "train":
            crop_len = rng.randint(min(self.crop_min_len, self.crop_max_len), self.crop_max_len)
        else:
            crop_len = self.crop_max_len
        start = rng.randint(0, n - crop_len)
        return start, start + crop_len

    # 2. synthetic [UNK]
    def _inject_unk(self, tokens: List[int], boundary: List[int], rng) -> Tuple[List[int], List[int]]:
        if self.unk_aug_max <= 0 or rng.random() >= self.unk_aug_prob:
            return tokens, boundary
        n_text = len(tokens) - 1
        n_gaps = int(round(rng.uniform(0, self.unk_aug_max) * n_text)) - sum(t == self.unk_id for t in tokens)
        removed = 0
        for _ in range(n_gaps):
            L = _geometric(rng, self.unk_aug_p, self.lacuna_max_len)
            if removed + L > n_text // 2:
                break
            starts = [
                s for s in range(1, len(tokens) - L + 1)
                if not any(t in self.special_ids for t in tokens[s:s + L])
                and tokens[s - 1] != self.unk_id
                and (s + L == len(tokens) or tokens[s + L] != self.unk_id)
            ]
            if not starts:
                continue
            s = rng.choice(starts)
            tokens = tokens[:s] + [self.unk_id] + tokens[s + L:]
            boundary = boundary[:s] + [-100] + boundary[s + L:]
            removed += L
        return tokens, boundary

    # 3. the [#] gap
    def _pick_gap(self, tokens: List[int], rng, np_rng) -> Optional[Tuple[List[int], bool]]:
        """(positions, hides more than one char) or None."""
        if rng.random() < self.edge_prob and len(tokens) > 10:
            edge_len = rng.randint(2, 20)
            if rng.random() < 0.5:
                start = 1
                while start < len(tokens) and tokens[start] in self.special_ids:
                    start += 1
                end = min(start + edge_len, len(tokens))
            else:
                end = len(tokens)
                start = max(1, end - edge_len)
            positions = [i for i in range(start, end) if tokens[i] not in self.special_ids]
            if positions:
                return positions, len(positions) > 1

        span_len = int(np_rng.geometric(self.unk_geometric_p)) - 1
        if span_len <= 0:
            return None
        candidates = [i for i in range(1, len(tokens)) if tokens[i] not in self.special_ids]
        if len(candidates) < span_len:
            return None
        for _ in range(50):  # a gap must be one unbroken run of positions
            s = rng.randint(0, len(candidates) - span_len)
            positions = candidates[s:s + span_len]
            if positions[-1] - positions[0] == span_len - 1:
                return positions, span_len > 1
        return None

    # 4. [-] lacunae
    def _sample_lacunae(self, tokens: List[int], excluded: set, rng) -> List[int]:
        ok = [i > 0 and tokens[i] in self.maskable_ids and i not in excluded for i in range(len(tokens))]
        n_free = sum(ok)
        if n_free == 0:
            return []
        lo = 0.0 if self.mode == "train" else self.eval_mask_rate_min
        budget = int(round(rng.uniform(lo, self.mask_rate_max) * n_free))
        if self.mode != "train":
            budget = max(budget, 1)  # every eval window scores at least one char

        taken = [False] * len(tokens)
        for _ in range(500):
            if budget <= 0:
                break
            p = self.lacuna_p_long if rng.random() < self.lacuna_long_weight else self.lacuna_p_short
            L = min(_geometric(rng, p, self.lacuna_max_len), budget)
            runs, start = [], None  # free stretches that touch no existing lacuna
            for i in range(len(tokens) + 1):
                allowed = (
                    i < len(tokens) and ok[i] and not taken[i]
                    and not (i > 0 and taken[i - 1]) and not (i + 1 < len(tokens) and taken[i + 1])
                )
                if allowed and start is None:
                    start = i
                elif not allowed and start is not None:
                    if i - start >= L:
                        runs.append((start, i))
                    start = None
            if not runs:
                if L == 1:
                    break
                continue
            s, e = rng.choice(runs)
            st = rng.randint(s, e - L)
            for j in range(st, st + L):
                taken[j] = True
            budget -= L
        return [i for i, t in enumerate(taken) if t]

    def sample_span(self, tokens: List[int], excluded: set, span_len: int, rng) -> List[int]:
        """One contiguous run of span_len maskable positions (used by evaluate.py)."""
        runs, start = [], None
        for i, tok in enumerate(tokens + [None]):
            ok = tok in self.maskable_ids and i not in excluded
            if ok and start is None:
                start = i
            elif not ok and start is not None:
                runs.append((start, i))
                start = None
        runs = [r for r in runs if r[1] - r[0] >= span_len]
        if span_len <= 0 or not runs:
            return []
        s, e = rng.choice(runs)
        st = rng.randint(s, e - span_len)
        return list(range(st, st + span_len))

    def process(self, tokens: List[int], boundary: List[int]):
        """One document (continua ids + boundary labels) -> (input ids, restore
        labels, unk labels, boundary labels, ngram ids)."""
        if self.mode == "valid":
            seed = self.eval_seed ^ hash(tuple(tokens))
            rng, np_rng = random.Random(seed), np.random.RandomState(seed & 0xFFFFFFFF)
        else:
            rng, np_rng = random, np.random

        s, e = self.crop_bounds(len(tokens), rng)
        tokens, boundary = tokens[s:e], boundary[s:e]
        if not tokens or tokens[0] != self.sos_id:
            tokens, boundary = [self.sos_id] + tokens, [-100] + boundary

        tokens, boundary = self._inject_unk(tokens, boundary, rng)
        gap = self._pick_gap(tokens, rng, np_rng)
        gap_pos = set(gap[0]) if gap else set()
        masked = set(self._sample_lacunae(tokens, {0} | gap_pos, rng))

        out_t, out_res, out_unk, out_bnd = [], [], [], []
        i = 0
        while i < len(tokens):
            if i in gap_pos:  # the whole gap becomes one [#]
                out_t.append(self.gap_id)
                out_res.append(-100)
                out_unk.append(1 if gap[1] else 0)
                out_bnd.append(-100)
                i = max(gap_pos) + 1
                continue
            if i in masked:
                out_t.append(self.mask_id)
                out_res.append(tokens[i])
            else:
                out_t.append(tokens[i])
                out_res.append(-100)
            out_unk.append(-100)
            out_bnd.append(boundary[i])
            i += 1
        return out_t, out_res, out_unk, out_bnd, self.ngram_ids(out_t)

    def __call__(self, features: List[Dict[str, Any]]) -> Dict[str, torch.Tensor]:
        rows = [self.process(*self.to_continua(f)) for f in features]
        batch = pad_batch(rows, self.pad_id, self.pad_to_len)
        batch.update(attribution_labels(features, self.date_bins))
        return batch


def pad_batch(rows, pad_id: int, pad_to_len: Optional[int] = None) -> Dict[str, torch.Tensor]:
    """rows: (input ids, restore labels, unk labels, boundary labels, ngram ids)."""
    T = max(len(r[0]) for r in rows)
    if pad_to_len is not None:
        T = max(T, pad_to_len)
    fills = {"input_ids": pad_id, "labels": -100, "unk_labels": -100, "boundary_labels": -100, "ngram_ids": PAD_NGRAM_ID}
    batch = {k: torch.full((len(rows), T), v, dtype=torch.long) for k, v in fills.items()}
    batch["attention_mask"] = torch.zeros((len(rows), T), dtype=torch.long)
    for b, row in enumerate(rows):
        n = len(row[0])
        for key, values in zip(("input_ids", "labels", "unk_labels", "boundary_labels", "ngram_ids"), row):
            batch[key][b, :n] = torch.tensor(values)
        batch["attention_mask"][b, :n] = 1
    return batch


def attribution_labels(features: List[Dict[str, Any]], date_bins: int = 20,
                       zero_date_is_missing: bool = False) -> Dict[str, torch.Tensor]:
    """date_labels (+ date_labels_mask: 0 = no date) and region_labels
    (-100 = no label). A missing date is None; test_b rows store an all-zero
    vector instead (zero_date_is_missing=True)."""
    out = {}
    if "date_labels" in features[0]:
        dates = [f.get("date_labels") for f in features]
        ok = [d is not None and (sum(d) > 0 or not zero_date_is_missing) for d in dates]
        out["date_labels"] = torch.tensor([d if k else [0.0] * date_bins for d, k in zip(dates, ok)], dtype=torch.float)
        out["date_labels_mask"] = torch.tensor(ok, dtype=torch.float)
    if "region_labels" in features[0]:
        out["region_labels"] = torch.tensor(
            [f["region_labels"] if f["region_labels"] is not None else -100 for f in features], dtype=torch.long
        )
    return out


def prepare_test_b(dataset: Any, collator: ContinuaCollator) -> Any:
    """test_b's real lacunae are baked in ([-] in input_ids, true chars in
    labels): convert each row to continua once (a masked space is dropped, a
    masked punctuation mark folded into `·`) and add ngram_ids and
    boundary_labels."""

    def convert(ex: Dict[str, Any]) -> Dict[str, Any]:
        keep = [j for j, t in enumerate(ex["input_ids"]) if t != collator.pad_id]
        toks, labels, boundary = collator.continua([ex["input_ids"][j] for j in keep], [ex["labels"][j] for j in keep])
        return {"input_ids": toks, "labels": labels, "attention_mask": [1] * len(toks),
                "boundary_labels": boundary, "ngram_ids": collator.ngram_ids(toks)}

    return dataset.map(convert)


class TestBCollator:
    """Pads prepared test_b rows into a batch (no damage is added)."""

    def __init__(self, pad_id: int, date_bins: int = 20):
        self.pad_id, self.date_bins = pad_id, date_bins

    def __call__(self, features: List[Dict[str, Any]]) -> Dict[str, torch.Tensor]:
        rows = [(f["input_ids"], f["labels"], f.get("unk_labels") or [-100] * len(f["input_ids"]),
                 f["boundary_labels"], f["ngram_ids"]) for f in features]
        batch = pad_batch(rows, self.pad_id)
        batch.update(attribution_labels(features, self.date_bins, zero_date_is_missing=True))
        return batch

	"""
	Text processing and damage simulation for ContinuaRoformer-slav.

	Text
	tokenize_text raw text -> char ids (lowercased; `[...]` special tokens kept)
	Continua char ids -> scriptio continua ids + per-char boundary labels:
	spaces removed (a space after a char becomes its boundary
	label 1), every punctuation char (Unicode P*) folded into the
	divider `·`, runs of dividers collapsed
	ngram_ids id of the char n-gram STARTING at each position, computed on
	the model input: an n-gram touching a special token
	([-], [#], [UNK], [SOS], ...) or the end gets BLANK, so a
	masked char never leaks through the n-gram channel

	Damage (ContinuaCollator), per window, in this order:
	1. crop random length crop_min..max_len in train, max_len in eval;
	in train, with short_crop_prob, a short crop of
	short_crop_min..short_crop_max chars
	2. [UNK] synthetic unrecoverable damage in every source: target
	density d ~ U(0, unk_aug_max) [UNK] per char, real [UNK]s
	counted; each removes Geom(unk_aug_p) chars, no target
	3. [#] one gap of unknown length (or, with edge_prob, a torn edge),
	length Geom(unk_geometric_p) - 1; target for the unk head:
	1 if it hides more than one char
	4. [-] lacunae covering a share U(0, mask_rate_max) of maskable
	chars (eval: U(eval_mask_rate_min, mask_rate_max)), length
	(1-w)Geom(lacuna_p_short) + w Geom(lacuna_p_long), never
	touching each other
	Eval mode ('valid') draws everything from a document-seeded RNG, so eval and
	test_a get identical damage on every call.
	"""

	import json
	import math
	import random
	import unicodedata
	from pathlib import Path
	from typing import Any, Dict, List, Optional, Set, Tuple

	import numpy as np
	import torch

	PAD_NGRAM_ID, OOV_NGRAM_ID, BLANK_NGRAM_ID, FIRST_NGRAM_ID = 0, 1, 2, 3
	DIVIDER = "·"
	# letters, space, word dividers (`·`, `:`) and `+` can be masked and predicted
	MASKABLE_CATEGORIES = ("Ll", "Lu", "Lo", "Zs", "Po", "Sm")


	# ---- text --------------------------------------------------------------------


	def load_char_vocab(path: str) -> Dict[str, int]:
	return json.loads(Path(path).read_text(encoding="utf-8"))


	def special_ids(char_vocab: Dict[str, int]) -> Set[int]:
	return {v for k, v in char_vocab.items() if k.startswith("[") and k.endswith("]")}


	def tokenize_text(text: str, vocab: Dict[str, int]) -> List[int]:
	unk = vocab.get("[UNK]", 1)
	out, i = [], 0
	while i < len(text):
	if text[i] == "[":
	end = text.find("]", i)
	if end != -1 and text[i:end + 1] in vocab:
	out.append(vocab[text[i:end + 1]])
	i = end + 1
	continue
	out.append(vocab.get(text[i].lower(), unk))
	i += 1
	return out


	class Continua:
	"""Spaced, punctuated ids -> continua ids (+ labels) + boundary labels
	(1 = a space followed this char, 0 = not, -100 = special token / last)."""

	def __init__(self, char_vocab: Dict[str, int]):
	self.space_id = char_vocab[" "]
	self.divider_id = char_vocab[DIVIDER]
	size = max(char_vocab.values()) + 1
	self.remap = np.arange(size)
	for k, v in char_vocab.items():
	if len(k) == 1 and unicodedata.category(k).startswith("P"):
	self.remap[v] = self.divider_id
	self.is_special = np.zeros(size, dtype=bool)
	self.is_special[list(special_ids(char_vocab))] = True
	# ids that can no longer occur in model input or as a target
	self.removed_ids = {self.space_id} \| {int(v) for v in np.nonzero(self.remap != np.arange(size))[0]}

	def __call__(self, tokens: List[int], labels: Optional[List[int]] = None):
	t = self.remap[np.asarray(tokens, dtype=np.int64)]
	lab = None
	if labels is not None:
	lab = np.asarray(labels, dtype=np.int64)
	lab = np.where(lab != -100, self.remap[np.maximum(lab, 0)], -100)
	true = t if lab is None else np.where(lab != -100, lab, t)

	is_space = true == self.space_id
	keep = np.nonzero(~is_space)[0]
	is_div = true[keep] == self.divider_id
	dup = np.zeros(len(keep), dtype=bool)
	dup[1:] = is_div[1:] & is_div[:-1]
	keep = keep[~dup]

	# boundary[k] = a space strictly between keep[k] and keep[k+1]
	cs = np.concatenate([[0], np.cumsum(is_space)])
	boundary = np.full(len(keep), -100, dtype=np.int64)
	if len(keep) > 1:
	boundary[:-1] = (cs[keep[1:]] - cs[keep[:-1] + 1] > 0).astype(np.int64)
	boundary[self.is_special[true[keep]]] = -100
	return t[keep].tolist(), (lab[keep].tolist() if lab is not None else None), boundary.tolist()


	def maskable_ids(char_vocab: Dict[str, int], continua: Continua) -> Set[int]:
	"""Chars that can be masked in continua text and that the restore head may
	predict: MASKABLE_CATEGORIES minus spaces and folded punctuation."""
	ids = {int(v) for k, v in char_vocab.items() if len(k) == 1 and unicodedata.category(k) in MASKABLE_CATEGORIES}
	return ids - continua.removed_ids


	def load_ngram_vocab(path: str) -> Tuple[Dict[Tuple[int, ...], int], Dict]:
	raw = json.loads(Path(path).read_text(encoding="utf-8"))
	return {tuple(int(t) for t in k.split(",")): int(v) for k, v in raw["ngrams"].items()}, raw


	def ngram_ids(tokens: List[int], vocab: Dict[Tuple[int, ...], int], n: int, specials: Set[int]) -> List[int]:
	bad = [t in specials for t in tokens]
	out = [BLANK_NGRAM_ID] * len(tokens)
	for i in range(len(tokens) - n + 1):
	if not any(bad[i:i + n]):
	out[i] = vocab.get(tuple(tokens[i:i + n]), OOV_NGRAM_ID)
	return out


	def get_source(metadata: Optional[str]) -> str:
	try:
	return json.loads(metadata).get("source_dataset") or "unknown"
	except (TypeError, ValueError, AttributeError):
	return "unknown"


	# ---- damage simulation ----------------------------------------------------------


	def _geometric(rng, p: float, cap: int) -> int:
	"""1, 2, ... with P(L=k) = p(1-p)^(k-1), capped."""
	return min(cap, 1 + int(math.log(1.0 - rng.random()) / math.log(1.0 - p)))


	class ContinuaCollator:
	def __init__(
	self,
	char_vocab: Dict[str, int],
	ngram_vocab: Dict[Tuple[int, ...], int],
	ngram_n: int = 3,
	mode: str = "train",
	crop_max_len: Optional[int] = 512,
	crop_min_len: int = 128,
	short_crop_prob: float = 0.2,
	short_crop_min: int = 30,
	short_crop_max: int = 100,
	unk_aug_prob: float = 1.0,
	unk_aug_max: float = 0.04,
	unk_aug_p: float = 0.2,
	unk_geometric_p: float = 0.25,
	edge_prob: float = 0.1,
	mask_rate_max: float = 0.30,
	eval_mask_rate_min: float = 0.02,
	lacuna_p_short: float = 0.6,
	lacuna_p_long: float = 0.2,
	lacuna_long_weight: float = 0.1,
	lacuna_max_len: int = 30,
	date_bins: int = 20,
	pad_to_len: Optional[int] = None,
	eval_seed: int = 42,
	) -> None:
	self.vocab = char_vocab
	self.pad_id = char_vocab["[PAD]"]
	self.mask_id = char_vocab["[-]"] # one missing char, restore target
	self.gap_id = char_vocab["[#]"] # gap of unknown length
	self.unk_id = char_vocab["[UNK]"] # unrecoverable damage, no target
	self.sos_id = char_vocab["[SOS]"]
	self.special_ids = special_ids(char_vocab)
	self.continua = Continua(char_vocab)
	self.maskable_ids = maskable_ids(char_vocab, self.continua)
	self.ngram_vocab, self.ngram_n = ngram_vocab, ngram_n

	self.mode = mode
	self.crop_max_len, self.crop_min_len = crop_max_len, crop_min_len
	self.short_crop_prob, self.short_crop_min, self.short_crop_max = short_crop_prob, short_crop_min, short_crop_max
	self.unk_aug_prob, self.unk_aug_max, self.unk_aug_p = unk_aug_prob, unk_aug_max, unk_aug_p
	self.unk_geometric_p, self.edge_prob = unk_geometric_p, edge_prob
	self.mask_rate_max, self.eval_mask_rate_min = mask_rate_max, eval_mask_rate_min
	self.lacuna_p_short, self.lacuna_p_long = lacuna_p_short, lacuna_p_long
	self.lacuna_long_weight, self.lacuna_max_len = lacuna_long_weight, lacuna_max_len
	self.date_bins, self.pad_to_len, self.eval_seed = date_bins, pad_to_len, eval_seed

	def ngram_ids(self, tokens: List[int]) -> List[int]:
	return ngram_ids(tokens, self.ngram_vocab, self.ngram_n, self.special_ids)

	def to_continua(self, feature: Dict[str, Any]) -> Tuple[List[int], List[int]]:
	raw = feature["input_ids"] if "input_ids" in feature else tokenize_text(feature["text"], self.vocab)
	tokens, _, boundary = self.continua(list(raw))
	return tokens, boundary

	# 1. crop
	def crop_bounds(self, n: int, rng) -> Tuple[int, int]:
	if self.mode == "train" and n > self.short_crop_max and rng.random() < self.short_crop_prob:
	crop_len = rng.randint(self.short_crop_min, self.short_crop_max)
	elif self.crop_max_len is None or n <= self.crop_max_len:
	return 0, n
	elif self.mode == "train":
	crop_len = rng.randint(min(self.crop_min_len, self.crop_max_len), self.crop_max_len)
	else:
	crop_len = self.crop_max_len
	start = rng.randint(0, n - crop_len)
	return start, start + crop_len

	# 2. synthetic [UNK]
	def _inject_unk(self, tokens: List[int], boundary: List[int], rng) -> Tuple[List[int], List[int]]:
	if self.unk_aug_max <= 0 or rng.random() >= self.unk_aug_prob:
	return tokens, boundary
	n_text = len(tokens) - 1
	n_gaps = int(round(rng.uniform(0, self.unk_aug_max) * n_text)) - sum(t == self.unk_id for t in tokens)
	removed = 0
	for _ in range(n_gaps):
	L = _geometric(rng, self.unk_aug_p, self.lacuna_max_len)
	if removed + L > n_text // 2:
	break
	starts = [
	s for s in range(1, len(tokens) - L + 1)
	if not any(t in self.special_ids for t in tokens[s:s + L])
	and tokens[s - 1] != self.unk_id
	and (s + L == len(tokens) or tokens[s + L] != self.unk_id)
	]
	if not starts:
	continue
	s = rng.choice(starts)
	tokens = tokens[:s] + [self.unk_id] + tokens[s + L:]
	boundary = boundary[:s] + [-100] + boundary[s + L:]
	removed += L
	return tokens, boundary

	# 3. the [#] gap
	def _pick_gap(self, tokens: List[int], rng, np_rng) -> Optional[Tuple[List[int], bool]]:
	"""(positions, hides more than one char) or None."""
	if rng.random() < self.edge_prob and len(tokens) > 10:
	edge_len = rng.randint(2, 20)
	if rng.random() < 0.5:
	start = 1
	while start < len(tokens) and tokens[start] in self.special_ids:
	start += 1
	end = min(start + edge_len, len(tokens))
	else:
	end = len(tokens)
	start = max(1, end - edge_len)
	positions = [i for i in range(start, end) if tokens[i] not in self.special_ids]
	if positions:
	return positions, len(positions) > 1

	span_len = int(np_rng.geometric(self.unk_geometric_p)) - 1
	if span_len <= 0:
	return None
	candidates = [i for i in range(1, len(tokens)) if tokens[i] not in self.special_ids]
	if len(candidates) < span_len:
	return None
	for _ in range(50): # a gap must be one unbroken run of positions
	s = rng.randint(0, len(candidates) - span_len)
	positions = candidates[s:s + span_len]
	if positions[-1] - positions[0] == span_len - 1:
	return positions, span_len > 1
	return None

	# 4. [-] lacunae
	def _sample_lacunae(self, tokens: List[int], excluded: set, rng) -> List[int]:
	ok = [i > 0 and tokens[i] in self.maskable_ids and i not in excluded for i in range(len(tokens))]
	n_free = sum(ok)
	if n_free == 0:
	return []
	lo = 0.0 if self.mode == "train" else self.eval_mask_rate_min
	budget = int(round(rng.uniform(lo, self.mask_rate_max) * n_free))
	if self.mode != "train":
	budget = max(budget, 1) # every eval window scores at least one char

	taken = [False] * len(tokens)
	for _ in range(500):
	if budget <= 0:
	break
	p = self.lacuna_p_long if rng.random() < self.lacuna_long_weight else self.lacuna_p_short
	L = min(_geometric(rng, p, self.lacuna_max_len), budget)
	runs, start = [], None # free stretches that touch no existing lacuna
	for i in range(len(tokens) + 1):
	allowed = (
	i < len(tokens) and ok[i] and not taken[i]
	and not (i > 0 and taken[i - 1]) and not (i + 1 < len(tokens) and taken[i + 1])
	)
	if allowed and start is None:
	start = i
	elif not allowed and start is not None:
	if i - start >= L:
	runs.append((start, i))
	start = None
	if not runs:
	if L == 1:
	break
	continue
	s, e = rng.choice(runs)
	st = rng.randint(s, e - L)
	for j in range(st, st + L):
	taken[j] = True
	budget -= L
	return [i for i, t in enumerate(taken) if t]

	def sample_span(self, tokens: List[int], excluded: set, span_len: int, rng) -> List[int]:
	"""One contiguous run of span_len maskable positions (used by evaluate.py)."""
	runs, start = [], None
	for i, tok in enumerate(tokens + [None]):
	ok = tok in self.maskable_ids and i not in excluded
	if ok and start is None:
	start = i
	elif not ok and start is not None:
	runs.append((start, i))
	start = None
	runs = [r for r in runs if r[1] - r[0] >= span_len]
	if span_len <= 0 or not runs:
	return []
	s, e = rng.choice(runs)
	st = rng.randint(s, e - span_len)
	return list(range(st, st + span_len))

	def process(self, tokens: List[int], boundary: List[int]):
	"""One document (continua ids + boundary labels) -> (input ids, restore
	labels, unk labels, boundary labels, ngram ids)."""
	if self.mode == "valid":
	seed = self.eval_seed ^ hash(tuple(tokens))
	rng, np_rng = random.Random(seed), np.random.RandomState(seed & 0xFFFFFFFF)
	else:
	rng, np_rng = random, np.random

	s, e = self.crop_bounds(len(tokens), rng)
	tokens, boundary = tokens[s:e], boundary[s:e]
	if not tokens or tokens[0] != self.sos_id:
	tokens, boundary = [self.sos_id] + tokens, [-100] + boundary

	tokens, boundary = self._inject_unk(tokens, boundary, rng)
	gap = self._pick_gap(tokens, rng, np_rng)
	gap_pos = set(gap[0]) if gap else set()
	masked = set(self._sample_lacunae(tokens, {0} \| gap_pos, rng))

	out_t, out_res, out_unk, out_bnd = [], [], [], []
	i = 0
	while i < len(tokens):
	if i in gap_pos: # the whole gap becomes one [#]
	out_t.append(self.gap_id)
	out_res.append(-100)
	out_unk.append(1 if gap[1] else 0)
	out_bnd.append(-100)
	i = max(gap_pos) + 1
	continue
	if i in masked:
	out_t.append(self.mask_id)
	out_res.append(tokens[i])
	else:
	out_t.append(tokens[i])
	out_res.append(-100)
	out_unk.append(-100)
	out_bnd.append(boundary[i])
	i += 1
	return out_t, out_res, out_unk, out_bnd, self.ngram_ids(out_t)

	def __call__(self, features: List[Dict[str, Any]]) -> Dict[str, torch.Tensor]:
	rows = [self.process(*self.to_continua(f)) for f in features]
	batch = pad_batch(rows, self.pad_id, self.pad_to_len)
	batch.update(attribution_labels(features, self.date_bins))
	return batch


	def pad_batch(rows, pad_id: int, pad_to_len: Optional[int] = None) -> Dict[str, torch.Tensor]:
	"""rows: (input ids, restore labels, unk labels, boundary labels, ngram ids)."""
	T = max(len(r[0]) for r in rows)
	if pad_to_len is not None:
	T = max(T, pad_to_len)
	fills = {"input_ids": pad_id, "labels": -100, "unk_labels": -100, "boundary_labels": -100, "ngram_ids": PAD_NGRAM_ID}
	batch = {k: torch.full((len(rows), T), v, dtype=torch.long) for k, v in fills.items()}
	batch["attention_mask"] = torch.zeros((len(rows), T), dtype=torch.long)
	for b, row in enumerate(rows):
	n = len(row[0])
	for key, values in zip(("input_ids", "labels", "unk_labels", "boundary_labels", "ngram_ids"), row):
	batch[key][b, :n] = torch.tensor(values)
	batch["attention_mask"][b, :n] = 1
	return batch


	def attribution_labels(features: List[Dict[str, Any]], date_bins: int = 20,
	zero_date_is_missing: bool = False) -> Dict[str, torch.Tensor]:
	"""date_labels (+ date_labels_mask: 0 = no date) and region_labels
	(-100 = no label). A missing date is None; test_b rows store an all-zero
	vector instead (zero_date_is_missing=True)."""
	out = {}
	if "date_labels" in features[0]:
	dates = [f.get("date_labels") for f in features]
	ok = [d is not None and (sum(d) > 0 or not zero_date_is_missing) for d in dates]
	out["date_labels"] = torch.tensor([d if k else [0.0] * date_bins for d, k in zip(dates, ok)], dtype=torch.float)
	out["date_labels_mask"] = torch.tensor(ok, dtype=torch.float)
	if "region_labels" in features[0]:
	out["region_labels"] = torch.tensor(
	[f["region_labels"] if f["region_labels"] is not None else -100 for f in features], dtype=torch.long
	)
	return out


	def prepare_test_b(dataset: Any, collator: ContinuaCollator) -> Any:
	"""test_b's real lacunae are baked in ([-] in input_ids, true chars in
	labels): convert each row to continua once (a masked space is dropped, a
	masked punctuation mark folded into `·`) and add ngram_ids and
	boundary_labels."""

	def convert(ex: Dict[str, Any]) -> Dict[str, Any]:
	keep = [j for j, t in enumerate(ex["input_ids"]) if t != collator.pad_id]
	toks, labels, boundary = collator.continua([ex["input_ids"][j] for j in keep], [ex["labels"][j] for j in keep])
	return {"input_ids": toks, "labels": labels, "attention_mask": [1] * len(toks),
	"boundary_labels": boundary, "ngram_ids": collator.ngram_ids(toks)}

	return dataset.map(convert)


	class TestBCollator:
	"""Pads prepared test_b rows into a batch (no damage is added)."""

	def __init__(self, pad_id: int, date_bins: int = 20):
	self.pad_id, self.date_bins = pad_id, date_bins

	def __call__(self, features: List[Dict[str, Any]]) -> Dict[str, torch.Tensor]:
	rows = [(f["input_ids"], f["labels"], f.get("unk_labels") or [-100] * len(f["input_ids"]),
	f["boundary_labels"], f["ngram_ids"]) for f in features]
	batch = pad_batch(rows, self.pad_id)
	batch.update(attribution_labels(features, self.date_bins, zero_date_is_missing=True))
	return batch