#!/usr/bin/env python
# coding=utf-8
# Copyright 2021 The HuggingFace Inc. team. All rights reserved.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#     http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
"""
Fine-tuning the library models for causal language modeling (GPT, GPT-2, CTRL, ...)
on a text file or a dataset without using HuggingFace Trainer.

Here is the full list of checkpoints on the hub that can be fine-tuned by this script:
https://huggingface.co/models?filter=text-generation
"""
# You can also adapt this script on your own causal language modeling task. Pointers for this are left as comments.

import argparse
import json
import logging
import math
import os
import random
from itertools import chain
from pathlib import Path
import numpy as np
import datasets
import torch
from datasets import load_dataset
from torch.utils.data import DataLoader
from tqdm.auto import tqdm
from schedulefree import AdamWScheduleFree
import transformers
from transformers import AutoTokenizer, default_data_collator

from model import GPT, GPTConfig
from model_reg import GPT_reg

from transformers.utils import check_min_version
from transformers.utils.versions import require_version


# Will error if the minimal version of Transformers is not installed. Remove at your own risks.
check_min_version("4.32.0.dev0")

logger = logging.getLogger(__name__)

require_version("datasets>=1.8.0", "To fix: pip install -r examples/pytorch/language-modeling/requirements.txt")


def parse_args():
    parser = argparse.ArgumentParser(description="Finetune a transformers model on a causal language modeling task")
    parser.add_argument("--dataset_name", type=str, default="wikitext", help="The name of the dataset to use (via the datasets library).")
    parser.add_argument("--dataset_config_name", type=str, default="wikitext-103-raw-v1", help="The configuration name of the dataset to use (via the datasets library).")
    parser.add_argument("--validation_split_percentage", default=5, help="The percentage of the train set used as validation set in case there's no validation split")
    parser.add_argument("--model_name_or_path", type=str, help="Path to pretrained model or model identifier from huggingface.co/models.", required=False)
    parser.add_argument("--tokenizer_name", type=str, default=None, help="Pretrained tokenizer name or path if not the same as model_name")
    parser.add_argument(
        "--use_slow_tokenizer",
        action="store_true",
        help="If passed, will use a slow tokenizer (not backed by the 🤗 Tokenizers library).",
    )
    parser.add_argument("--per_device_train_batch_size", type=int, default=8, help="Batch size (per device) for the training dataloader.")
    parser.add_argument("--per_device_eval_batch_size", type=int, default=8, help="Batch size (per device) for the evaluation dataloader.")
    parser.add_argument("--learning_rate", type=float, default=1e-4, help="Initial learning rate (after the potential warmup period) to use.")
    parser.add_argument("--weight_decay", type=float, default=0.05, help="Weight decay to use.")
    parser.add_argument("--num_train_epochs", type=int, default=20, help="Total number of training epochs to perform.")
    parser.add_argument("--max_train_steps", type=int, default=None, help="Total number of training steps to perform. If provided, overrides num_train_epochs.")
    parser.add_argument("--num_warmup_steps", type=int, default=2000, help="Number of steps for the warmup in the lr scheduler.")
    parser.add_argument("--output_dir", type=str, default=None, help="Where to store the final model.")
    parser.add_argument("--seed", type=int, default=0, help="A seed for reproducible training.")
    parser.add_argument(
        "--block_size",
        type=int,
        default=None,
        help=(
            "Optional input sequence length after tokenization. The training dataset will be truncated in block of"
            " this size for training. Default to the model max input length for single sentence inputs (take into"
            " account special tokens)."
        ),
    )
    parser.add_argument("--preprocessing_num_workers", type=int, default=None, help="The number of processes to use for the preprocessing.")
    parser.add_argument("--overwrite_cache", action="store_true", help="Overwrite the cached training and evaluation sets")    
    parser.add_argument("--checkpointing_steps", type=str, default=None, help="Whether the various states should be saved at the end of every n steps, or 'epoch' for each epoch.")
    parser.add_argument("--resume_from_checkpoint", type=str, default=None, help="If the training should continue from a checkpoint folder.")
    parser.add_argument("--reg", type=float, default=0.0, help="reg_param")
    parser.add_argument("--dropout", type=float, default=0.0, help="dropout_probability")
    parser.add_argument("--device", type=str, default="cuda" if torch.cuda.is_available() else "cpu", help="Device to use for training")

    args = parser.parse_args()

    # Sanity checks
    if args.dataset_name is None:
        raise ValueError("Need either a dataset name or a training/validation file.")

    return args

