import torch
import torch.nn.functional as F
from einops import rearrange
from model import MegabyteConfig, Megabyte

V = 512         # vocabulary size, input bytes have 256 characters, and the extra 256 are reserved for special tokens.
P = 4           # patch size
D_G = 512       # global model dimension
D_L = 128       # local model dimension
T = 1024        # sequence length
B = 2           # batch size
K = T//P        # number of patches
PAD_ID = 257    # padding token id
EOS_ID = 258    # end of sequence token id

config = MegabyteConfig(
    V=V,
    P=P,
    D_G=D_G,
    D_L=D_L,
    T_MAX=T,
    initializer_range=0.02, # Parameter initialization value range
    g_nlayers=4,            # number of global model layers
    g_nheads=32,            # number of global model attention heads
    l_nlayers=2,            # number of local model attention layers
    l_nheads=2,             # number of local model attention heads
    pad_id=PAD_ID,
    eos_id=EOS_ID,
)
megabyte = Megabyte(config)
input_ids = torch.randint(0, 255, (B, T))
# Autoregressive learning, megabyte will learn from the inputs input[:, :-1], labels input[:, :], and learn to predict the next token.
loss = megabyte(input_ids, return_loss=True).loss
loss.backward()

print(loss.norm())

...
from model.megabyte_transformers import MegabyteLMHeadModel, MegabyteTokenizer
lm_head_megabyte = MegabyteLMHeadModel.from_native_megabyte(megabyte)
tokenizer = MegabyteTokenizer(
    eos_token_id=lm_head_megabyte.config.eos_token_id,
)

inputs = tokenizer("Today is", return_tensors="pt")
outputs = lm_head_megabyte.generate(
    **inputs,
    max_new_tokens=5,
    return_dict_in_generate=True,
    output_scores=True,
)

texts = tokenizer.decode(outputs.sequences)
print(texts)