import glob
import math
import sys
import time
from pathlib import Path
from typing import Optional, Tuple, Union
import math
import lightning as L
import torch
import torch.nn as nn
from lightning.fabric.strategies import FSDPStrategy, XLAStrategy
from torch.utils.data import DataLoader
from functools import partial
# support running without installing as a package
wd = Path(__file__).parent.parent.resolve()
sys.path.append(str(wd))
# from apex.optimizers import FusedAdam #torch optimizer has a cuda backend, which is faster actually
from lit_gpt.model import GPT, Block, Config

from lit_gpt.model_prefix import GPTprefix
from lit_gpt.model_window import GPTwindow

from lit_gpt.model_pos import GPTpos
from lit_gpt.model_relpos import GPTrelpos
from lit_gpt.model_postln import GPTpostln
from lit_gpt.model_kv_bias import GPTkvbias
from lit_gpt.model_attn_sim import GPTattnsim
from lit_gpt.model_attn_kernel import GPTattnkernel

from lit_gpt.packed_dataset import CombinedDataset, PackedDataset
from lit_gpt.speed_monitor import SpeedMonitorFabric as Monitor
from lit_gpt.speed_monitor import estimate_flops, measure_flops
from lit_gpt.utils import chunked_cross_entropy, get_default_supported_precision, num_parameters, step_csv_logger, lazy_load
from pytorch_lightning.loggers import WandbLogger
from lit_gpt import FusedCrossEntropyLoss
import random
import yaml
import os

# tinyllama_1M, tinyllama_60M, tinyllama_1_1b
model_name = "tinyllama_1_1b"

# Experimental settings
reset_embedding = False
group_level_sampling = False
only_save_model = False

# Hyperparameters
total_devices = 8
num_of_devices = 8
num_of_nodes = total_devices // num_of_devices if total_devices >= num_of_devices else 1
# optimal tokens should be 10^20
global_batch_size = 512
learning_rate = 4e-4
min_lr = 1e-5
decay_lr = True

micro_batch_size = 4
max_step = 25000
warmup_steps = 1000
log_step_interval = 10
eval_iters = 50
save_step_interval = 1000
eval_step_interval = 1000
# -100 is the default ignore index
# ignore_token_id = -100

weight_decay = 1e-1
beta1 = 0.9
beta2 = 0.95
grad_clip = 1.0

batch_size = global_batch_size // total_devices
gradient_accumulation_steps = batch_size // micro_batch_size
assert gradient_accumulation_steps > 0
warmup_iters = warmup_steps * gradient_accumulation_steps


max_iters = max_step * gradient_accumulation_steps
lr_decay_iters = max_iters
log_iter_interval = log_step_interval * gradient_accumulation_steps


# Be careful about the weights, it should be something as the len(dataset) * actual reweighting
train_data_config = [
]
val_data_config = [
]

hparams = {k: v for k, v in locals().items() if isinstance(v, (int, float, str)) and not k.startswith("_")}
# get a random name
random_name = ''.join(random.choices('abcdefghijklmnopqrstuvwxyz', k=10))
logger = step_csv_logger("out", random_name, flush_logs_every_n_steps=log_iter_interval)
# log hyper-parameters into wandb
wandb_logger = WandbLogger()

