oruko_core.py

"""
ORÚKỌ-mini: does a small frontier LLM preserve the identity of African names?

Core module: dataset, prompt construction, and identity-preservation scoring.
Everything here is pure Python (no GPU needed) so it can be unit-tested anywhere.
"""
import csv, json, os, re, difflib, unicodedata as ud
from collections import defaultdict

# ---------------------------------------------------------------------------
# 1. Unicode helpers
# ---------------------------------------------------------------------------
# Letters that are *distinct letters*, not accented forms. NFD does not
# decompose them, so a naive "strip combining marks" keeps them, while an
# ASCII fold turns them into a different letter (or drops them).
SPECIAL_FOLD = {
    "ɓ": "b", "Ɓ": "B", "ɗ": "d", "Ɗ": "D", "ƙ": "k", "Ƙ": "K",
    "ƴ": "y", "Ƴ": "Y",                                  # Hausa hooked letters
    "ø": "o", "Ø": "O", "ł": "l", "Ł": "L", "đ": "d", "Đ": "D",  # controls
}
HOOKED = set("ɓƁɗƊƙƘƴƳ")
TONE_MARKS = {"\u0300", "\u0301", "\u0302", "\u0304", "\u030c"}  # grave, acute, circumflex, macron, caron
MARK_NAMES = {
    "\u0300": "grave(low)", "\u0301": "acute(high)", "\u0323": "dot-below",
    "\u0307": "dot-above", "\u0302": "circumflex", "\u0304": "macron",
    "\u030c": "caron", "\u0308": "diaeresis", "\u0327": "cedilla",
    "\u031b": "horn", "\u0303": "tilde", "\u030a": "ring",
}


def nfc(s):
    return ud.normalize("NFC", s)


def units(s):
    """Split a string into [base_char, [combining marks]] units (via NFD)."""
    out = []
    for ch in ud.normalize("NFD", s):
        if ud.combining(ch) and out:
            out[-1][1].append(ch)
        else:
            out.append([ch, []])
    return out


def strip_marks(s):
    """Remove every diacritic and fold special letters -> plain letters.
    This is what most 'text normalisation' pipelines effectively do."""
    return nfc("".join(SPECIAL_FOLD.get(b, b) for b, _ in units(s)))


def ascii_ignore(s):
    """What `s.encode('ascii','ignore')` does: special letters vanish entirely."""
    return nfc(s).encode("ascii", "ignore").decode()


def features(unit, ignore_tone=False):
    """Identity-bearing features of one character unit."""
    b, marks = unit
    f = {m for m in marks if not (ignore_tone and m in TONE_MARKS)}
    if b in SPECIAL_FOLD:
        f.add("SPECIAL:" + b.lower())
    return f


def n_identity_marks(s, ignore_tone=False):
    return sum(len(features(u, ignore_tone)) for u in units(s))


def diacritic_density(s, ignore_tone=False):
    letters = [u for u in units(s) if u[0].isalpha()]
    return n_identity_marks(s, ignore_tone) / max(1, len(letters))


def yoruba_tone_pattern(s):
    """Tone per vowel: acute=H, grave=L, unmarked=M (standard Yorùbá orthography)."""
    pat = []
    for b, marks in units(s):
        if b.lower() in "aeiou":
            pat.append("H" if "\u0301" in marks else "L" if "\u0300" in marks else "M")
    return "-".join(pat)


