Diacritics: Automatic Restoration of Diacritics in Slovak Text

"""Slovak OAI 2026 home round, task 3: restoring Slovak diacritics with SlovakBERT.

usage: python solution.py prep        # part 1: skwiki dump -> sentence pairs, splits, stats.json
       python solution.py zeroshot    # part 2: candidate scoring with the frozen masked LM
       python solution.py finetune    # part 3: SlovakBERT + per-character diacritic head
Inputs: skwiki_part.xml.bz2 (the first 120 MB of skwiki-latest-pages-articles.xml.bz2; the
20251120 dump named in the statement is no longer online), gerulata/slovakbert in ./sbert.
"""
import bz2, itertools, json, random, re, sys, time, unicodedata, zlib
import numpy as np, torch, torch.nn as nn, torch.nn.functional as F
torch.set_num_threads(2)
DEV = torch.device("mps" if torch.backends.mps.is_available() else "cpu")

# ---- the statement's map, and its inverse as (base letter, mark) classes -----------------------
DIA = "áäčďéíĺľňóôŕšťúýžÁÄČĎÉÍĹĽŇÓÔŔŠŤÚÝŽ"
BASE = "aacdeillnoorstuyzAACDEILLNOORSTUYZ"
STRIP = str.maketrans(DIA, BASE)
strip = lambda s: s.translate(STRIP)
MARK_OF = {}  # diacritic char -> mark id; 1 acute, 2 caron/hacek, 3 umlaut, 4 circumflex
for ch in DIA:
    name = unicodedata.name(ch)
    MARK_OF[ch] = 1 if "ACUTE" in name else 2 if ("CARON" in name or ch in "ľĽďťĎŤ") else 3 if "DIAERESIS" in name else 4
APPLY = {(b, MARK_OF[d]): d for d, b in zip(DIA, BASE)}           # (base, mark) -> diacritic char
CAND = set(BASE)                                                    # positions that are scored
VARIANTS = {b: [b] + [APPLY[(b, m)] for m in range(1, 5) if (b, m) in APPLY] for b in set(BASE)}

def char_accuracy(preds, refs):
    """The official metric: accuracy on characters that could carry a diacritic."""
    ok = n = 0
    for p, r in zip(preds, refs):
        for a, b in zip(p, r):
            if strip(b) in CAND:
                n += 1; ok += a == b
    return ok / n

# ---------------------------------- part 1: data ------------------------------------------------
def prep(max_articles=40000):
    import mwparserfromhell
    buf, pages = "", []
    dec = bz2.BZ2Decompressor()
    with open("skwiki_part.xml.bz2", "rb") as f:        # truncated archive: decompress until it breaks
        try:
            while len(pages) < max_articles:
                chunk = f.read(1 << 20)
                if not chunk: break
                buf += dec.decompress(chunk).decode("utf-8", "ignore")
                parts = buf.split("</page>"); buf = parts.pop()
                for p in parts:
                    if "<ns>0</ns>" in p and "<redirect" not in p:
                        m = re.search(r"<text[^>]*>(.*?)</text>", p, re.S)
                        if m: pages.append(m.group(1))
        except (EOFError, OSError):
            pass
    import html
    sents, seen = [], set()
    SPLIT = re.compile(r"(?<=[.!?])\s+(?=[A-ZÁÄČĎÉÍĹĽŇÓÔŔŠŤÚÝŽ])")
    for art_id, wikitext in enumerate(pages):
        code = mwparserfromhell.parse(html.unescape(wikitext))
        for t in code.filter_templates(recursive=False) + [t for t in code.filter_tags(recursive=False)
                if str(t.tag) in ("ref", "table", "math", "gallery", "code")]:
            try: code.remove(t)                                       # infoboxes, citations, tables, formulas
            except ValueError: pass                                   # already removed with its parent
        text = code.strip_code(normalize=True, collapse=True)
        for line in text.split("\n"):
            line = line.strip()
            if len(line) < 50 or line.startswith(("Kategória:", "Súbor:", "|", "*", "#", "!", "{")): continue
            for s in SPLIT.split(line):
                s = re.sub(r"\s+", " ", s).strip()
                letters = sum(c.isalpha() for c in s)
                # keep real sentences: 50-2000 chars, mostly letters, ends like a sentence, no leftover markup
                if 50 <= len(s) <= 2000 and letters > 0.6 * len(s) and s[-1] in ".!?" and not re.search(r"[\[\]{}<>=|]", s) and s not in seen:
                    seen.add(s); sents.append((art_id, s))
    split_of = lambda a: "test" if zlib.crc32(str(a).encode()) % 20 == 0 else "val" if zlib.crc32(str(a).encode()) % 20 == 1 else "train"
    data = {"train": [], "val": [], "test": []}
    for a, s in sents: data[split_of(a)].append({"input": strip(s), "target": s})
    json.dump(data, open("pairs.json", "w"), ensure_ascii=False)
    allt = [p["target"] for v in data.values() for p in v]
    stats = {"articles_used": len(pages), "total_sentence_pairs": len(allt), "train_pairs": len(data["train"]),
             "val_pairs": len(data["val"]), "test_pairs": len(data["test"]),
             "avg_sentence_length_chars": round(float(np.mean([len(s) for s in allt])), 1),
             "total_characters": int(sum(len(s) for s in allt)),
             "diacritic_characters_count": int(sum(c in MARK_OF for s in allt for c in s)),
             "unique_words": len({w for s in allt for w in re.findall(r"\w+", s.lower())})}
    json.dump(stats, open("stats.json", "w"), indent=2); print(stats)