def setup(
    data_seed: int = 3406,
    devices: int = num_of_devices,
    train_data_dir: Path = Path("data/redpajama_sample"),
    val_data_dir: Optional[Path] = None,
    data_yaml_file: Optional[Path] = None,
    precision: Optional[str] = None,
    tpu: bool = False,
    resume: Union[bool, Path] = False,
    out_name: str = "default_model",
    load_from: Optional[Path] = None,
    gpu_memory: Optional[int] = None,
    mode: str = "base",
    downsample_ratio: Optional[float] = 1.0,
    ignore_loss: Optional[int] = 0,
    fix_position: Optional[int] = 1,        # added by Xiangming
) -> None:
    # may modify the global train_data_config
    global train_data_config
    global val_data_config
    
    # if train config exists as a yaml file, load it
    if data_yaml_file is not None:
        data_yaml_file = Path(data_yaml_file)
        if data_yaml_file.exists():
            print("loading config from {}".format(data_yaml_file))
            with open(data_yaml_file, "r") as f:
                # template yaml file is as
                # train_file: weight
                config = yaml.safe_load(f)
            if "data_seed" in config:
                data_seed = int(config["data_seed"])
                print("update data_seed to {}".format(data_seed))
            if "train" in config:
                train_config = []
                for k, v in config["train"].items():
                    train_config.append((k, float(v)))
                # update the config
                train_data_config = train_config
            if "valid" in config:
                val_config = []
                for k, v in config["valid"].items():
                    # TODO: by deafult we use separate validation set
                    val_config.append([(k, float(v))])
                val_data_config = val_config
                    # 093 train and valid
            del config["train"]
            del config["valid"]
            # see if any local variable is in the config, if so, update it
            for k, v in config.items():
                if k in globals():
                    print("update {} to {}".format(k, v))
                    globals()[k] = v
            # use the new value to update
            if "num_of_devices" in config.keys():
                # update devices
                devices = num_of_devices
            # if global batch size is changed, update the batch size
            if "global_batch_size" in config.keys():
                globals()["batch_size"] = global_batch_size // total_devices
                globals()["gradient_accumulation_steps"] = batch_size // micro_batch_size
                globals()["warmup_iters"] = warmup_steps * gradient_accumulation_steps
                globals()["max_iters"] = max_step * gradient_accumulation_steps
                globals()["lr_decay_iters"] = max_iters
                globals()["log_iter_interval"] = log_step_interval * gradient_accumulation_steps                
        else:
            print("config {} does not exist, skip loading".format(data_yaml_file))
    
    precision = precision or get_default_supported_precision(training=True, tpu=tpu)

    if devices > 1:
        if tpu:
            # For multi-host TPU training, the device count for Fabric is limited to the count on a single host.
            devices = "auto"
            strategy = XLAStrategy(sync_module_states=False)
        else:
            strategy = FSDPStrategy(
                # nn.Embedding
                auto_wrap_policy={Block,nn.Embedding},
                # activation_checkpointing_policy={Block},
                state_dict_type="full",
                sharding_strategy="FULL_SHARD",
                limit_all_gathers=True,
                cpu_offload=False,
            )
    else:
        strategy = "auto"
        
    fabric = L.Fabric(devices=devices, 
                      strategy=strategy, 
                      precision=precision, 
                      loggers=[logger, wandb_logger],
                      num_nodes=num_of_nodes)

    fabric.print("precision: {}".format(precision))
    fabric.print("Use gpu memory: {}".format(gpu_memory))

    hparams = {k: v for k, v in globals().items() if isinstance(v, (int, float, str)) and not k.startswith("_")}
    fabric.print(hparams)
    wandb_logger.log_hyperparams(hparams)
    # log the train & val data config
    wandb_logger.log_hyperparams({"train_data_config": train_data_config, "val_data_config": val_data_config})
    #fabric.launch(main, train_data_dir, val_data_dir, resume)
    main(fabric, data_seed, train_data_dir, val_data_dir, resume, out_name, load_from, mode, downsample_ratio, ignore_loss, fix_position)


