source/data.py

from __future__ import annotations

import numpy as np

KEYS = 16
VALUES = 10
VALUE_SENTINEL = VALUES


def generate_bindings(
    samples: int,
    pairs: int,
    distractors: int,
    seed: int,
) -> tuple[np.ndarray, np.ndarray, np.ndarray, np.ndarray]:
    if pairs > KEYS:
        raise ValueError("Pairs cannot exceed the unique key vocabulary.")
    rng = np.random.default_rng(seed)
    length = pairs + distractors + 1
    keys = np.zeros((samples, length), dtype=np.int64)
    values = np.zeros((samples, length), dtype=np.int64)
    writes = np.zeros((samples, length), dtype=np.float32)
    targets = np.zeros(samples, dtype=np.int64)
    for row in range(samples):
        bound_keys = rng.choice(KEYS, pairs, replace=False)
        bound_values = rng.integers(0, VALUES, size=pairs)
        events = [
            (int(key), int(value), 1.0)
            for key, value in zip(bound_keys, bound_values, strict=True)
        ]
        events.extend(
            [
                (
                    int(rng.integers(0, KEYS)),
                    int(rng.integers(0, VALUES)),
                    0.0,
                )
                for _ in range(distractors)
            ]
        )
        rng.shuffle(events)
        for step, (key, value, write) in enumerate(events):
            keys[row, step] = key
            values[row, step] = value
            writes[row, step] = write
        query = int(rng.integers(0, pairs))
        keys[row, -1] = bound_keys[query]
        values[row, -1] = VALUE_SENTINEL
        targets[row] = bound_values[query]
    return keys, values, writes, targets