RAM_Altman_GPT / src /create_RAM_Altman_GPT.py
brunuu's picture
1.0
7de15fc verified
Raw History Blame Contribute Delete
1.15 kB
# RAM_Altman_GPT
# A tiny GPT-2 with that answers everything with RAM!!
from transformers import GPT2LMHeadModel
from config import DATA_FILE, DEVICE, OUTPUT_DIR, set_seeds
from tokenizer_utils import build_tokenizer
from model_utils import build_model
from train import load_token_ids, train_loop
def main():
set_seeds(0)
# tokenizer
hf_tokenizer, vocab_size = build_tokenizer()
print(f"vocabulary: {vocab_size} tokens | device: {DEVICE}")
# model
model: GPT2LMHeadModel
context_length: int
model, context_length = build_model(hf_tokenizer, vocab_size, DEVICE)
print(f"context: {context_length} | parameters: {model.num_parameters():,} | RAM required: ALL")
# data
token_ids, weights = load_token_ids(hf_tokenizer, DATA_FILE, DEVICE)
print(f"dataset: {len(token_ids):,} tokens ({(weights > 0).float().mean():.0%} are RAM-ALTMAN replies, the only part graded)")
# train
train_loop(model, token_ids, weights, context_length)
model.save_pretrained(OUTPUT_DIR)
hf_tokenizer.save_pretrained(OUTPUT_DIR)
print(f"saved to {OUTPUT_DIR}")
if __name__ == "__main__":
main()