def main(fabric, data_seed, train_data_dir, val_data_dir, resume, out_name, load_from, mode, downsample_ratio, ignore_loss, fix_position):
    monitor = Monitor(fabric, window_size=2, time_unit="seconds", log_iter_interval=log_iter_interval)
    out_dir = Path("checkpoints") / out_name

    if fabric.global_rank == 0:
        out_dir.mkdir(parents=True, exist_ok=True)

    config = Config.from_name(model_name)

    train_dataloader, val_dataloaders = create_dataloaders(
        batch_size=micro_batch_size,
        block_size=config.block_size,
        fabric=fabric,
        train_data_dir=train_data_dir,
        val_data_dir=val_data_dir,
        seed=data_seed,
        downsample_ratio=downsample_ratio,
    )

    train_dataloader = fabric.setup_dataloaders(train_dataloader)
    if val_dataloaders is not None:
        for i in range(len(val_dataloaders)):
            val_dataloaders[i] = fabric.setup_dataloaders(val_dataloaders[i])

    fabric.seed_everything(data_seed)  # same seed for every process to init model (FSDP)

    fabric.print(f"Loading model with {config.__dict__}")
    t0 = time.perf_counter()
    with fabric.init_module(empty_init=True if num_of_devices > 1 else False):
        if mode == "prefix":
            model = GPTprefix(config)
        elif mode == "window":
            model = GPTwindow(config)
        elif mode == "pos":
            model = GPTpos(config)
        elif mode == "relpos":
            model = GPTrelpos(config)
        elif mode == "postln":
            model = GPTpostln(config)
        elif mode == "kv_bias":
            model = GPTkvbias(config)
        elif mode == "attnsim":
            model = GPTattnsim(config)
        elif mode == "attnkernel":
            model = GPTattnkernel(config)
        else:
            model = GPT(config)
        model.apply(partial(model._init_weights ,n_layer=config.n_layer))
        # load pretrained model
        if load_from is not None:
            # use torch.load to load the model
            print("loading model from {}".format(load_from))
            state_dict = torch.load(load_from, map_location=fabric.device)
            if "model" in state_dict:
                state_dict = state_dict["model"]
            model.load_state_dict(state_dict, strict=True, assign=True)
        # if reset embedding, reset the embedding layer
        if reset_embedding:
            torch.nn.init.normal_(model.transformer.wte.weight,
                                  mean=0.0, std=math.sqrt(2.0 / 5 / model.transformer.wte.weight.size(1)))
            # reset the output layer, also
            torch.nn.init.normal_(model.lm_head.weight,
                                  mean=0.0, std=math.sqrt(2.0 / 5 / model.lm_head.weight.size(1)))
            for n, p in model.named_parameters():
                if "wte" not in n and "lm_head" not in n:
                    p.requires_grad = False
                else:
                    print("resetting {}".format(n))
                    p.requires_grad = True


    fabric.print(f"Time to instantiate model: {time.perf_counter() - t0:.02f} seconds.")
    fabric.print(f"Total parameters {num_parameters(model):,}")

    model = fabric.setup(model)
    
    optimizer = torch.optim.AdamW(
        filter(lambda p: p.requires_grad, model.parameters()),
        lr=learning_rate, weight_decay=weight_decay, betas=(beta1, beta2), fused=True
    )
    # optimizer = FusedAdam(model.parameters(), lr=learning_rate, weight_decay=weight_decay, betas=(beta1, beta2),adam_w_mode=True)
    optimizer = fabric.setup_optimizers(optimizer)

    state = {"model": model, "optimizer": optimizer, "hparams": hparams, "iter_num": 0, "step_count": 0}

    if resume is True:
        resume = sorted(out_dir.glob("*.pth"))
        
    if resume:
        # take the last checkpoint
        resume = resume[-1]
        fabric.print(f"Resuming training from {resume}")
        fabric.load(resume, state)

    train_time = time.perf_counter()
    train(fabric, state, train_dataloader, val_dataloaders, monitor, resume, out_dir, mode, ignore_loss, fix_position)
    fabric.print(f"Training time: {(time.perf_counter()-train_time):.2f}s")
    if fabric.device.type == "cuda":
        fabric.print(f"Memory used: {torch.cuda.max_memory_allocated() / 1e9:.02f} GB")