def set_seed(seed):
    random.seed(seed)
    np.random.seed(seed)
    torch.manual_seed(seed)
    torch.cuda.manual_seed_all(seed)

def main():
    args = parse_args()

    # Sending telemetry. Tracking the example usage helps us better allocate resources to maintain them. The
    # information sent is the one passed as arguments along with your Python/PyTorch versions.
    #send_example_telemetry("run_clm_no_trainer", args)

    logging.basicConfig(
        format="%(asctime)s - %(levelname)s - %(name)s - %(message)s",
        datefmt="%m/%d/%Y %H:%M:%S",
        level=logging.INFO,
    )

    # If passed along, set the training seed now.
    if args.seed is not None:
        set_seed(args.seed)

    if args.output_dir is not None:
        os.makedirs(args.output_dir, exist_ok=True)
    
    raw_datasets = load_dataset(args.dataset_name, args.dataset_config_name)
    if "validation" not in raw_datasets.keys():
        raw_datasets["validation"] = load_dataset(
            args.dataset_name,
            args.dataset_config_name,
            split=f"train[:{args.validation_split_percentage}%]",
        )
        raw_datasets["train"] = load_dataset(
            args.dataset_name,
            args.dataset_config_name,
            split=f"train[{args.validation_split_percentage}%:]",
        )

    config = GPTConfig(
        block_size=1024,
        vocab_size=50304,
        n_layer=12,
        n_head=12,
        n_embd=768,
        dropout=args.dropout,
        bias=True
    )

    if args.tokenizer_name:
        tokenizer = AutoTokenizer.from_pretrained(args.tokenizer_name, use_fast=not args.use_slow_tokenizer)
    elif args.model_name_or_path:
        tokenizer = AutoTokenizer.from_pretrained(args.model_name_or_path, use_fast=not args.use_slow_tokenizer)
    else:
        raise ValueError(
            "You are instantiating a new tokenizer from scratch. This is not supported by this script."
            "You can do it from another script, save it, and load it from here, using --tokenizer_name."
        )

    logger.info("Initializing nano GPT-2 model")
    if args.reg > 0:
        model = GPT_reg(config)
    else:
        model = GPT(config)
    device = torch.device(args.device)
    model.to(device)
    
    # We resize the embeddings only when necessary to avoid index errors. If you are creating a model from scratch

    # embedding_size = model.transformer.wte.weight.shape[0]
    # if len(tokenizer) > embedding_size:
    #     old_weight = model.transformer.wte.weight
    #     model.transformer.wte = nn.Embedding(len(tokenizer), old_weight.shape[1])
    #     model.transformer.wte.weight.data[:old_weight.shape[0], :] = old_weight
    #     model.lm_head = nn.Linear(old_weight.shape[1], len(tokenizer), bias=False)
    #     model.lm_head.weight = model.transformer.wte.weight 

    # Preprocessing the datasets.
    # First we tokenize all the texts.
    column_names = raw_datasets["train"].column_names
    text_column_name = "text" if "text" in column_names else column_names[0]

    def tokenize_function(examples):
        return tokenizer(examples[text_column_name])

    tokenized_datasets = raw_datasets.map(
        tokenize_function,
        batched=True,
        num_proc=args.preprocessing_num_workers,
        remove_columns=column_names,
        load_from_cache_file=not args.overwrite_cache,
        desc="Running tokenizer on dataset",
    )

    if args.block_size is None:
        block_size = tokenizer.model_max_length
        if block_size > 1024:
            logger.warning(
                "The chosen tokenizer supports a `model_max_length` that is longer than the default `block_size` value"
                " of 1024. If you would like to use a longer `block_size` up to `tokenizer.model_max_length` you can"
                " override this default with `--block_size xxx`."
            )
        block_size = 1024
    else:
        if args.block_size > tokenizer.model_max_length:
            logger.warning(
                f"The block_size passed ({args.block_size}) is larger than the maximum length for the model"
                f"({tokenizer.model_max_length}). Using block_size={tokenizer.model_max_length}."
            )
        block_size = min(args.block_size, tokenizer.model_max_length)

    # Main data processing function that will concatenate all texts from our dataset and generate chunks of block_size.
    def group_texts(examples):
        # Concatenate all texts.
        concatenated_examples = {k: list(chain(*examples[k])) for k in examples.keys()}
        total_length = len(concatenated_examples[list(examples.keys())[0]])
        # We drop the small remainder, and if the total_length < block_size  we exclude this batch and return an empty dict.
        # We could add padding if the model supported it instead of this drop, you can customize this part to your needs.
        total_length = (total_length // block_size) * block_size
        # Split by chunks of max_len.
        result = {
            k: [t[i : i + block_size] for i in range(0, total_length, block_size)]
            for k, t in concatenated_examples.items()
        }
        result["labels"] = result["input_ids"].copy()
        return result


    lm_datasets = tokenized_datasets.map(
        group_texts,
        batched=True,
        num_proc=args.preprocessing_num_workers,
        load_from_cache_file=not args.overwrite_cache,
        desc=f"Grouping texts in chunks of {block_size}",
    )

    train_dataset = lm_datasets["train"]
    eval_dataset = lm_datasets["validation"]

    # # Log a few random samples from the training set:
    # for index in random.sample(range(len(train_dataset)), 3):
    #     logger.info(f"Sample {index} of the training set: {train_dataset[index]}.")

    # DataLoaders creation:
    train_dataloader = DataLoader(
        train_dataset, shuffle=True, collate_fn=default_data_collator, batch_size=args.per_device_train_batch_size
    )
    eval_dataloader = DataLoader(
        eval_dataset, collate_fn=default_data_collator, batch_size=args.per_device_eval_batch_size
    )

    # Scheduler and math around the number of training steps.
    overrode_max_train_steps = False
    num_update_steps_per_epoch = math.ceil(len(train_dataloader))
    if args.max_train_steps is None:
        args.max_train_steps = args.num_train_epochs * num_update_steps_per_epoch
        overrode_max_train_steps = True

    optimizer = AdamWScheduleFree(model.parameters(), lr=args.learning_rate, weight_decay =0.05, warmup_steps =2000)
    
    num_update_steps_per_epoch = math.ceil(len(train_dataloader))
    if overrode_max_train_steps:
        args.max_train_steps = args.num_train_epochs * num_update_steps_per_epoch
    # Afterwards we recalculate our number of training epochs
    args.num_train_epochs = math.ceil(args.max_train_steps / num_update_steps_per_epoch)

    checkpointing_steps = args.checkpointing_steps
    if checkpointing_steps is not None and checkpointing_steps.isdigit():
        checkpointing_steps = int(checkpointing_steps)

    logger.info("***** Running training *****")
    logger.info(f"  Num examples = {len(train_dataset)}")
    logger.info(f"  Num Epochs = {args.num_train_epochs}")
    logger.info(f"  Instantaneous batch size per device = {args.per_device_train_batch_size}")
    logger.info(f"  Total train batch size (w. parallel, distributed & accumulation) = {args.per_device_train_batch_size}")
    logger.info(f"  Total optimization steps = {args.max_train_steps}")
    # Only show the progress bar once on each machine.
    progress_bar = tqdm(range(args.max_train_steps))
    completed_steps = 0
    starting_epoch = 0

    # Load checkpoint if resuming
    if args.resume_from_checkpoint:
        if os.path.isfile(os.path.join(args.resume_from_checkpoint, "pytorch_model.bin")):
            logger.info(f"Loading model from {args.resume_from_checkpoint}")
            checkpoint = torch.load(os.path.join(args.resume_from_checkpoint, "pytorch_model.bin"), map_location=device)
            model.load_state_dict(checkpoint["model_state_dict"])
            
            # Try to determine the epoch from the checkpoint path
            path = os.path.basename(args.resume_from_checkpoint)
            if "epoch_" in path:
                starting_epoch = int(path.replace("epoch_", "")) + 1
                completed_steps = starting_epoch * num_update_steps_per_epoch
            elif "step_" in path:
                completed_steps = int(path.replace("step_", ""))
                starting_epoch = completed_steps // num_update_steps_per_epoch
            
            progress_bar.update(completed_steps)
    
    # Initialize training log
    log_file_path = os.path.join(args.output_dir, "training_log.txt")
    
    for epoch in range(starting_epoch, args.num_train_epochs):
        optimizer.train()
        model.train()
        total_train_loss = 0.0
        train_loss = 0.0
        
        for step, batch in enumerate(train_dataloader):
            if args.resume_from_checkpoint and epoch == starting_epoch and step < completed_steps % num_update_steps_per_epoch:
                continue
            batch = {k: v.to(device) for k, v in batch.items()}
            
            logits, loss_, regloss = model(batch["input_ids"], batch["labels"])
            loss = loss_ + args.reg * regloss
            # We keep track of the loss at each epoch
            loss.backward()
            
            train_loss += loss_.detach().float()
            total_train_loss += loss.detach().float()
                
            optimizer.step()
            optimizer.zero_grad()

            progress_bar.update(1)
            completed_steps += 1

            if isinstance(checkpointing_steps, int):
                if completed_steps % checkpointing_steps == 0:
                    output_dir = f"step_{completed_steps }"
                    if args.output_dir is not None:
                        output_dir = os.path.join(args.output_dir, output_dir)
                    torch.save(
                        {
                            "model_state_dict": model.state_dict(),
                            "optimizer_state_dict": optimizer.state_dict(),
                            "completed_steps": completed_steps,
                            "epoch": epoch,
                        },
                        os.path.join(output_dir, "pytorch_model.bin")
                    )
                        
            if completed_steps >= args.max_train_steps:
                break
        total_tokens = 0
        total_loss = 0.0
        optimizer.eval()
        model.eval()
        losses = []
        for step, batch in enumerate(eval_dataloader):
            with torch.no_grad():
                batch = {k: v.to(device) for k, v in batch.items()}
                logits, loss, regloss = model(batch["input_ids"], batch["labels"])
            mask = batch["labels"].ne(-100) 
            num_tokens = mask.sum().item()
            total_loss += loss.item() * num_tokens
            total_tokens += num_tokens
        
        eval_loss = total_loss / total_tokens
        perplexity = math.exp(eval_loss)

        logger.info(f"epoch {epoch}: perplexity: {perplexity} eval_loss: {eval_loss}")

        if args.output_dir is not None:
            # Save the model checkpoint
            torch.save(
                {
                    "model_state_dict": model.state_dict(),
                    "config": vars(model.config),
                    "vocab_size": model.transformer.wte.weight.shape[0],
                    "perplexity": perplexity,
                    "completed_steps": completed_steps,
                    "epoch": epoch,
                },
                os.path.join(args.output_dir, "pytorch_model.bin")
            )
            
            # Save the tokenizer
            tokenizer.save_pretrained(args.output_dir)
            with open(log_file_path, "a") as f:
                train_loss = train_loss / len(train_dataloader)
                total_train_loss = total_train_loss / len(train_dataloader)
                f.write(f"epoch: {epoch}, step: {completed_steps}, train_loss: {train_loss:.6f}, total_train_loss: {total_train_loss:.6f}, eval_loss: {eval_loss:.6f}, perplexity: {perplexity:.6f}\n")                
        if args.checkpointing_steps == "epoch":
            output_dir = f"epoch_{epoch}"
            if args.output_dir is not None:
                output_dir = os.path.join(args.output_dir, output_dir)
            os.makedirs(output_dir, exist_ok=True)
            
            # Save model for the epoch
            torch.save(
                {
                    "model_state_dict": model.state_dict(),
                    "optimizer_state_dict": optimizer.state_dict(),
                    "epoch": epoch,
                    "completed_steps": completed_steps,
                },
                os.path.join(output_dir, "pytorch_model.bin")
            )

    if args.output_dir is not None:
        # Save the final model
        torch.save(
            {
                "model_state_dict": model.state_dict(),
                "config": vars(model.config),
                "vocab_size": model.transformer.wte.weight.shape[0],
                "perplexity": perplexity
            },
            os.path.join(args.output_dir, "pytorch_model.bin")
        )
        
        # Save tokenizer and results
        tokenizer.save_pretrained(args.output_dir)
        with open(os.path.join(args.output_dir, "all_results.json"), "w") as f:
            json.dump({"perplexity": perplexity}, f)
        
        logger.info(f"Model saved to {args.output_dir}")

if __name__ == "__main__":
    main()
