GPT2_MEDIUM_CONFIG: dict[str, int | float | bool] = {
    "vocab_size": 50257,
    "context_length": 1024,
    "emb_dim": 1024,
    "n_heads": 16,
    "n_layers": 24,
    "drop_rate": 0.1,
    "qkv_bias": False
}
GPT_CONFIG_124M: dict[str, int | float | bool] = {
    "vocab_size": 50257,    # Vocabulary size
    "context_length": 1024,  # Context length
    "emb_dim": 768,         # Embedding dimension
    "n_heads": 12,          # Number of attention heads
    "n_layers": 12,         # Number of layers
    "drop_rate": 0.1,       # Dropout rate
    "qkv_bias": False       # Query-Key-Value bias
}
GPT_CONFIG_355M: dict[str, int | float | bool] = {
    "vocab_size": 50257,
    "context_length": 1024,
    "emb_dim": 1024,
    "n_heads": 16,
    "n_layers": 24,
    "drop_rate": 0.1,
    "qkv_bias": False
}
GPT_CONFIG_774M: dict[str, int | float | bool] = {
    "vocab_size": 50257,
    "context_length": 1024,
    "emb_dim": 1280,
    "n_heads": 20,
    "n_layers": 36,
    "drop_rate": 0.1,
    "qkv_bias": False
}
GPT_CONFIG_1500M: dict[str, int | float | bool] = {
    "vocab_size": 50257,
    "context_length": 1024,
    "emb_dim": 1600,
    "n_heads": 25,
    "n_layers": 48,
    "drop_rate": 0.1,
    "qkv_bias": False
}

HYPER_PARAMETERS = {
    # 学习率
    "lr": 0.0004,
    # 每一步都会额外减去一部分参数本身的大小
    "weight_decay": 0.1,
    # 训练集遍历次数
    "num_epochs": 40,
    # 每几次迭代进行一次评估
    "eval_freq": 5,
    # 每次评估使用几个batch(全量验证太慢)
    "eval_iter": 1
}
