from transformers import GPT2Config, GPT2LMHeadModel, GPT2TokenizerFast
from config import EMBEDDING_SIZE, NUM_HEADS, NUM_LAYERS, TARGET_PARAMS

def context_length_for(vocab_size: int) -> int:
    tokens_per_layer = 12 * EMBEDDING_SIZE + 13
    return TARGET_PARAMS // EMBEDDING_SIZE - vocab_size - 2 - NUM_LAYERS * tokens_per_layer

def build_model(hf_tokenizer: GPT2TokenizerFast, vocab_size: int, device: str) -> tuple[GPT2LMHeadModel, int]:
    context_length = context_length_for(vocab_size)

    config = GPT2Config(
        vocab_size = vocab_size,
        n_positions = context_length,
        n_ctx = context_length,
        n_embd = EMBEDDING_SIZE,
        n_layer = NUM_LAYERS,
        n_head = NUM_HEADS,
        n_inner = 4 * EMBEDDING_SIZE, # llama.cpp expects the FFN to be 4x the embedding size
        bos_token_id = hf_tokenizer.eos_token_id,
        eos_token_id = hf_tokenizer.eos_token_id,
    )

    model = GPT2LMHeadModel(config)
    total = model.num_parameters()
    assert total == TARGET_PARAMS, f"expected {TARGET_PARAMS:,} parameters, got {total:,}"
    model.to(device)
    return model, context_length