# Use GPT-2 Medium instead of GPT-2 Small for instruction fine-tuning

from gpt_download import download_and_load_gpt2
from chap05 import GPTModel, load_weights_into_gpt


BASE_CONFIG = {
    "vocab_size": 50257,
    "context_length": 1024,
    "drop_rate": 0.1,
    "qkv_bias": True
}

model_configs = {
    "gpt2-medium (355M)": {
        "emb_dim": 1024,
        "n_layers": 24,
        "n_heads": 16
    }
}

CHOOSEN_MODEL = "gpt2-medium (355M)"
BASE_CONFIG.update(model_configs[CHOOSEN_MODEL])

model_size = CHOOSEN_MODEL.split("(")[-1].lstrip(")").rstrip(")")
settings, params = download_and_load_gpt2(
    model_size=model_size,
    models_dir="gpt2"
)

model = GPTModel(BASE_CONFIG)
load_weights_into_gpt(model, params)
model.eval()

        __
        __