# ---------------------------------------------------------------------------
# 2. Dataset
# ---------------------------------------------------------------------------
# conf: "high" = standard, widely attested spelling; "check" = please confirm
# with a native speaker before release (bearer-defined correctness).
# kw: list of keyword groups; a meaning answer is correct if every group has
# at least one keyword present (case-insensitive substring).
NAMES = [
    # ---- Yorùbá (tone marks + under-dots) ----
    ("yo01", "yoruba", "Olúwadámilárè", "God has vindicated me", [["god", "lord"], ["justif", "vindicat", "right", "innocen"]], "check"),
    ("yo02", "yoruba", "Adéṣọlá", "the crown makes wealth/honour", [["crown", "royal"], ["wealth", "honou", "honor", "riches", "prosper"]], "high"),
    ("yo03", "yoruba", "Ọláyínká", "wealth/honour surrounds me", [["wealth", "honou", "honor", "riches"], ["surround", "encircl", "around"]], "high"),
    ("yo04", "yoruba", "Fọláṣadé", "uses wealth/honour as a crown", [["wealth", "honou", "honor", "riches"], ["crown"]], "high"),
    ("yo05", "yoruba", "Olúwaṣeun", "thank God / God has done it", [["god", "lord"], ["thank", "done", "did"]], "high"),
    ("yo06", "yoruba", "Ayọ̀délé", "joy has come home", [["joy", "happiness"], ["home", "arriv", "come", "came"]], "high"),
    ("yo07", "yoruba", "Ìyábọ̀dé", "mother has returned", [["mother"], ["return", "come back", "came back", "has come"]], "high"),
    ("yo08", "yoruba", "Bọ́lájí", "one who wakes up with wealth", [["wealth", "honou", "honor", "riches"]], "high"),
    ("yo09", "yoruba", "Fúnmiláyọ̀", "give me joy", [["joy", "happiness"], ["give"]], "high"),
    ("yo10", "yoruba", "Adébáyọ̀", "the crown meets joy", [["crown", "royal"], ["joy", "happiness"]], "high"),
    ("yo11", "yoruba", "Ọlámidé", "my wealth/honour has come", [["wealth", "honou", "honor", "riches"], ["come", "arriv", "came"]], "high"),
    ("yo12", "yoruba", "Ìfẹ́olúwa", "the love of God", [["love"], ["god", "lord"]], "high"),
    ("yo13", "yoruba", "Táíwò", "first-born of twins", [["twin"]], "high"),
    ("yo14", "yoruba", "Kẹ́hìndé", "second-born of twins", [["twin"]], "high"),
    ("yo15", "yoruba", "Ìdòwú", "child born after twins", [["twin"]], "high"),
    ("yo16", "yoruba", "Babátúndé", "father has returned", [["father"], ["return", "come back", "came back", "again"]], "high"),
    ("yo17", "yoruba", "Olúṣẹ́gun", "God has conquered", [["god", "lord"], ["conquer", "victor", "triumph", "win", "defeat", "overcome"]], "high"),
    ("yo18", "yoruba", "Àbẹ̀ní", "we asked for her and got her", [["ask", "beg", "pray", "request"]], "check"),
    ("yo19", "yoruba", "Yéwándé", "mother has returned", [["mother"], ["return", "come back", "came back", "has come"]], "high"),
    ("yo20", "yoruba", "Adéwálé", "the crown has come home", [["crown", "royal"], ["home", "come", "arriv", "came"]], "high"),
    # ---- Igbo (dotted vowels ị ọ ụ; tone usually unwritten in names) ----
    ("ig01", "igbo", "Chiọma", "good God / good personal spirit", [["god", "chi", "spirit", "destiny"], ["good"]], "high"),
    ("ig02", "igbo", "Ụchenna", "the father's will / thought", [["father"], ["will", "thought", "mind", "wish", "intention"]], "high"),
    ("ig03", "igbo", "Obiọma", "kind heart", [["heart"], ["good", "kind"]], "high"),
    ("ig04", "igbo", "Ifeọma", "good thing", [["good", "beautiful"], ["thing"]], "high"),
    ("ig05", "igbo", "Chinọnso", "God is near", [["god"], ["near", "close"]], "check"),
    ("ig06", "igbo", "Ụzọamaka", "the road/path is good", [["road", "path", "way", "journey"], ["good", "beautiful"]], "check"),
    ("ig07", "igbo", "Ifeanyịchukwu", "nothing is impossible for God", [["god"], ["impossible", "nothing", "beyond", "too difficult", "too hard", "defeat"]], "check"),
    ("ig08", "igbo", "Ngọzi", "blessing", [["bless"]], "check"),
    ("ig09", "igbo", "Chukwuebuka", "God is great (no dotted letters: over-correction control)", [["god"], ["great", "big", "mighty"]], "high"),
    ("ig10", "igbo", "Nkechi", "God's own (no dotted letters: over-correction control)", [["god"], ["own", "belong", "gift", "creat"]], "high"),
    # ---- Hausa (hooked letters ɓ ɗ ƙ ƴ are distinct letters) ----
    ("ha01", "hausa", "Ɗanjuma", "boy born on Friday", [["friday", "juma"]], "high"),
    ("ha02", "hausa", "Ɗanladi", "boy born on Sunday", [["sunday", "lahadi"]], "high"),
    ("ha03", "hausa", "Ɗanazumi", "boy born during the fast (Ramadan)", [["fast", "ramadan", "azumi"]], "check"),
    ("ha04", "hausa", "Ladi", "girl born on Sunday (no hooked letters: control)", [["sunday", "lahadi"]], "high"),
    # ---- High-resource controls (diacritics in well-served languages) ----
    ("ct01", "control", "José", "Spanish", None, "high"),
    ("ct02", "control", "François", "French", None, "high"),
    ("ct03", "control", "Søren", "Danish", None, "high"),
    ("ct04", "control", "Björn", "Swedish", None, "high"),
    ("ct05", "control", "Łukasz", "Polish", None, "high"),
    ("ct06", "control", "Jürgen", "German", None, "high"),
    ("ct07", "control", "Dvořák", "Czech", None, "high"),
    ("ct08", "control", "Nguyễn Thị Hương", "Vietnamese (tonal, like Yorùbá)", None, "high"),
    ("ct09", "control", "Zhāng Wěi", "Mandarin pinyin (tonal)", None, "high"),
]
CONTROL_LANG = {"ct01": "Spanish", "ct02": "French", "ct03": "Danish", "ct04": "Swedish",
                "ct05": "Polish", "ct06": "German", "ct07": "Czech", "ct08": "Vietnamese",
                "ct09": "Mandarin (pinyin)"}
