MAX_MODEL_CONTEXT = 256
TRAIN_CONTEXTS = [32, 64, 128, 256]

MODEL_CONFIG = {
    "vocab_size": 50_257,
    "context_length": MAX_MODEL_CONTEXT,
    "emb_dim": 128,
    "n_heads": 4,
    "n_layers": 4,
    "drop_rate": 0.1,
    "qkv_bias": False,
}

        __
        __