def delete_all_except_last(folder_path, file_extension=".pth"):
    # 获取文件夹中以特定扩展名结尾的所有文件
    files = [f for f in os.listdir(folder_path) if f.endswith(file_extension)]
    
    if len(files) >= 2:
        files.sort(key=lambda x: os.path.getmtime(os.path.join(folder_path, x)))
        
        for file_name in files[:-1]:
            file_path = os.path.join(folder_path, file_name)
            os.remove(file_path)
            print(f"Deleted: {file_path}")
    else:
        print("Not enough files to delete.")


def train(fabric, state, train_dataloader, val_dataloaders, monitor, resume, out_dir, mode, ignore_loss, fix_position):
    model = state["model"]
    optimizer = state["optimizer"]

    if val_dataloaders is not None:
        # for i in range(len(val_dataloaders)):
        #     validate(fabric, model, val_dataloaders[i])  # sanity check
        # validate(fabric, model, val_dataloaders[0])
        pass

    with torch.device("meta"):
        if mode == "prefix":
            meta_model = GPTprefix(model.config)
        elif mode == "window":
            meta_model = GPTwindow(model.config)
        elif mode == "pos":
            meta_model = GPTpos(model.config)
        elif mode == "relpos":
            meta_model = GPTrelpos(model.config)
        elif mode == "postln":
            meta_model = GPTpostln(model.config)
        elif mode == "kv_bias":
            meta_model = GPTkvbias(model.config)
        elif mode == "attnsim":
            meta_model = GPTattnsim(model.config)
        elif mode == "attnkernel":
            meta_model = GPTattnkernel(model.config)
        else:
            meta_model = GPT(model.config)
        # "estimated" is not as precise as "measured". Estimated is optimistic but widely used in the wild.
        # When comparing MFU or FLOP numbers with other projects that use estimated FLOPs,
        # consider passing `SpeedMonitor(flops_per_batch=estimated_flops)` instead
        estimated_flops = estimate_flops(meta_model) * micro_batch_size
        fabric.print(f"Estimated TFLOPs: {estimated_flops * fabric.world_size / 1e12:.2f}")
        x = torch.randint(0, 1, (micro_batch_size, model.config.block_size))
        # measured_flos run in meta. Will trigger fusedRMSNorm error
        #measured_flops = measure_flops(meta_model, x)
        #fabric.print(f"Measured TFLOPs: {measured_flops * fabric.world_size / 1e12:.2f}")
        del meta_model, x

    total_lengths = 0
    total_t0 = time.perf_counter()

    if fabric.device.type == "xla":
        import torch_xla.core.xla_model as xm

        xm.mark_step()
    
    
    initial_iter = state["iter_num"]
    curr_iter = 0
    
    loss_func = FusedCrossEntropyLoss()
    for  train_data in train_dataloader:
        # resume loader state. This is not elegant but it works. Should rewrite it in the future.
        if resume:
            if curr_iter < initial_iter:
                curr_iter += 1
                fabric.print("skip iter {}".format(curr_iter))
                continue
            else:
                resume = False
                curr_iter = -1
                fabric.barrier()
                fabric.print("resume finished, taken {} seconds".format(time.perf_counter() - total_t0))
        if state["iter_num"] >= max_iters:
            break
        
        if group_level_sampling:
            # sample a one-hot weight from the train_data_config
            weights = [weight for _, weight in train_data_config]
            sum_weights = sum(weights)
            weights = [el / sum_weights for el in weights]
            idx = random.choices(range(len(train_data_config)), weights=weights, k=1)[0]
            # construct the one-hot weight
            weight = [0 for _ in range(len(train_data_config))]
            weight[idx] = 1.0
            # update the weight in the train_loader
            train_dataloader.dataset._weights = weight
            # fabric.print("update weight to {}".format(weight))
        
        # determine and set the learning rate for this iteration
        lr = get_lr(state["iter_num"]) if decay_lr else learning_rate
        for param_group in optimizer.param_groups:
            param_group["lr"] = lr

        iter_t0 = time.perf_counter()

        if mode == "fix":
            train_data = train_data[:, fix_position] = 1 # padding token

        input_ids = train_data[:, 0 : model.config.block_size].contiguous()
        targets = train_data[:, 1 : model.config.block_size + 1].contiguous()
        is_accumulating = (state["iter_num"] + 1) % gradient_accumulation_steps != 0
        with fabric.no_backward_sync(model, enabled=is_accumulating):
            logits = model(input_ids)
            # print(logits[:, ignore_loss:, :].shape)
            loss = loss_func(logits[:, ignore_loss:, :].contiguous(), targets[:, ignore_loss:].contiguous())
            # loss = loss_func(logits, targets)
            # loss = chunked_cross_entropy(logits, targets, chunk_size=0)
            fabric.backward(loss / gradient_accumulation_steps)

        if not is_accumulating:
            grad_norm = fabric.clip_gradients(model, optimizer, max_norm=grad_clip)
            fabric.log_dict({
                "gradient_norm": grad_norm.item()
            })
            optimizer.step()
            optimizer.zero_grad()
            state["step_count"] += 1
        elif fabric.device.type == "xla":
            xm.mark_step()
        state["iter_num"] += 1
        # input_id: B L 
        total_lengths += input_ids.size(1)
        t1 = time.perf_counter()
        fabric.print(
                f"iter {state['iter_num']} learning rate {lr} step {state['step_count']}: loss {loss.item():.4f}, iter time:"
                f" {(t1 - iter_t0) * 1000:.2f}ms{' (optimizer.step)' if not is_accumulating else ''}"
                f" remaining time: {(t1 - total_t0) / (state['iter_num'] - initial_iter) * (max_iters - state['iter_num']) / 3600:.2f} hours. " 
                # print days as well
                f" or {(t1 - total_t0) / (state['iter_num'] - initial_iter) * (max_iters - state['iter_num']) / 3600 / 24:.2f} days. "
            )
 
        monitor.on_train_batch_end(
            state["iter_num"] * micro_batch_size,
            t1 - total_t0,
            # this assumes that device FLOPs are the same and that all devices have the same batch size
            fabric.world_size,
            state["step_count"],
            flops_per_batch=estimated_flops,
            lengths=total_lengths,
            train_loss = loss.item()
        )

            
        if val_dataloaders is not None and not is_accumulating and state["step_count"] % eval_step_interval == 0:
            try:            
                val_names = [name[0][0].replace("valid_", "") for name in val_data_config]
            except Exception as e:
                print(e)
                val_names = [str(i) for i in range(len(val_dataloaders))]
            for i in range(len(val_dataloaders)):
                t0 = time.perf_counter()
                val_loss = validate(fabric, model, val_dataloaders[i], val_names[i], mode=mode, fix_position=fix_position)
                t1 = time.perf_counter() - t0
                monitor.eval_end(t1)
                fabric.print(f"step {state['iter_num']}: {val_names[i]} val loss {val_loss:.4f}, val time: {t1 * 1000:.2f}ms")
                fabric.log_dict({f"metric/{val_names[i]}_val_loss": val_loss.item(), "total_tokens": model.config.block_size * (state["iter_num"] + 1) * micro_batch_size * fabric.world_size}, state["step_count"])
                fabric.log_dict({f"metric/{val_names[i]}_val_ppl": math.exp(val_loss.item()), "total_tokens": model.config.block_size * (state["iter_num"] + 1) * micro_batch_size * fabric.world_size}, state["step_count"])
                fabric.barrier()

        if not is_accumulating and state["step_count"] % save_step_interval == 0:
            checkpoint_path = out_dir / f"iter-{state['step_count']:06d}-ckpt.pth"
            # if fabric.global_rank == 0:
            #     # delete all the checkpoints except the last one
            #     delete_all_except_last(out_dir)
                
            fabric.print(f"Saving checkpoint to {str(checkpoint_path)!r}")
            if only_save_model:
                saved_state = state["model"]
            else:
                saved_state = state
            fabric.save(checkpoint_path, saved_state)

        