def load_pairs(split, n=None, seed=0):
    d = json.load(open("pairs.json"))[split]
    random.Random(seed).shuffle(d)
    return d[:n] if n else d

# ---------------------------------- part 2: zero-shot -------------------------------------------
def build_piece_index(tok):
    """Map each STRIPPED vocabulary piece to the real pieces it can come from ('dnes' -> 'dnes', 'dneš'...)."""
    first, cont = {}, {}
    for t, i in tok.get_vocab().items():
        sdec = tok.convert_tokens_to_string([t])
        if not sdec.strip() or any(ch in sdec for ch in "<>"): continue
        if sdec.startswith(" "): first.setdefault(strip(sdec[1:]), []).append((sdec[1:], i))
        else: cont.setdefault(strip(sdec), []).append((sdec, i))
    return first, cont

def candidates(w, first, cont, slack=1, cap=300):
    """All spellings of w that SlovakBERT's vocabulary writes with at most (fewest pieces + slack) pieces,
    as (text, piece ids). Dynamic programming over the split points of the word."""
    n = len(w); best = {0: [("", [])]}
    for i in range(n):
        if i not in best: continue
        for j in range(i + 1, n + 1):
            for txt, tid in (first if i == 0 else cont).get(w[i:j], []):
                for pre, ids in best[i]:
                    best.setdefault(j, []).append((pre + txt, ids + [tid]))
        if i in best and i > 0: best[i] = best[i][:cap]
    if n not in best: return [(w, None)]
    kmin = min(len(ids) for _, ids in best[n])
    seen, out = set(), []
    for txt, ids in sorted(best[n], key=lambda x: len(x[1])):
        if len(ids) <= kmin + slack and txt not in seen:
            seen.add(txt); out.append((txt, ids))
    return out[:cap]

