def make_char_dataset(text, block_size):
    chars = sorted(set(text))
    stoi = {ch: i for i, ch in enumerate(chars)}
    itos = {i: ch for ch, i in stoi.items()}
    data = torch.tensor([stoi[c] for c in text], dtype=torch.long)
    X, Y = [], []
    for i in range(0, len(data) - block_size - 1, block_size):
        X.append(data[i:i + block_size])
        Y.append(data[i + 1:i + block_size + 1])
    return torch.stack(X), torch.stack(Y), stoi, itos


if __name__ == "__main__":
    torch.manual_seed(0)
    # Public-domain text (Project Gutenberg style excerpt) stands
    # in for a full corpus in this minimal, self-contained example.
    text = (
        "Four score and seven years ago our fathers brought forth "
        "on this continent a new nation conceived in liberty and "
        "dedicated to the proposition that all men are created "
        "equal. " * 30
    )
    block_size = 32
    X, Y, stoi, itos = make_char_dataset(text, block_size)

    model = MiniGPT(
        vocab_size=len(stoi), d_model=128, n_heads=4,
        n_layers=4, max_len=block_size,
    )
    optimizer = torch.optim.AdamW(model.parameters(), lr=3e-4)

    for step in range(300):
        logits = model(X)
        loss = F.cross_entropy(
            logits.reshape(-1, len(stoi)), Y.reshape(-1)
        )
        optimizer.zero_grad()
        loss.backward()
        optimizer.step()
        if step % 100 == 0:
            print(f"step {step}  loss {loss.item():.4f}")

    prompt = torch.tensor(
        [[stoi[c] for c in "Four score"]], dtype=torch.long
    )
    generated = model.generate(prompt, max_new_tokens=60)
    text_out = "".join(itos[i] for i in generated[0].tolist())
    print("\nGenerated:", text_out)