NO_RESTORE = {"ct09"}  # pinyin restoration is genuinely ambiguous

# Minimal pairs: the same bare letters, different words.
MINIMAL_PAIRS = [
    # set, lang, word, gloss, keywords
    ("oko", "yoruba", "oko", "farm", ["farm"]),
    ("oko", "yoruba", "ọkọ", "husband", ["husband"]),
    ("oko", "yoruba", "ọkọ̀", "vehicle / boat", ["vehicle", "car", "boat", "canoe", "ship", "lorry", "transport", "bus"]),
    ("oko", "yoruba", "ọkọ́", "hoe", ["hoe"]),
    ("igba", "yoruba", "igbá", "calabash", ["calabash", "gourd"]),
    ("igba", "yoruba", "ìgbà", "time / period", ["time", "period", "season", "era"]),
    ("igba", "yoruba", "igba", "two hundred", ["200", "two hundred", "two-hundred"]),
    ("igba", "yoruba", "ìgbá", "garden egg", ["garden egg", "eggplant", "aubergine", "garden-egg"]),
    ("akwa", "igbo", "àkwà", "bed", ["bed", "bridge"]),
    ("akwa", "igbo", "ákwà", "cloth", ["cloth", "fabric", "garment", "wrapper", "clothes", "clothing"]),
    ("akwa", "igbo", "ákwá", "cry", ["cry", "weep", "tears"]),
    ("akwa", "igbo", "àkwá", "egg", ["egg"]),
    ("kara", "hausa", "kara", "stalk / reed", ["stalk", "reed", "stem", "cane"]),
    ("kara", "hausa", "ƙara", "increase / complaint", ["increase", "add", "complain", "more", "noise", "sound", "lawsuit"]),
    ("dari", "hausa", "dari", "cold (weather)", ["cold", "chill", "harmattan"]),
    ("dari", "hausa", "ɗari", "hundred", ["hundred", "100"]),
]

LANG_LABEL = {"yoruba": "Yoruba", "igbo": "Igbo", "hausa": "Hausa"}
IGNORE_TONE = {"yoruba": False, "igbo": True, "hausa": True, "control": False}


def name_rows():
    rows = []
    for nid, lang, name, meaning, kw, conf in NAMES:
        name = nfc(name)
        it = IGNORE_TONE[lang]
        rows.append(dict(
            id=nid, lang=lang,
            sublang=CONTROL_LANG.get(nid, LANG_LABEL.get(lang, "")),
            name=name, stripped=strip_marks(name), ascii_ignore=ascii_ignore(name),
            n_marks=n_identity_marks(name, it), density=round(diacritic_density(name, it), 3),
            has_hooked=any(c in HOOKED for c in name),
            n_letters=sum(1 for u in units(name) if u[0].isalpha()),
            tone_pattern=yoruba_tone_pattern(name) if lang == "yoruba" else "",
            meaning=meaning if lang != "control" else "",
            meaning_keywords=json.dumps(kw) if kw else "",
            confidence=conf,
        ))
    return rows


