1import numpy as np
 2
 3META = {'name':'masked_multitoken_modes','domain':'sequence','description':'Binary multi-token sequences from two global modes, trained with masked-token reconstruction.'}
 4N = 8
 5P_MODE = 0.8
 6
 7def get_dataset(seed, n_train, n_test):
 8    rng = np.random.default_rng(seed)
 9    def sample(n):
10        mode = rng.random(n) < P_MODE
11        alt = np.tile(np.array([0,1,0,1,0,1,0,1], dtype=np.float32), (n,1))
12        x = np.where(mode[:,None], 1.0, alt)
13        noise = rng.random((n,N)) < .03
14        return np.where(noise, 1.0-x, x).astype(np.float32)
15    return {'xtr':sample(n_train), 'ytr':sample(n_train), 'xte':sample(n_test),
16            'yte':sample(n_test), 'task':'regression', 'metric':'bce'}