import numpy as np

META = {'name':'masked_multitoken_modes','domain':'sequence','description':'Binary multi-token sequences from two global modes, trained with masked-token reconstruction.'}
N = 8
P_MODE = 0.8

def get_dataset(seed, n_train, n_test):
    rng = np.random.default_rng(seed)
    def sample(n):
        mode = rng.random(n) < P_MODE
        alt = np.tile(np.array([0,1,0,1,0,1,0,1], dtype=np.float32), (n,1))
        x = np.where(mode[:,None], 1.0, alt)
        noise = rng.random((n,N)) < .03
        return np.where(noise, 1.0-x, x).astype(np.float32)
    return {'xtr':sample(n_train), 'ytr':sample(n_train), 'xte':sample(n_test),
            'yte':sample(n_test), 'task':'regression', 'metric':'bce'}