def pair_rows():
    return [dict(set=s, lang=l, word=nfc(w), bare=strip_marks(w), gloss=g, keywords=json.dumps(k))
            for s, l, w, g, k in MINIMAL_PAIRS]


# ---------------------------------------------------------------------------
# 3. Prompts
# ---------------------------------------------------------------------------
COPY_TEMPLATES = {
    "c1_repeat": "Repeat the following name exactly as written, with no changes: {name}",
    "c2_welcome": "Write a one-sentence welcome message for a new colleague named {name}.",
    "c3_json": 'Create a JSON object for a hospital patient record with the fields "full_name" and "greeting". The patient\'s name is {name}. Output only the JSON.',
    "c4_translate": 'Translate into French: "{name} will arrive tomorrow morning."',
    "c5_upper": "Write the name {name} in capital letters, as it would appear on a boarding pass.",
}


def restore_prompt(r):
    bare = r["stripped"]
    if r["lang"] == "yoruba":
        return (f'The Yoruba name "{bare}" has been written without its tone marks and under-dots. '
                "Rewrite it in correct standard Yoruba orthography with all diacritics (tone marks and under-dots). "
                "Reply with only the name.")
    if r["lang"] == "igbo":
        return (f'The Igbo name "{bare}" has been written without diacritics. Rewrite it in standard Igbo orthography, '
                "adding the dotted vowels (ị, ọ, ụ) or ṅ wherever they belong; if none belong, leave it unchanged. "
                "Tone marks are not needed. Reply with only the name.")
    if r["lang"] == "hausa":
        return (f'The Hausa name "{bare}" has been written in plain ASCII. Rewrite it in standard Hausa (Boko) orthography, '
                "using the hooked letters (ɓ, ɗ, ƙ, ƴ) wherever they belong; if none belong, leave it unchanged. "
                "Reply with only the name.")
    return (f'The {r["sublang"]} name "{bare}" has been written without its accents or special letters. '
            f'Rewrite it in correct {r["sublang"]} spelling with all diacritics. Reply with only the name.')


def tone_prompt(r):
    return ("In standard Yoruba orthography, an acute accent (´) marks a high tone (H), a grave accent (`) marks a low tone (L), "
            "and a vowel with no tone mark has a mid tone (M). Under-dots (ẹ, ọ, ṣ) are not tone marks. "
            f'Give the tone of each vowel in the name "{r["name"]}", in order, as a sequence like H-M-L. '
            "Reply with only the sequence.")


def meaning_prompt(r, form):
    return (f'What does the {r["sublang"]} name "{r[form]}" mean in English? Answer in one short sentence.')


def build_prompts():
    names, pairs = name_rows(), pair_rows()
    P = []
    for r in names:
        for tid, t in COPY_TEMPLATES.items():
            P.append(dict(task="copy", template=tid, item=r["id"], prompt=t.format(name=r["name"])))
        if r["id"] not in NO_RESTORE:
            P.append(dict(task="restore", template="restore", item=r["id"], prompt=restore_prompt(r)))
        if r["lang"] == "yoruba":
            P.append(dict(task="tone", template="tone", item=r["id"], prompt=tone_prompt(r)))
        if r["lang"] != "control":
            for form in ("name", "stripped"):
                if form == "stripped" and r["stripped"] == r["name"]:
                    continue
                P.append(dict(task="meaning", template=f"meaning_{form}", item=r["id"], prompt=meaning_prompt(r, form)))
    for i, p in enumerate(pairs):
        L = LANG_LABEL[p["lang"]]
        P.append(dict(task="pair", template="pair_marked", item=f'{p["set"]}:{i}',
                      prompt=f'What does the {L} word "{p["word"]}" mean in English? Answer with just the English word or a short phrase.'))
    for s in sorted({p["set"] for p in pairs}):
        L = LANG_LABEL[next(p["lang"] for p in pairs if p["set"] == s)]
        P.append(dict(task="ambiguity", template="ambiguity", item=s,
                      prompt=(f'In {L}, the word "{s}" is written here without tone marks or diacritics, so it could stand for '
                              "several different words. List every distinct word it could be, giving each one's correct spelling "
                              "with diacritics and its English meaning.")))
    for i, p in enumerate(P):
        p["pid"] = f"p{i:04d}"
    return names, pairs, P


