import torch
from datasets import load_dataset
from transformers import (
    GPT2Config,
    GPT2LMHeadModel,
    GPT2TokenizerFast,
    Trainer,
    TrainingArguments,
    DataCollatorForLanguageModeling,
)

# ===== Config =====
OUTPUT_DIR = "./TinyGPT-8M"
MODEL_NAME = "gpt2"  # tokenizer only

# ===== Tokenizer =====
tokenizer = GPT2TokenizerFast.from_pretrained(MODEL_NAME)

if tokenizer.pad_token is None:
    tokenizer.pad_token = tokenizer.eos_token

# ===== Dataset =====
print("Loading TinyStories...")
dataset = load_dataset("roneneldan/TinyStories", split="train")

# Use a tiny subset for fast training (<5 min)
dataset = dataset.select(range(10000))

def tokenize_function(examples):
    return tokenizer(
        examples["text"],
        truncation=True,
        padding="max_length",
        max_length=128,
    )

print("Tokenizing...")
tokenized_dataset = dataset.map(
    tokenize_function,
    batched=True,
    remove_columns=dataset.column_names,
)

# ===== Tiny GPT (~7M params) =====
config = GPT2Config(
    vocab_size=len(tokenizer),
    n_positions=128,
    n_ctx=128,
    n_embd=128,
    n_layer=4,
    n_head=4,
    bos_token_id=tokenizer.bos_token_id,
    eos_token_id=tokenizer.eos_token_id,
)

model = GPT2LMHeadModel(config)

# Resize embeddings if tokenizer size differs
model.resize_token_embeddings(len(tokenizer))

# Print parameter count
num_params = sum(p.numel() for p in model.parameters())
print(f"Parameters: {num_params:,}")

# ===== Data collator =====
data_collator = DataCollatorForLanguageModeling(
    tokenizer=tokenizer,
    mlm=False,
)

# ===== Training arguments =====
training_args = TrainingArguments(
    output_dir=OUTPUT_DIR,
    num_train_epochs=10,
    per_device_train_batch_size=32,
    gradient_accumulation_steps=1,
    learning_rate=5e-4,
    weight_decay=0.01,
    logging_steps=20,
    save_strategy="no",
    report_to="none",
    fp16=torch.cuda.is_available(),
    bf16=False,
    dataloader_num_workers=2,
    remove_unused_columns=False,
)

# ===== Trainer =====
trainer = Trainer(
    model=model,
    args=training_args,
    train_dataset=tokenized_dataset,
    data_collator=data_collator,
)

# ===== Train =====
print("Starting training...")
trainer.train()

# ===== Save =====
print("Saving model...")
trainer.save_model(OUTPUT_DIR)
tokenizer.save_pretrained(OUTPUT_DIR)

print(f"Model saved to: {OUTPUT_DIR}")

# ===== Test generation =====
prompt = "Once upon a time"
inputs = tokenizer(prompt, return_tensors="pt")

device = "cuda" if torch.cuda.is_available() else "cpu"
model.to(device)
inputs = {k: v.to(device) for k, v in inputs.items()}

output = model.generate(
    **inputs,
    max_new_tokens=50,
    do_sample=True,
    temperature=0.8,
    top_p=0.95,
)

print("\nSample output:")
print(tokenizer.decode(output[0], skip_special_tokens=True))