@torch.no_grad()
def validate(fabric: L.Fabric, model: torch.nn.Module, val_dataloader: DataLoader, name: str = None, mode: str = 'base', fix_position: int = 1) -> torch.Tensor:
    fabric.print(f"Validating {name} ...")
    model.eval()

    losses = torch.zeros(eval_iters, device=fabric.device)
    for k, val_data in enumerate(val_dataloader):
        # fabric.print("val data: {}".format(val_data))
        if mode == "fix":
            val_data = val_data[:, fix_position] = 1 # padding token
        if k >= eval_iters:
            break
        input_ids = val_data[:, 0 : model.config.block_size].contiguous()
        targets = val_data[:, 1 : model.config.block_size + 1].contiguous()
        logits = model(input_ids)
        loss = chunked_cross_entropy(logits, targets, chunk_size=0)
    
        # loss_func = FusedCrossEntropyLoss()
        # loss = loss_func(logits, targets)
        losses[k] = loss.item()
    
    # print top 100 losses
    # fabric.print("top 100 losses: {}".format(losses[:100]))
    out = losses.mean()

    model.train()
    return out


def create_train_dataloader(
    batch_size: int, block_size: int, data_dir: Path, fabric, shuffle: bool = True, seed: int = 12345, split="train", downsample_ratio: Optional[float] = 1.0,
) -> DataLoader:
    # datasets = []
    data_config = train_data_config if split == "train" else val_data_config
    # check the validness
    for idx in range(len(data_config) - 1, -1, -1):
        prefix = data_config[idx][0]
        filenames = sorted(glob.glob(str(data_dir / f"{prefix}-*")))        
        if len(filenames) < total_devices:
            fabric.print("skip dataset {}".format(prefix))
            del data_config[idx]
            continue
    
    all_filenames = []
    for idx in range(len(data_config)):
        prefix = data_config[idx][0]
        filenames = sorted(glob.glob(str(data_dir / f"{prefix}-*")))
        fabric.print("create dataset {}".format(prefix))
        all_filenames.extend(filenames)
    
    fabric.print("total dataset, the length of filenames {}".format(len(all_filenames)))  # 9749
    random.seed(seed)
    random.shuffle(all_filenames)
    # down-sample
    downsampled_length = int(len(all_filenames) * downsample_ratio) if int(len(all_filenames) * downsample_ratio / total_devices) > 0 else total_devices
    all_filenames = all_filenames[:downsampled_length]
    print(all_filenames)
    fabric.print("downsample dataset with the ratio of {}, the length of filenames {}".format(downsample_ratio, len(all_filenames)))
    combined_dataset = PackedDataset(
        all_filenames,
        # n_chunks control the buffer size. 
        # Note that the buffer size also impacts the random shuffle
        # (PackedDataset is an IterableDataset. So the shuffle is done by prefetch a buffer and shuffle the buffer)
        n_chunks=1,
        block_size=block_size,
        shuffle=shuffle,
        seed=seed+fabric.global_rank,
        num_processes=fabric.world_size,
        process_rank=fabric.global_rank
    )
        # datasets.append(dataset)

    if not combined_dataset:
        raise RuntimeError(
            f"No data found at {data_dir}. Make sure you ran prepare_redpajama.py to create the dataset."
        )

    # weights = [weight for _, weight in data_config]
    # sum_weights = sum(weights)
    # weights = [el / sum_weights for el in weights]

    # combined_dataset = CombinedDataset(datasets=datasets, seed=seed, weights=weights)
    return DataLoader(combined_dataset, batch_size=batch_size, shuffle=False, pin_memory=True)