def zeroshot(n_test=150):
    from transformers import AutoTokenizer, AutoModelForMaskedLM
    tok = AutoTokenizer.from_pretrained("sbert"); mlm = AutoModelForMaskedLM.from_pretrained("sbert").eval().to(DEV)
    first, cont = build_piece_index(tok); cache = {}
    test = load_pairs("test", n_test)
    t0, preds = time.time(), []
    for n_done, ex in enumerate(test):
        if n_done % 25 == 0: print(f"  {n_done} sentences, {time.time() - t0:.0f}s", flush=True)
        words = re.findall(r"\w+|\W+", ex["input"])
        out = list(words)
        for i, w in enumerate(words):
            if not w[0].isalpha(): continue
            if w not in cache: cache[w] = candidates(w, first, cont)
            cands = cache[w]
            if len(cands) == 1:                      # one spelling only: the vocabulary decides alone
                out[i] = cands[0][0]; continue
            # Score each spelling with the frozen masked LM: replace the word by as many <mask> as the
            # spelling has BPE pieces and add the log-probabilities of its pieces (one forward pass per
            # piece count). The sum also acts as a lexicon prior: real words need few, likely pieces.
            score = {}
            for k in sorted({len(ids) for _, ids in cands}):
                left = "".join(out[:i]).rstrip(" ")
                ctx = left + "".join([tok.mask_token] * k) + "".join(words[i + 1:])
                enc = tok(ctx, return_tensors="pt", truncation=True, max_length=256).to(DEV)
                pos = (enc.input_ids[0] == tok.mask_token_id).nonzero()[:, 0]
                if len(pos) != k: continue
                with torch.no_grad():
                    lp = mlm(**enc).logits[0, pos].log_softmax(-1).cpu()   # one GPU->CPU copy
                for c, ids in cands:
                    if len(ids) == k:
                        score[c] = max(score.get(c, -1e9), sum(lp[j, t].item() for j, t in enumerate(ids)))
            if score:
                out[i] = max(score, key=score.get)
        preds.append("".join(out))
    refs = [e["target"] for e in test]
    print(f"zero-shot on {n_test} test sentences: no diacritics {char_accuracy([e['input'] for e in test], refs):.4f}, "
          f"SlovakBERT candidates {char_accuracy(preds, refs):.4f}, {time.time() - t0:.0f}s")
    with open("predictions_zeroshot.tsv", "w") as f:
        for k, p in enumerate(preds, 1): f.write(f"{k}\t{p}\n")

# ---------------------------------- part 3: fine-tuning -----------------------------------------
class DiacriticTagger(nn.Module):
    """SlovakBERT encodes the stripped text; each character is classified (none / acute / caron /
    umlaut / circumflex) from the hidden state of its token + the character + its offset in the token."""
    def __init__(self):
        super().__init__()
        from transformers import AutoModel
        self.enc = AutoModel.from_pretrained("sbert")
        self.char = nn.Embedding(128, 64); self.off = nn.Embedding(32, 32)
        self.head = nn.Sequential(nn.Linear(768 + 96, 256), nn.GELU(), nn.Linear(256, 5))
    def forward(self, ids, att, tok_idx, chars, offs):
        h = self.enc(input_ids=ids, attention_mask=att).last_hidden_state           # (B, T, 768)
        ht = torch.gather(h, 1, tok_idx[..., None].expand(-1, -1, 768))             # (B, C, 768)
        return self.head(torch.cat([ht, self.char(chars), self.off(offs.clamp(max=31))], -1))

def encode_batch(tok, inputs, targets=None, max_chars=400):
    enc = tok(inputs, return_offsets_mapping=True, padding=True, truncation=True, max_length=160, return_tensors="pt")
    B, C = len(inputs), min(max_chars, max(len(s) for s in inputs))
    tok_idx = torch.zeros(B, C, dtype=torch.long); chars = torch.zeros(B, C, dtype=torch.long)
    offs = torch.zeros(B, C, dtype=torch.long); y = torch.full((B, C), -100, dtype=torch.long)
    for b, s in enumerate(inputs):
        for t, (a, e) in enumerate(enc.offset_mapping[b].tolist()):
            for c in range(a, min(e, C)):
                tok_idx[b, c], offs[b, c] = t, c - a
                chars[b, c] = min(ord(s[c]), 127)
                if targets is not None and s[c] in CAND:
                    y[b, c] = MARK_OF.get(targets[b][c], 0)
    return enc.input_ids, enc.attention_mask, tok_idx, chars, offs, y

def restore(model, tok, inputs, bs=64):
    out = []
    model.eval()
    with torch.no_grad():
        for k in range(0, len(inputs), bs):
            batch = inputs[k:k + bs]
            ids, att, ti, ch, of, _ = encode_batch(tok, batch)
            logits = model(ids.to(DEV), att.to(DEV), ti.to(DEV), ch.to(DEV), of.to(DEV)).cpu()
            for b, s in enumerate(batch):
                chars = list(s)
                for c in range(min(len(s), logits.shape[1])):
                    if s[c] in CAND:
                        order = logits[b, c].argsort(descending=True).tolist()
                        m = next(m for m in order if m == 0 or (s[c], m) in APPLY)   # best VALID mark for this letter
                        chars[c] = APPLY[(s[c], m)] if m else s[c]
                out.append("".join(chars))
    return out