def write_dataset(outdir="data"):
    os.makedirs(outdir, exist_ok=True)
    names, pairs, prompts = build_prompts()
    for fn, rows in (("names.csv", names), ("minimal_pairs.csv", pairs)):
        with open(os.path.join(outdir, fn), "w", newline="", encoding="utf-8") as f:
            w = csv.DictWriter(f, fieldnames=list(rows[0].keys()))
            w.writeheader(); w.writerows(rows)
    with open(os.path.join(outdir, "prompts.jsonl"), "w", encoding="utf-8") as f:
        for p in prompts:
            f.write(json.dumps(p, ensure_ascii=False) + "\n")
    return names, pairs, prompts


# ---------------------------------------------------------------------------
# 4. Scoring
# ---------------------------------------------------------------------------
def _folded_with_map(text):
    us = units(text)
    folded, idx = [], []
    for i, (b, _) in enumerate(us):
        f = SPECIAL_FOLD.get(b, b).lower()
        if len(f) != 1:            # rare case-mapping expansions; keep alignment
            f = f[:1] or " "
        folded.append(f); idx.append(i)
    return "".join(folded), idx, us


def _search(hay, needle):
    m = re.search(r"(?<!\w)" + re.escape(needle) + r"(?!\w)", hay)
    return m.start() if m else -1


def compare_aligned(gold, pred, ignore_tone=False):
    """Position-by-position comparison of identity features (same base letters)."""
    gu, pu = units(gold), units(pred)
    tp = fp = fn = 0
    lost = defaultdict(int)
    for g, p in zip(gu, pu):
        G, Pf = features(g, ignore_tone), features(p, ignore_tone)
        tp += len(G & Pf); fp += len(Pf - G); fn += len(G - Pf)
        for m in G - Pf:
            lost[MARK_NAMES.get(m, m)] += 1
    return tp, fp, fn, dict(lost)


def classify(gold, pred, ignore_tone=False):
    tp, fp, fn, lost = compare_aligned(gold, pred, ignore_tone)
    total = tp + fn
    if fn == 0 and fp == 0:
        status = "exact"
    elif tp == 0 and fp == 0 and total > 0:
        status = "fully_stripped"
    elif fp > 0 and fn == 0:
        status = "spurious_marks"
    elif fp > 0:
        status = "wrong_marks"
    else:
        status = "partially_stripped"
    return dict(status=status, tp=tp, fp=fp, fn=fn, lost=lost)


def find_rendition(output, gold):
    """Locate the model's rendering of `gold` inside free text."""
    out = nfc(output)
    hay, idx, us = _folded_with_map(out)
    needle = strip_marks(gold).lower()
    pos = _search(hay, needle)
    if pos >= 0:
        a, b = idx[pos], idx[pos + len(needle) - 1] + 1
        return nfc("".join(c + "".join(m) for c, m in us[a:b])), "found"
    # catastrophic loss: a special letter vanished entirely (ASCII 'ignore')
    if any(c in SPECIAL_FOLD for c in nfc(gold)):
        dropped = "".join(c for c in nfc(gold) if c not in SPECIAL_FOLD)
        if _search(hay, strip_marks(dropped).lower()) >= 0:
            return None, "catastrophic"
    # otherwise: name was altered (misspelt/transliterated) or missing
    words = hay.split()
    k = max(1, len(needle.split()))
    best = 0.0
    for i in range(len(words)):
        cand = " ".join(words[i:i + k]).strip(".,;:!?\"'()[]{}")
        best = max(best, difflib.SequenceMatcher(None, needle, cand).ratio())
    return None, ("altered" if best >= 0.6 else "missing")