def create_val_dataloader(
    batch_size: int, block_size: int, data_dir: Path, fabric, shuffle: bool = True, seed: int = 12345, split="train"
) -> DataLoader:
    
    val_data_loaders = []
    
    # check the validness
    for idx in range(len(val_data_config) - 1, -1, -1):
        data_config = val_data_config[idx]
        delete_val_flag = False
        for prefix, _ in data_config:
            filenames = sorted(glob.glob(str(data_dir / f"{prefix}-*")))
            if len(filenames) < total_devices:
                fabric.print("skip val dataset {}".format(prefix))
                delete_val_flag = True
                break
        if delete_val_flag:
            del val_data_config[idx]


    for data_config in val_data_config:
        datasets = []
        for prefix, _ in data_config:
            filenames = sorted(glob.glob(str(data_dir / f"{prefix}-*")))
            random.seed(seed)
            random.shuffle(filenames)

            dataset = PackedDataset(
                filenames,
                # n_chunks control the buffer size. 
                # Note that the buffer size also impacts the random shuffle
                # (PackedDataset is an IterableDataset. So the shuffle is done by prefetch a buffer and shuffle the buffer)
                n_chunks=1,
                block_size=block_size,
                shuffle=shuffle,
                seed=seed+fabric.global_rank,
                num_processes=fabric.world_size,
                process_rank=fabric.global_rank,
            )
            datasets.append(dataset)

        if not datasets:
            raise RuntimeError(
                f"No data found at {data_dir}. Make sure you ran prepare_redpajama.py to create the dataset."
            )

        weights = [weight for _, weight in data_config]
        sum_weights = sum(weights)
        weights = [el / sum_weights for el in weights]

        check_flag = True
        for dataset in datasets:
            if len(dataset._filenames) == 0:
                check_flag = False
                break
        
        if check_flag:
            combined_dataset = CombinedDataset(datasets=datasets, seed=seed, weights=weights)
            val_data_loaders.append(DataLoader(combined_dataset, batch_size=batch_size, shuffle=False, pin_memory=True))
            fabric.print("create val dataset {}".format(data_config))
        else:
            fabric.print("there are something wrong with the val dataset {}".format(data_config))
    return val_data_loaders


