RAM_Altman_GPT / src /model_utils.py
brunuu's picture
1.0
7de15fc verified
Raw History Blame Contribute Delete
1.13 kB
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