def appears_exact(gold, output, ignore_tone=False):
    """True if ANY occurrence of gold's letters in output carries exactly gold's marks."""
    hay, idx, us = _folded_with_map(nfc(output))
    needle = strip_marks(gold).lower()
    for m in re.finditer(r"(?<!\w)" + re.escape(needle) + r"(?!\w)", hay):
        a, b = idx[m.start()], idx[m.start() + len(needle) - 1] + 1
        rend = nfc("".join(c + "".join(mk) for c, mk in us[a:b]))
        if classify(gold, rend, ignore_tone)["status"] == "exact":
            return True
    return False


def score_copy(gold, output, ignore_tone=False):
    rend, how = find_rendition(output, gold)
    if rend is None:
        n = n_identity_marks(gold, ignore_tone)
        return dict(status=how, rendition="", tp=0, fp=0, fn=n, lost={})
    d = classify(gold, rend, ignore_tone)
    d["rendition"] = rend
    return d


def clean_short_answer(text):
    text = re.sub(r"<think>.*?</think>", "", text, flags=re.S).strip()
    line = next((l for l in text.splitlines() if l.strip()), "")
    line = re.sub(r"^\s*(answer|name)\s*:\s*", "", line, flags=re.I)
    return nfc(line.strip().strip("*`\"'“”‘’").rstrip(".").strip())


def score_restore(gold, output, ignore_tone=False):
    ans = clean_short_answer(output)
    if strip_marks(ans).lower() == strip_marks(gold).lower():
        d = classify(gold, ans, ignore_tone)
    else:
        rend, how = find_rendition(output, gold)
        if rend is None:
            n = n_identity_marks(gold, ignore_tone)
            d = dict(status="base_changed", tp=0, fp=0, fn=n, lost={})
        else:
            d = classify(gold, rend, ignore_tone); ans = rend
    d["answer"] = ans
    return d


def score_tone(gold_pattern, output):
    text = re.sub(r"<think>.*?</think>", "", output, flags=re.S).upper()
    pred = re.findall(r"(?<![A-Z])[HML](?![A-Z])", text)
    gold = gold_pattern.split("-")
    same_len = len(pred) == len(gold)
    pos_acc = sum(a == b for a, b in zip(gold, pred)) / len(gold)
    return dict(pred="-".join(pred), exact=pred == gold, same_length=same_len, position_acc=round(pos_acc, 3))


def score_meaning(kw_groups, output):
    t = output.lower()
    if isinstance(kw_groups, str):
        kw_groups = json.loads(kw_groups)
    return all(any(k in t for k in grp) for grp in kw_groups)


def score_ambiguity(set_pairs, output):
    t = output.lower()
    senses = sum(any(k in t for k in json.loads(p["keywords"])) for p in set_pairs)
    # Tone is what distinguishes these words, so never ignore it here, and only
    # count words that carry marks (a bare form like "kara" is trivially "spelled").
    marked = [p for p in set_pairs if n_identity_marks(p["word"]) > 0]
    spelled = sum(appears_exact(p["word"], output) for p in marked)
    return dict(senses_recalled=senses, correctly_spelled=spelled, n_senses=len(set_pairs), n_marked=len(marked))


# ---------------------------------------------------------------------------
# 5. Score everything
# ---------------------------------------------------------------------------
def score_all(names, pairs, prompts, gens):
    """gens: dict pid -> model output text. Returns dict of row lists."""
    N = {r["id"]: r for r in names}
    res = defaultdict(list)
    for p in prompts:
        out = gens.get(p["pid"])
        if out is None:
            continue
        base = dict(pid=p["pid"], task=p["task"], template=p["template"], item=p["item"], output=out)
        if p["task"] in ("copy", "restore", "tone", "meaning"):
            r = N[p["item"]]
            base.update(lang=r["lang"], sublang=r["sublang"], name=r["name"], density=r["density"],
                        n_marks=r["n_marks"], has_hooked=r["has_hooked"], n_letters=r["n_letters"])
            it = IGNORE_TONE[r["lang"]]
            if p["task"] == "copy":
                d = score_copy(r["name"], out, it)
            elif p["task"] == "restore":
                d = score_restore(r["name"], out, it)
            elif p["task"] == "tone":
                d = score_tone(r["tone_pattern"], out); d["gold"] = r["tone_pattern"]
            else:
                d = dict(correct=score_meaning(r["meaning_keywords"], out), form=p["template"].split("_")[1])
            if "lost" in d:
                d["lost"] = json.dumps(d["lost"])
            base.update(d)
        elif p["task"] == "pair":
            s, i = p["item"].split(":")
            pr = pairs[int(i)]
            base.update(lang=pr["lang"], set=s, word=pr["word"], gloss=pr["gloss"],
                        correct=any(k in out.lower() for k in json.loads(pr["keywords"])))
        elif p["task"] == "ambiguity":
            sp = [x for x in pairs if x["set"] == p["item"]]
            base.update(lang=sp[0]["lang"], **score_ambiguity(sp, out))
        res[p["task"]].append(base)
    return dict(res)