def create_dataloaders(
    batch_size: int,
    block_size: int,
    fabric,
    train_data_dir: Path = Path("data/redpajama_sample"),
    val_data_dir: Optional[Path] = None,
    seed: int = 12345,
    downsample_ratio: Optional[float] = 1.0,
) -> Tuple[DataLoader, DataLoader]:
    # Increase by one because we need the next word as well
    effective_block_size = block_size + 1
    train_dataloader = create_train_dataloader(
        batch_size=batch_size,
        block_size=effective_block_size,
        fabric=fabric,
        data_dir=train_data_dir,
        shuffle=True,
        seed=seed,
        split="train",
        downsample_ratio=downsample_ratio,
    )
    val_dataloader = (
        create_val_dataloader(
            batch_size=batch_size,
            block_size=effective_block_size,
            fabric=fabric,
            data_dir=val_data_dir,
            shuffle=False,
            seed=seed,
            split="validation"
        )
        if val_data_dir
        else None
    )
    return train_dataloader, val_dataloader


# learning rate decay scheduler (cosine with warmup)
def get_lr(it):
    # 1) linear warmup for warmup_iters steps
    if it < warmup_iters:
        return learning_rate * it / warmup_iters
    # 2) if it > lr_decay_iters, return min learning rate
    if it > lr_decay_iters:
        return min_lr
    # 3) in between, use cosine decay down to min learning rate
    decay_ratio = (it - warmup_iters) / (lr_decay_iters - warmup_iters)
    assert 0 <= decay_ratio <= 1
    coeff = 0.5 * (1.0 + math.cos(math.pi * decay_ratio))  # coeff ranges 0..1
    return min_lr + coeff * (learning_rate - min_lr)


if __name__ == "__main__":
    # Uncomment this line if you see an error: "Expected is_sm80 to be true, but got false"
    # torch.backends.cuda.enable_flash_sdp(False)
    torch.set_float32_matmul_precision("high")

    from jsonargparse import CLI

    CLI(setup)
