# 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()