def main()

in tasks/add-or-mul/generate_data.py [0:0]


def main(train_length, test_span):
    try:
        shutil.rmtree(str(train_length))
    except:
        pass

    pathlib.Path(str(train_length)).mkdir()

    generate_fpa(root=f'{train_length}/fpa/', train_length=train_length, test_span=test_span)

    additive = lambda x: x + train_length
    memorization = lambda _: 2 * train_length
    multiplicative = lambda x: 2 * x

    generate_mdl(root=f'{train_length}/mem/', train_length=train_length, test_span=test_span, rule=memorization)
    generate_mdl(root=f'{train_length}/add/', train_length=train_length, test_span=test_span, rule=additive)
    generate_mdl(root=f'{train_length}/mul/', train_length=train_length, test_span=test_span, rule=multiplicative)