def finetune(n_train=70000, epochs=1, bs=32, lr=5e-5):
    from transformers import AutoTokenizer
    tok = AutoTokenizer.from_pretrained("sbert")
    train, val, test = load_pairs("train", n_train), load_pairs("val", 2000), load_pairs("test", 3000)
    train = [p for p in train if len(p["input"]) <= 400]
    model = DiacriticTagger().to(DEV)
    opt = torch.optim.AdamW(model.parameters(), lr=lr, weight_decay=0.01)
    steps = epochs * len(train) // bs
    sch = torch.optim.lr_scheduler.LambdaLR(opt, lambda s: min(1, s / (0.05 * steps)) * max(0.0, 1 - s / steps))
    t0 = time.time(); step = 0
    for ep in range(epochs):
        random.Random(ep).shuffle(train)
        for k in range(0, len(train) - bs + 1, bs):
            batch = train[k:k + bs]
            ids, att, ti, ch, of, y = encode_batch(tok, [p["input"] for p in batch], [p["target"] for p in batch])
            model.train()
            logits = model(ids.to(DEV), att.to(DEV), ti.to(DEV), ch.to(DEV), of.to(DEV))
            loss = F.cross_entropy(logits.reshape(-1, 5), y.to(DEV).reshape(-1), ignore_index=-100)
            opt.zero_grad(); loss.backward(); opt.step(); sch.step(); step += 1
            if step % 500 == 0:
                acc = char_accuracy(restore(model, tok, [p["input"] for p in val[:500]]), [p["target"] for p in val[:500]])
                print(f"step {step}/{steps} loss {loss.item():.4f} val char-acc {acc:.4f} {time.time() - t0:.0f}s", flush=True)
    torch.save(model.state_dict(), "tagger.pt")
    # Lexicon-constrained decoding: every word spelling seen in the part-1 training sentences is a
    # candidate (plus the tagger's own spelling); score = sum of the tagger's log-probabilities of the
    # marks + alpha * log relative frequency of that spelling. alpha is chosen on the validation set.
    import collections
    lex = collections.defaultdict(collections.Counter)
    for p in load_pairs("train"):
        for w in re.findall(r"\w+", p["target"]): lex[strip(w)][w] += 1
    def char_logp(inputs, bs=64):
        out = []
        model.eval()
        with torch.no_grad():
            for k in range(0, len(inputs), bs):
                b = inputs[k:k + bs]; ids, att, ti, ch, of, _ = encode_batch(tok, b)
                lp = model(ids.to(DEV), att.to(DEV), ti.to(DEV), ch.to(DEV), of.to(DEV)).log_softmax(-1).cpu().numpy()
                out += [lp[i, :len(s)] for i, s in enumerate(b)]
        return out
    def decode(s, lp, own, alpha):
        res = list(s)
        for m in re.finditer(r"\w+", s):
            w, a = m.group(0), m.start()
            cands, tot = set(lex.get(w, {})) | {own[a:a + len(w)]}, sum(lex[w].values()) if w in lex else 0
            def score(c):
                sc = sum(lp[a + j][MARK_OF.get(c[j], 0)] for j in range(len(w)) if w[j] in CAND and a + j < len(lp))
                return sc + (alpha * np.log((lex[w].get(c, 0) + 0.5) / (tot + 1)) if tot else 0)
            res[a:a + len(w)] = max(cands, key=score)
        return "".join(res)
    results = {}
    for name, d in [("val", val), ("test", test)]:
        inp, refs = [p["input"] for p in d], [p["target"] for p in d]
        own, L = restore(model, tok, inp), char_logp(inp)
        print(f"fine-tuned {name}: char accuracy {char_accuracy(own, refs):.4f}")
        alphas = [0.5, 1, 2, 4] if name == "val" else [best_alpha]
        accs = {al: char_accuracy([decode(s, l, o, al) for s, l, o in zip(inp, L, own)], refs) for al in alphas}
        print(f"  + lexicon, {name}: " + ", ".join(f"alpha {al}: {v:.4f}" for al, v in accs.items()))
        if name == "val": best_alpha = max(accs, key=accs.get)
    pred = [decode(s, l, o, best_alpha) for s, l, o in zip(inp, L, own)]
    with open("predictions_finetuned.tsv", "w") as f:
        for k, p in enumerate(pred, 1): f.write(f"{k}\t{p}\n")

if __name__ == "__main__":
    {"prep": prep, "zeroshot": zeroshot, "finetune": finetune}[sys.argv[1]]()