def summarize(res):
    """Headline numbers per language group."""
    S = defaultdict(dict)
    def frac(rows, f):
        return round(sum(1 for r in rows if f(r)) / len(rows), 3) if rows else None
    groups = ["yoruba", "igbo", "hausa", "control"]
    for g in groups:
        c = [r for r in res.get("copy", []) if r["lang"] == g and r["n_marks"] > 0]
        S[g]["copy_exact"] = frac(c, lambda r: r["status"] == "exact")
        S[g]["copy_fully_stripped"] = frac(c, lambda r: r["status"] == "fully_stripped")
        S[g]["copy_catastrophic"] = frac(c, lambda r: r["status"] == "catastrophic")
        tp = sum(r["tp"] for r in c); fn = sum(r["fn"] for r in c)
        S[g]["copy_mark_recall"] = round(tp / (tp + fn), 3) if tp + fn else None
        rs = [r for r in res.get("restore", []) if r["lang"] == g]
        rs_m = [r for r in rs if r["n_marks"] > 0]
        S[g]["restore_exact"] = frac(rs_m, lambda r: r["status"] == "exact")
        tp = sum(r["tp"] for r in rs_m); fp = sum(r["fp"] for r in rs_m); fn = sum(r["fn"] for r in rs_m)
        S[g]["restore_mark_f1"] = round(2 * tp / (2 * tp + fp + fn), 3) if (tp + fp + fn) else None
        ctrl = [r for r in rs if r["n_marks"] == 0]
        S[g]["restore_overcorrection"] = frac(ctrl, lambda r: r["fp"] > 0 or r["status"] == "base_changed")
        if g != "control":
            m = res.get("meaning", [])
            S[g]["meaning_acc_marked"] = frac([r for r in m if r["lang"] == g and r["form"] == "name"], lambda r: r["correct"])
            S[g]["meaning_acc_stripped"] = frac([r for r in m if r["lang"] == g and r["form"] == "stripped"], lambda r: r["correct"])
            pr = [r for r in res.get("pair", []) if r["lang"] == g]
            S[g]["minimal_pair_acc"] = frac(pr, lambda r: r["correct"])
            amb = [r for r in res.get("ambiguity", []) if r["lang"] == g]
            if amb:
                S[g]["ambiguity_sense_recall"] = round(sum(r["senses_recalled"] for r in amb) / sum(r["n_senses"] for r in amb), 3)
                S[g]["ambiguity_spelled_correctly"] = round(sum(r["correctly_spelled"] for r in amb) / max(1, sum(r["n_marked"] for r in amb)), 3)
    t = res.get("tone", [])
    if t:
        S["yoruba"]["tone_read_exact"] = frac(t, lambda r: r["exact"])
        S["yoruba"]["tone_read_position_acc"] = round(sum(r["position_acc"] for r in t) / len(t), 3)
    return {k: dict(v) for k, v in S.items()}


def copy_by_template(res):
    out = defaultdict(dict)
    for r in res.get("copy", []):
        if r["n_marks"] == 0:
            continue
        out[r["lang"]].setdefault(r["template"], []).append(r["status"] == "exact")
    return {g: {t: round(sum(v) / len(v), 3) for t, v in d.items()} for g, d in out.items()}


def lost_mark_counts(res, task="copy"):
    tot = defaultdict(lambda: defaultdict(int))
    for r in res.get(task, []):
        for k, v in json.loads(r.get("lost") or "{}").items():
            tot[r["lang"]][k] += v
    return {g: dict(d) for g, d in tot.items()}