task.py

"""A task that cannot be done by looking at the last few characters.

Each line pairs keys with values, then asks for one of them again:

    a3 f7 c1 f -> 7

Answering needs the model to find the earlier occurrence of the queried key and
copy what followed it. The distance back is variable, so no fixed offset works;
the model has to match on content. This is the smallest thing I could construct
that a fixed-window model is structurally unable to do and attention can.
"""

import random

KEYS = "abcdef"
VALUES = "0123456789"
PAIRS = 3


def make_line(rng):
    keys = rng.sample(KEYS, PAIRS)
    values = [rng.choice(VALUES) for _ in range(PAIRS)]
    query = rng.choice(range(PAIRS))
    prompt = " ".join(k + v for k, v in zip(keys, values))
    return f"{prompt} {keys[query]}", values[query]


def make_corpus(n, seed=1):
    """Returns (text, list of (prompt, answer)) — the text is the lines joined."""
    rng = random.Random(seed)
    items = [make_line(rng) for _ in range(n)]
    text = "".join(f"{p}{a}\n" for p, a in items)
    return text, items


def accuracy(predict, items):
    """predict(prompt) -> character. Fraction of answers exactly right."""
    right = 0
    for prompt, answer in items:
        if predict(prompt) == answer:
            right += 1
    return right / len(items)


if __name__ == "__main__":
    text, items = make_corpus(6, seed=3)
    print(text.rstrip())
    print()
    print("chance accuracy is 1/10 =", 0.1)