spb/wpai0001_uqo

1"""Tier-1 tasks for open-weight models: verifiable input→output mappings presented as2``Input: x\\nOutput:`` prompts (few-shot with k demonstrations, or zero-shot).34Families: semantic (antonyms), factual (country → capital), invented morphology (pseudo-word suffix5rule, no prior), formal (add three to a two-digit number). Items are fixed lists (hashed in manifests);6demonstrations and queries are disjoint. Alignment text is template-generated English unrelated to the7tasks. Exact-match scoring on the first generated line, case-insensitive, with listed alternatives.8"""9from __future__ import annotations1011import hashlib12import random13from dataclasses import dataclass1415ANTONYMS = [("hot", "cold"), ("big", "small"), ("fast", "slow"), ("up", "down"), ("day", "night"), ("open", "closed"), ("happy", "sad"),16            ("light", "dark"), ("tall", "short"), ("rich", "poor"), ("young", "old"), ("early", "late"), ("wet", "dry"), ("full", "empty"),17            ("strong", "weak"), ("hard", "soft"), ("loud", "quiet"), ("clean", "dirty"), ("thick", "thin"), ("wide", "narrow"), ("true", "false"),18            ("inside", "outside"), ("push", "pull"), ("buy", "sell"), ("win", "lose"), ("love", "hate"), ("begin", "end"), ("alive", "dead"),19            ("cheap", "expensive"), ("deep", "shallow"), ("easy", "hard"), ("first", "last"), ("front", "back"), ("heavy", "light"), ("high", "low"),20            ("near", "far"), ("north", "south"), ("east", "west"), ("left", "right"), ("more", "less"), ("add", "subtract"), ("always", "never"),21            ("above", "below"), ("accept", "reject"), ("arrive", "depart"), ("asleep", "awake"), ("brave", "cowardly"), ("bright", "dim"),22            ("calm", "stormy"), ("come", "go"), ("correct", "wrong"), ("create", "destroy"), ("cruel", "kind"), ("friend", "enemy"),23            ("give", "take"), ("good", "bad"), ("guilty", "innocent"), ("increase", "decrease"), ("laugh", "cry"), ("major", "minor"),24            ("male", "female"), ("maximum", "minimum"), ("noisy", "quiet"), ("on", "off"), ("over", "under"), ("past", "future"),25            ("polite", "rude"), ("question", "answer"), ("raw", "cooked"), ("remember", "forget"), ("rough", "smooth"), ("safe", "dangerous"),26            ("same", "different"), ("sweet", "sour"), ("teacher", "student"), ("top", "bottom"), ("visible", "invisible"), ("war", "peace"),27            ("wild", "tame"), ("yes", "no")]28ALT_ANTONYMS = {"hard": ["difficult"], "light": ["dark", "heavy"], "cold": ["cool"], "small": ["little", "tiny"], "quiet": ["silent"], "wrong": ["incorrect"],29                "bad": ["evil"], "stormy": ["rough"], "less": ["fewer"], "dim": ["dark", "dull"], "cowardly": ["cowardly", "timid", "fearful"], "kind": ["gentle"]}3031CAPITALS = [("France", "Paris"), ("Germany", "Berlin"), ("Italy", "Rome"), ("Spain", "Madrid"), ("Portugal", "Lisbon"), ("Japan", "Tokyo"),32            ("China", "Beijing"), ("Canada", "Ottawa"), ("Australia", "Canberra"), ("Egypt", "Cairo"), ("Kenya", "Nairobi"), ("Peru", "Lima"),33            ("Chile", "Santiago"), ("Argentina", "Buenos Aires"), ("Brazil", "Brasilia"), ("Mexico", "Mexico City"), ("Russia", "Moscow"),34            ("Turkey", "Ankara"), ("Greece", "Athens"), ("Sweden", "Stockholm"), ("Norway", "Oslo"), ("Finland", "Helsinki"), ("Denmark", "Copenhagen"),35            ("Poland", "Warsaw"), ("Austria", "Vienna"), ("Hungary", "Budapest"), ("Ireland", "Dublin"), ("Iran", "Tehran"), ("Iraq", "Baghdad"),36            ("India", "New Delhi"), ("Thailand", "Bangkok"), ("Vietnam", "Hanoi"), ("Indonesia", "Jakarta"), ("Philippines", "Manila"),37            ("South Korea", "Seoul"), ("Cuba", "Havana"), ("Colombia", "Bogota"), ("Venezuela", "Caracas"), ("Morocco", "Rabat"), ("Nigeria", "Abuja"),38            ("Ethiopia", "Addis Ababa"), ("Ghana", "Accra"), ("Belgium", "Brussels"), ("Netherlands", "Amsterdam"), ("Switzerland", "Bern"),39            ("Czech Republic", "Prague"), ("Romania", "Bucharest"), ("Bulgaria", "Sofia"), ("Ukraine", "Kyiv"), ("Pakistan", "Islamabad"),40            ("Bangladesh", "Dhaka"), ("Nepal", "Kathmandu"), ("Saudi Arabia", "Riyadh"), ("Jordan", "Amman"), ("Lebanon", "Beirut"),41            ("Syria", "Damascus"), ("Afghanistan", "Kabul"), ("New Zealand", "Wellington"), ("Iceland", "Reykjavik"), ("Malaysia", "Kuala Lumpur")]42ALT_CAPITALS = {"Brasilia": ["Brasília"], "Bogota": ["Bogotá"], "Kyiv": ["Kiev"], "Reykjavik": ["Reykjavík"]}4344VOWELS = "aeiou"45CONS = "bdfgklmnprstvz"464748def _pseudoword(rng: random.Random) -> str:49    n = rng.randint(2, 3)50    w = ""51    for _ in range(n):52        w += rng.choice(CONS) + rng.choice(VOWELS)53    if rng.random() < 0.5:54        w += rng.choice(CONS)55    return w565758def morph_items(n: int, seed: int) -> list[tuple[str, str]]:59    """Invented rule: stems ending in a vowel take '-na', stems ending in a consonant take '-ek'."""60    rng = random.Random(f"morph|{seed}")61    out, seen = [], set()62    while len(out) < n:63        w = _pseudoword(rng)64        if w in seen:65            continue66        seen.add(w)67        out.append((w, w + ("na" if w[-1] in VOWELS else "ek")))68    return out697071def addthree_items(n: int, seed: int) -> list[tuple[str, str]]:72    rng = random.Random(f"add3|{seed}")73    xs = rng.sample(range(10, 97), n)74    return [(str(x), str(x + 3)) for x in xs]757677@dataclass78class Tier1Task:79    name: str80    demos: list[tuple[str, str]]81    queries: list[tuple[str, str]]82    alternatives: dict[str, list[str]]8384    @property85    def hash(self) -> str:86        return hashlib.sha256(repr((self.name, self.demos, self.queries)).encode()).hexdigest()[:16]878889def build_tasks(seed: int = 0, n_demo_pool: int = 16, split: str = "dev") -> dict[str, Tier1Task]:90    """Fixed partition: demonstration pool, then the remaining items split into two disjoint halves91    (``dev`` for exploration/selection, ``confirmatory`` for the frozen protocol). The shuffle seed is92    part of the task hash recorded in manifests."""93    rng = random.Random(f"tier1|{seed}")94    tasks = {}95    for name, items, alts in (("antonym", list(ANTONYMS), ALT_ANTONYMS), ("capital", list(CAPITALS), ALT_CAPITALS),96                              ("morph", morph_items(n_demo_pool + 80, seed), {}), ("addthree", addthree_items(n_demo_pool + 70, seed), {})):97        items = list(items)98        rng.shuffle(items)99        demos, rest = items[:n_demo_pool], items[n_demo_pool:]100        half = len(rest) // 2101        queries = rest[:half] if split == "dev" else rest[half : 2 * half]102        tasks[name] = Tier1Task(name, demos, queries, alts)103    return tasks104105106def few_shot_prompt(demos: list[tuple[str, str]], query: str, k: int, rng: random.Random) -> str:107    chosen = rng.sample(demos, k)108    body = "".join(f"Input: {a}\nOutput: {b}\n\n" for a, b in chosen)109    return body + f"Input: {query}\nOutput:"110111112def zero_shot_prompt(query: str) -> str:113    return f"Input: {query}\nOutput:"114115116def is_correct(generated: str, target: str, alternatives: dict[str, list[str]]) -> bool:117    g = generated.strip().split("\n")[0].strip().strip(".").lower()118    ok = {target.lower()} | {a.lower() for a in alternatives.get(target, [])}119    return g in ok120121122# ---------------- unrelated alignment text (synthetic, task-free) ----------------123_SUBJ = ["The committee", "A small boat", "My neighbour", "The old library", "An engineer", "The river", "Our teacher", "The market", "A tired driver",124         "The orchestra", "The farmer", "A young fox", "The council", "Her brother", "The museum", "A quiet village", "The airline", "The bakery"]125_VERB = ["announced", "repaired", "painted", "measured", "visited", "ignored", "described", "borrowed", "organised", "closed", "counted", "followed",126         "moved", "sold", "watched", "finished", "opened", "studied"]127_OBJ = ["the bridge", "three letters", "a wooden fence", "the schedule", "several maps", "the garden", "a broken clock", "the report", "two bicycles",128        "the harbour", "an empty box", "the contract", "the stairs", "a long poem", "the engine", "a glass jar", "the window", "the recipe"]129_TAIL = ["before noon.", "during the storm.", "without any help.", "on Tuesday.", "after the meeting.", "for the second time.", "in silence.",130         "near the station.", "every winter.", "despite the rain.", "", "last year.", "with great care.", "by mistake."]131132133def alignment_sentences(n: int, seed: int) -> list[str]:134    rng = random.Random(f"align|{seed}")135    out, seen = [], set()136    while len(out) < n:137        s = f"{rng.choice(_SUBJ)} {rng.choice(_VERB)} {rng.choice(_OBJ)} {rng.choice(_TAIL)}".strip()138        if rng.random() < 0.3:139            s += f" Then it {rng.choice(_VERB)} {rng.choice(_OBJ)}."140        if s in seen:141            continue142        seen.add(s)143        out.append(s)144    return out145146147_WORDS = ["river", "lamp", "seven", "blue", "engine", "quiet", "paper", "forty", "garden", "signal", "copper", "window", "orbit", "velvet", "ladder",148          "mirror", "pepper", "harbour", "cotton", "thunder", "marble", "yellow", "pencil", "basket", "summer", "winter", "hollow", "tiger", "silver", "rocket"]149150151def alignment_prompts(n: int, seed: int) -> list[str]:152    """Task-free, FORMAT-MATCHED alignment prompts: the same ``Input: …\nOutput:`` frame as the tasks, with153    0–3 preceding demonstration pairs whose outputs are random (no consistent rule) and a random query drawn154    from sentences, words, numbers or names. Nothing here encodes any of the four tasks."""155    rng = random.Random(f"alignprompt|{seed}")156    sents = alignment_sentences(n * 2, seed + 1)157    out, seen = [], set()158159    def item():160        kind = rng.random()161        if kind < 0.35:162            return rng.choice(sents)163        if kind < 0.6:164            return rng.choice(_WORDS)165        if kind < 0.8:166            return str(rng.randint(1, 9999))167        return rng.choice(_SUBJ).split()[-1] + " " + rng.choice(_WORDS)168169    while len(out) < n:170        k = rng.choice([0, 0, 1, 2, 3])171        body = "".join(f"Input: {item()}\nOutput: {item()}\n\n" for _ in range(k))172        pr = body + f"Input: {item()}\nOutput:"173        if pr in seen:174            continue175        seen.add(pr)176        out.append(pr)177    return out178