usrbinenv python codingutf8 2021 the huggingface team 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 orglicenseslicense2 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 pretrainingfinetuning the gptneo model for causal language modeling on a text file or a dataset using model parallelism arguments pertaining to which modelconfigtokenizer we are going to finetune or train from scratch arguments pertaining to what data we are going to input our model for training and eval returns batches of size batchsize from truncated dataset sharded over all local devices shuffle batches if shuffle is true returns a linear warmup lineardecay learning rate function stepsperepoch traindssize trainbatchsize numtrainsteps stepsperepoch numtrainepochs warmupfn optax linearscheduleinitvalue0 0 endvaluelearningrate transitionstepsnumwarmupsteps decayfn optax linearschedule initvaluelearningrate endvalue0 transitionstepsnumtrainsteps numwarmupsteps schedulefn optax joinschedulesscheduleswarmupfn decayfn boundariesnumwarmupsteps return schedulefn def main see all possible arguments in srctransformerstrainingargs py or by passing the help flag to this script we now keep distinct sets of args for a cleaner separation of concerns parser hfargumentparsermodelarguments datatrainingarguments trainingarguments if lensys argv 2 and sys argv1 endswith json if we pass only one argument to the script and it s the path to a json file let s parse it to get our arguments modelargs dataargs trainingargs parser parsejsonfilejsonfileos path abspathsys argv1 else modelargs dataargs trainingargs parser parseargsintodataclasses if os path existstrainingargs outputdir and os listdirtrainingargs outputdir and trainingargs dotrain and not trainingargs overwriteoutputdir raise valueerror foutput directory trainingargs outputdir already exists and is not empty use overwriteoutputdir to overcome make one log on every process with the configuration for debugging logging basicconfig formatasctimes levelnames names messages datefmtmdy h m s levellogging info setup logging we only want one process per machine to log things on the screen logger setlevellogging info if jax processindex 0 else logging error if jax processindex 0 datasets utils logging setverbositywarning transformers utils logging setverbosityinfo else datasets utils logging setverbosityerror transformers utils logging setverbosityerror set the verbosity to info of the transformers logger on main process only logger infoftrainingevaluation parameters trainingargs get the datasets you can either provide your own csvjsontxt training and evaluation files see below or just provide the name of one of the public datasets available on the hub at https huggingface codatasets the dataset will be downloaded automatically from the datasets hub for csvjson files this script will use the column called text or the first column if no column called text is found you can easily tweak this behavior see below if dataargs datasetname is not none downloading and loading a dataset from the hub dataset loaddataset dataargs datasetname dataargs datasetconfigname cachedirmodelargs cachedir keepinmemoryfalse if validation not in dataset keys datasetvalidation loaddataset dataargs datasetname dataargs datasetconfigname splitftrain dataargs validationsplitpercentage cachedirmodelargs cachedir datasettrain loaddataset dataargs datasetname dataargs datasetconfigname splitftraindataargs validationsplitpercentage cachedirmodelargs cachedir else datafiles if dataargs trainfile is not none datafilestrain dataargs trainfile if dataargs validationfile is not none datafilesvalidation dataargs validationfile extension dataargs trainfile split 1 if extension txt extension text dataset loaddatasetextension datafilesdatafiles cachedirmodelargs cachedir see more about loading any type of standard or custom dataset from files python dict pandas dataframe etc at https huggingface codocsdatasetsloadingdatasets load pretrained config and tokenizer if modelargs configname config autoconfig frompretrainedmodelargs configname cachedirmodelargs cachedir elif modelargs modelnameorpath config autoconfig frompretrainedmodelargs modelnameorpath cachedirmodelargs cachedir else config configmappingmodelargs modeltype logger warningyou are instantiating a new config instance from scratch if modelargs tokenizername tokenizer autotokenizer frompretrained modelargs tokenizername cachedirmodelargs cachedir usefastmodelargs usefasttokenizer elif modelargs modelnameorpath tokenizer autotokenizer frompretrained modelargs modelnameorpath cachedirmodelargs cachedir usefastmodelargs usefasttokenizer 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 tokenizername if trainingargs dotrain columnnames datasettrain columnnames else columnnames datasetvalidation columnnames textcolumnname text if text in columnnames else columnnames0 since this will be pickled to avoid lazymodule error in hasher force logger loading before tokenizefunction toklogger transformers utils logging getloggertransformers tokenizationutilsbase def tokenizefunctionexamples with captureloggertoklogger as cl output tokenizerexamplestextcolumnname clm input could be much much longer than blocksize if token indices sequence length is longer than the in cl out toklogger warning please ignore the warning above this long input will be chunked into smaller bits before being passed to the model return output tokenizeddatasets dataset map tokenizefunction batchedtrue numprocdataargs preprocessingnumworkers removecolumnscolumnnames loadfromcachefilenot dataargs overwritecache if dataargs blocksize is none blocksize tokenizer modelmaxlength if blocksize config maxpositionembeddings logger warning fthe tokenizer picked seems to have a very large modelmaxlength tokenizer modelmaxlength fusing blocksizemin1024 config maxpositionembeddings instead you can change that default value by passing blocksize xxx blocksize min1024 config maxpositionembeddings else if dataargs blocksize tokenizer modelmaxlength logger warning fthe blocksize passed dataargs blocksize is larger than the maximum length for the model ftokenizer modelmaxlength using blocksizetokenizer modelmaxlength blocksize mindataargs blocksize tokenizer modelmaxlength main data processing function that will concatenate all texts from our dataset and generate chunks of blocksize def grouptextsexamples concatenate all texts concatenatedexamples k listchainexamplesk for k in examples keys totallength lenconcatenatedexampleslistexamples keys0 we drop the small remainder we could add padding if the model supported it instead of this drop you can customize this part to your needs if totallength blocksize totallength totallength blocksize blocksize split by chunks of maxlen result k ti i blocksize for i in range0 totallength blocksize for k t in concatenatedexamples items resultlabels resultinputids copy return result note that with batchedtrue this map processes 1 000 texts together so grouptexts throws away a remainder for each of those groups of 1 000 texts you can adjust that batchsize here but a higher value might be slower to preprocess to speed up this part we use multiprocessing see the documentation of the map method for more information https huggingface codocsdatasetsprocessmap lmdatasets tokenizeddatasets map grouptexts batchedtrue numprocdataargs preprocessingnumworkers loadfromcachefilenot dataargs overwritecache if trainingargs dotrain if train not in tokenizeddatasets raise valueerrordotrain requires a train dataset traindataset lmdatasetstrain if dataargs maxtrainsamples is not none maxtrainsamples minlentraindataset dataargs maxtrainsamples traindataset traindataset selectrangemaxtrainsamples if trainingargs doeval if validation not in tokenizeddatasets raise valueerrordoeval requires a validation dataset evaldataset lmdatasetsvalidation if dataargs maxevalsamples is not none maxevalsamples minlenevaldataset dataargs maxevalsamples evaldataset evaldataset selectrangemaxevalsamples enable tensorboard only on the master node hastensorboard istensorboardavailable if hastensorboard and jax processindex 0 try from flax metrics tensorboard import summarywriter summarywriter summarywriterlogdirpathtrainingargs outputdir except importerror as ie hastensorboard false logger warning funable to display metrics through tensorboard because some package are not installed ie else logger warning unable to display metrics through tensorboard because the package is not installed please run pip install tensorboard to enable initialize our training rng jax random prngkeytrainingargs seed rng dropoutrng jax random splitrng store some constant numepochs inttrainingargs numtrainepochs trainbatchsize inttrainingargs perdevicetrainbatchsize jax devicecount evalbatchsize inttrainingargs perdeviceevalbatchsize jax devicecount stepsperepoch lentraindataset trainbatchsize totaltrainsteps stepsperepoch numepochs todo weights should be initialized in pjitted fun this won t work for really large models todo when loading from pretrained model we need to make sure the vocab is divisible by numpartitions gpt2 s vocab is odd we need to resize it for finetuning model flaxautomodelforcausallm frompretrained modelargs modelnameorpath seedtrainingargs seed dtypegetattrjnp modelargs dtype create learning rate schedule lineardecaylrschedulefn createlearningratefn lentraindataset trainbatchsize trainingargs numtrainepochs trainingargs warmupsteps trainingargs learningrate optimizer optax adamw learningratelineardecaylrschedulefn b1trainingargs adambeta1 b2trainingargs adambeta2 epstrainingargs adamepsilon weightdecaytrainingargs weightdecay def getinitialstateparams state optimizer initparams return tuplestate params get partitionspec for model params paramspec setpartitionsunfreezemodel params get the pytree for optstate we don t actually initialize the optstate yet paramsshapes jax treeutil treemaplambda x x shape model params stateshapes jax evalshapegetinitialstate paramsshapes get partitionspec for optstate this is very specific to adamw todo optax returns different state for different optimizers how can we handle this generically or maybe we don t since in our examples we just use adamw or adafactor def getoptspecx if isinstancex dict return paramspec return none optstatespec paramspec jax treeutil treemap getoptspec stateshapes isleaflambda x isinstancex dict optax emptystate pjit the getinitialstate function to shard params and init optimizer state in sharded way pgetinitialstate pjit getinitialstate inaxisresourcesnone outaxisresourcesoptstatespec paramspec hack move the inital params to cpu to free up device memory todo allow loading weights on cpu in pretrained model model params jax treeutil treemaplambda x np asarrayx model params mesh defination meshdevices np arrayjax devices reshape1 jax localdevicecount actually initialize the optstate with meshmeshdevices dp mp optstate params pgetinitialstatefreezemodel params crossentropy with z loss def lossfnlogits labels zloss0 shiftlogits logits 1 shiftlabels labels 1 shiftlabels onehotshiftlabels shiftlogits shape1 shiftlogits shiftlogits jax lax stopgradientshiftlogits maxaxis1 keepdimstrue logz jnp logjnp sumjnp expshiftlogits axis1 keepdimstrue logsoftmax shiftlogits logz loss jnp sumshiftlabels logsoftmax axis1 loss 1e4 jnp squarelogz squeeze1 zloss return loss mean define gradient update step fn todo try to use trainstate instead of passing params and optstate individually def trainstepparams optstate dropoutrng batch step dropoutrng newdropoutrng jax random splitdropoutrng def computelossparams labels batch poplabels logits modelbatch paramsparams dropoutrngdropoutrng traintrue0 loss lossfnlogits labels zloss1 0 return loss gradfn jax valueandgradcomputeloss loss grads gradfnparams updates newoptstate optimizer updategrads optstate params newparams optax applyupdatesparams updates metrics loss loss learningrate lineardecaylrschedulefnstep return newparams tuplenewoptstate newdropoutrng metrics step 1 define eval fn def evalstepinputids labels params logits modelinputidsinputids paramsparams trainfalse0 loss lossfnlogits labels metrics return loss loss ptrainstep pjit trainstep inaxisresourcesparamspec optstatespec none none none outaxisresourcesparamspec optstatespec none none none donateargnums0 1 pevalstep pjit evalstep inaxisresourcesnone none paramspec outaxisresourcesnone logger info running training logger infof num examples lentraindataset logger infof num epochs numepochs logger infof instantaneous batch size per device trainingargs perdevicetrainbatchsize logger infof total train batch size w parallel distributed trainbatchsize logger infof total optimization steps totaltrainsteps traintime 0 trainmetrics epochs tqdmrangenumepochs descfepoch 1numepochs position0 globalstep 0 we are not doing 2d parallelism yet this just does model parallelism with meshmeshdevices dp mp for in epochs training trainstart time time create sampling rng rng inputrng jax random splitrng generate an epoch by shuffling sampling indices from the train dataset trainmetrics trainloader dataloaderinputrng traindataset trainbatchsize shuffletrue stepsperepoch lentraindataset trainbatchsize train for in tqdmrangestepsperepoch desctraining position1 leavefalse batch nexttrainloader params optstate dropoutrng trainmetric globalstep ptrainstep params optstate dropoutrng batch globalstep trainmetrics appendtrainmetric curstep globalstep if curstep trainingargs loggingsteps 0 and curstep 0 save metrics traintime time time trainstart if hastensorboard and jax processindex 0 writetrainmetricsummarywriter trainmetrics traintime curstep epochs write fstep curstep loss trainmetric loss learning rate f trainmetric learningrate trainmetrics if curstep trainingargs evalsteps 0 and curstep 0 evaluating evalmetrics evalloader dataloaderinputrng evaldataset evalbatchsize evalsteps lenevaldataset evalbatchsize for in tqdmrangeevalsteps descevaluating position2 leavefalse batch nextevalloader metrics pevalstepbatchinputids batchlabels params evalmetrics appendmetrics normalize eval metrics evalmetrics stackforestevalmetrics evalmetrics jax treeutil treemapjnp mean evalmetrics try evalmetricsperplexity math expevalmetricsloss except overflowerror evalmetricsperplexity floatinf logger info fstep curstep eval loss evalmetrics loss eval perplexity f evalmetrics perplexity if curstep trainingargs savesteps 0 and curstep 0 save checkpoint after each epoch and push checkpoint to the hub if jax processindex 0 params jax devicegetparams model savepretrained trainingargs outputdir paramsparams pushtohubtrainingargs pushtohub commitmessagefsaving weights and logs of step curstep if name main main usr bin env python coding utf 8 2021 the huggingface team 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 pre training fine tuning the gptneo model for causal language modeling on a text file or a dataset using model parallelism arguments pertaining to which model config tokenizer we are going to fine tune or train from scratch arguments pertaining to what data we are going to input our model for training and eval returns batches of size batch_size from truncated dataset sharded over all local devices shuffle batches if shuffle is true skip incomplete batch returns a linear warmup linear_decay learning rate function see all possible arguments in src transformers training_args py or by passing the help flag to this script we now keep distinct sets of args for a cleaner separation of concerns if we pass only one argument to the script and it s the path to a json file let s parse it to get our arguments make one log on every process with the configuration for debugging setup logging we only want one process per machine to log things on the screen set the verbosity to info of the transformers logger on main process only get the datasets you can either provide your own csv json txt training and evaluation files see below or just provide the name of one of the public datasets available on the hub at https huggingface co datasets the dataset will be downloaded automatically from the datasets hub for csv json files this script will use the column called text or the first column if no column called text is found you can easily tweak this behavior see below downloading and loading a dataset from the hub see more about loading any type of standard or custom dataset from files python dict pandas dataframe etc at https huggingface co docs datasets loading_datasets load pretrained config and tokenizer since this will be pickled to avoid _lazymodule error in hasher force logger loading before tokenize_function clm input could be much much longer than block_size main data processing function that will concatenate all texts from our dataset and generate chunks of block_size concatenate all texts we drop the small remainder we could add padding if the model supported it instead of this drop you can customize this part to your needs split by chunks of max_len note that with batched true this map processes 1 000 texts together so group_texts throws away a remainder for each of those groups of 1 000 texts you can adjust that batch_size here but a higher value might be slower to preprocess to speed up this part we use multiprocessing see the documentation of the map method for more information https huggingface co docs datasets process map enable tensorboard only on the master node initialize our training store some constant todo weights should be initialized in pjitted fun this won t work for really large models todo when loading from pre trained model we need to make sure the vocab is divisible by num_partitions gpt2 s vocab is odd we need to resize it for fine tuning create learning rate schedule get partitionspec for model params get the pytree for opt_state we don t actually initialize the opt_state yet get partitionspec for opt_state this is very specific to adamw todo optax returns different state for different optimizers how can we handle this generically or maybe we don t since in our examples we just use adamw or adafactor pjit the get_initial_state function to shard params and init optimizer state in sharded way hack move the inital params to cpu to free up device memory todo allow loading weights on cpu in pre trained model mesh defination actually initialize the opt_state cross entropy with z loss define gradient update step fn todo try to use trainstate instead of passing params and opt_state individually define eval fn metrics we are not doing 2d parallelism yet this just does model parallelism training create sampling rng generate an epoch by shuffling sampling indices from the train dataset train save metrics evaluating normalize eval metrics save checkpoint after each epoch and push checkpoint to the hub
import logging
import math
import os
import sys
import time
from dataclasses import dataclass, field
from itertools import chain
from pathlib import Path
from typing import Callable, Optional

import datasets
import jax
import jax.numpy as jnp
import numpy as np
import optax
from datasets import Dataset, load_dataset
from flax.core.frozen_dict import freeze, unfreeze
from flax.training.common_utils import onehot, stack_forest
from jax.experimental.maps import mesh
from jax.experimental.pjit import pjit
from partitions import set_partitions
from tqdm import tqdm

import transformers
from transformers import (
    CONFIG_MAPPING,
    FLAX_MODEL_FOR_CAUSAL_LM_MAPPING,
    AutoConfig,
    AutoTokenizer,
    FlaxAutoModelForCausalLM,
    HfArgumentParser,
    TrainingArguments,
    is_tensorboard_available,
)
from transformers.testing_utils import CaptureLogger


logger = logging.getLogger(__name__)

MODEL_CONFIG_CLASSES = list(FLAX_MODEL_FOR_CAUSAL_LM_MAPPING.keys())
MODEL_TYPES = tuple(conf.model_type for conf in MODEL_CONFIG_CLASSES)


@dataclass
class ModelArguments:
    

    model_name_or_path: Optional[str] = field(
        default=None,
        metadata={
            "help": (
                "The model checkpoint for weights initialization. Don't set if you want to train a model from scratch."
            )
        },
    )
    model_type: Optional[str] = field(
        default=None,
        metadata={"help": "If training from scratch, pass a model type from the list: " + ", ".join(MODEL_TYPES)},
    )
    config_name: Optional[str] = field(
        default=None, metadata={"help": "Pretrained config name or path if not the same as model_name"}
    )
    tokenizer_name: Optional[str] = field(
        default=None, metadata={"help": "Pretrained tokenizer name or path if not the same as model_name"}
    )
    cache_dir: Optional[str] = field(
        default=None, metadata={"help": "Where do you want to store the pretrained models downloaded from s3"}
    )
    use_fast_tokenizer: bool = field(
        default=True,
        metadata={"help": "Whether to use one of the fast tokenizer (backed by the tokenizers library) or not."},
    )
    dtype: Optional[str] = field(
        default="float32",
        metadata={
            "help": (
                "Floating-point format in which the model weights should be initialized and trained. Choose one of"
                " `[float32, float16, bfloat16]`."
            )
        },
    )


@dataclass
class DataTrainingArguments:
    

    dataset_name: Optional[str] = field(
        default=None, metadata={"help": "The name of the dataset to use (via the datasets library)."}
    )
    dataset_config_name: Optional[str] = field(
        default=None, metadata={"help": "The configuration name of the dataset to use (via the datasets library)."}
    )
    train_file: Optional[str] = field(default=None, metadata={"help": "The input training data file (a text file)."})
    validation_file: Optional[str] = field(
        default=None,
        metadata={"help": "An optional input evaluation data file to evaluate the perplexity on (a text file)."},
    )
    max_train_samples: Optional[int] = field(
        default=None,
        metadata={
            "help": (
                "For debugging purposes or quicker training, truncate the number of training examples to this "
                "value if set."
            )
        },
    )
    max_eval_samples: Optional[int] = field(
        default=None,
        metadata={
            "help": (
                "For debugging purposes or quicker training, truncate the number of evaluation examples to this "
                "value if set."
            )
        },
    )
    overwrite_cache: bool = field(
        default=False, metadata={"help": "Overwrite the cached training and evaluation sets"}
    )
    validation_split_percentage: Optional[int] = field(
        default=5,
        metadata={
            "help": "The percentage of the train set used as validation set in case there's no validation split"
        },
    )
    block_size: Optional[int] = field(
        default=None,
        metadata={
            "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)."
            )
        },
    )
    overwrite_cache: bool = field(
        default=False, metadata={"help": "Overwrite the cached training and evaluation sets"}
    )
    preprocessing_num_workers: Optional[int] = field(
        default=None,
        metadata={"help": "The number of processes to use for the preprocessing."},
    )

    def __post_init__(self):
        if self.dataset_name is None and self.train_file is None and self.validation_file is None:
            raise ValueError("Need either a dataset name or a training/validation file.")
        else:
            if self.train_file is not None:
                extension = self.train_file.split(".")[-1]
                assert extension in ["csv", "json", "txt"], "`train_file` should be a csv, a json or a txt file."
            if self.validation_file is not None:
                extension = self.validation_file.split(".")[-1]
                assert extension in ["csv", "json", "txt"], "`validation_file` should be a csv, a json or a txt file."


def data_loader(rng: jax.random.PRNGKey, dataset: Dataset, batch_size: int, shuffle: bool = False):
    
    steps_per_epoch = len(dataset) // batch_size

    if shuffle:
        batch_idx = jax.random.permutation(rng, len(dataset))
    else:
        batch_idx = jnp.arange(len(dataset))

    batch_idx = batch_idx[: steps_per_epoch * batch_size]  
    batch_idx = batch_idx.reshape((steps_per_epoch, batch_size))

    for idx in batch_idx:
        batch = dataset[idx]
        batch = {k: jnp.array(v) for k, v in batch.items()}
        yield batch


def write_train_metric(summary_writer, train_metrics, train_time, step):
    summary_writer.scalar("train_time", train_time, step)

    train_metrics = stack_forest(train_metrics)
    for key, vals in train_metrics.items():
        tag = f"train_{key}"
        for i, val in enumerate(vals):
            summary_writer.scalar(tag, val, step - len(vals) + i + 1)


def write_eval_metric(summary_writer, eval_metrics, step):
    for metric_name, value in eval_metrics.items():
        summary_writer.scalar(f"eval_{metric_name}", value, step)


def create_learning_rate_fn(
    train_ds_size: int, train_batch_size: int, num_train_epochs: int, num_warmup_steps: int, learning_rate: float
) -> Callable[[int], jnp.ndarray]:
    
    steps_per_epoch = train_ds_size // train_batch_size
    num_train_steps = steps_per_epoch * num_train_epochs
    warmup_fn = optax.linear_schedule(init_value=0.0, end_value=learning_rate, transition_steps=num_warmup_steps)
    decay_fn = optax.linear_schedule(
        init_value=learning_rate, end_value=0, transition_steps=num_train_steps - num_warmup_steps
    )
    schedule_fn = optax.join_schedules(schedules=[warmup_fn, decay_fn], boundaries=[num_warmup_steps])
    return schedule_fn


def main():
    
    
    

    parser = HfArgumentParser((ModelArguments, DataTrainingArguments, TrainingArguments))
    if len(sys.argv) == 2 and sys.argv[1].endswith(".json"):
        
        
        model_args, data_args, training_args = parser.parse_json_file(json_file=os.path.abspath(sys.argv[1]))
    else:
        model_args, data_args, training_args = parser.parse_args_into_dataclasses()

    if (
        os.path.exists(training_args.output_dir)
        and os.listdir(training_args.output_dir)
        and training_args.do_train
        and not training_args.overwrite_output_dir
    ):
        raise ValueError(
            f"Output directory ({training_args.output_dir}) already exists and is not empty. "
            "Use --overwrite_output_dir to overcome."
        )

    
    logging.basicConfig(
        format="%(asctime)s - %(levelname)s - %(name)s - %(message)s",
        datefmt="%m/%d/%Y %H:%M:%S",
        level=logging.INFO,
    )
    
    logger.setLevel(logging.INFO if jax.process_index() == 0 else logging.ERROR)
    if jax.process_index() == 0:
        datasets.utils.logging.set_verbosity_warning()
        transformers.utils.logging.set_verbosity_info()
    else:
        datasets.utils.logging.set_verbosity_error()
        transformers.utils.logging.set_verbosity_error()

    
    logger.info(f"Training/evaluation parameters {training_args}")

    
    
    
    
    
    
    if data_args.dataset_name is not None:
        
        dataset = load_dataset(
            data_args.dataset_name, data_args.dataset_config_name, cache_dir=model_args.cache_dir, keep_in_memory=False
        )

        if "validation" not in dataset.keys():
            dataset["validation"] = load_dataset(
                data_args.dataset_name,
                data_args.dataset_config_name,
                split=f"train[:{data_args.validation_split_percentage}%]",
                cache_dir=model_args.cache_dir,
            )
            dataset["train"] = load_dataset(
                data_args.dataset_name,
                data_args.dataset_config_name,
                split=f"train[{data_args.validation_split_percentage}%:]",
                cache_dir=model_args.cache_dir,
            )
    else:
        data_files = {}
        if data_args.train_file is not None:
            data_files["train"] = data_args.train_file
        if data_args.validation_file is not None:
            data_files["validation"] = data_args.validation_file
        extension = data_args.train_file.split(".")[-1]
        if extension == "txt":
            extension = "text"
        dataset = load_dataset(extension, data_files=data_files, cache_dir=model_args.cache_dir)
    
    

    
    if model_args.config_name:
        config = AutoConfig.from_pretrained(model_args.config_name, cache_dir=model_args.cache_dir)
    elif model_args.model_name_or_path:
        config = AutoConfig.from_pretrained(model_args.model_name_or_path, cache_dir=model_args.cache_dir)
    else:
        config = CONFIG_MAPPING[model_args.model_type]()
        logger.warning("You are instantiating a new config instance from scratch.")

    if model_args.tokenizer_name:
        tokenizer = AutoTokenizer.from_pretrained(
            model_args.tokenizer_name, cache_dir=model_args.cache_dir, use_fast=model_args.use_fast_tokenizer
        )
    elif model_args.model_name_or_path:
        tokenizer = AutoTokenizer.from_pretrained(
            model_args.model_name_or_path, cache_dir=model_args.cache_dir, use_fast=model_args.use_fast_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."
        )

    if training_args.do_train:
        column_names = dataset["train"].column_names
    else:
        column_names = dataset["validation"].column_names
    text_column_name = "text" if "text" in column_names else column_names[0]

    
    tok_logger = transformers.utils.logging.get_logger("transformers.tokenization_utils_base")

    def tokenize_function(examples):
        with CaptureLogger(tok_logger) as cl:
            output = tokenizer(examples[text_column_name])
        
        if "Token indices sequence length is longer than the" in cl.out:
            tok_logger.warning(
                "^^^^^^^^^^^^^^^^ Please ignore the warning above - this long input will be chunked into smaller bits"
                " before being passed to the model."
            )
        return output

    tokenized_datasets = dataset.map(
        tokenize_function,
        batched=True,
        num_proc=data_args.preprocessing_num_workers,
        remove_columns=column_names,
        load_from_cache_file=not data_args.overwrite_cache,
    )

    if data_args.block_size is None:
        block_size = tokenizer.model_max_length
        if block_size > config.max_position_embeddings:
            logger.warning(
                f"The tokenizer picked seems to have a very large `model_max_length` ({tokenizer.model_max_length}). "
                f"Using block_size={min(1024, config.max_position_embeddings)} instead. You can change that default value by passing --block_size xxx."
            )
            block_size = min(1024, config.max_position_embeddings)
    else:
        if data_args.block_size > tokenizer.model_max_length:
            logger.warning(
                f"The block_size passed ({data_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(data_args.block_size, tokenizer.model_max_length)

    
    def group_texts(examples):
        
        concatenated_examples = {k: list(chain(*examples[k])) for k in examples.keys()}
        total_length = len(concatenated_examples[list(examples.keys())[0]])
        
        
        if total_length >= block_size:
            total_length = (total_length // block_size) * block_size
        
        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=data_args.preprocessing_num_workers,
        load_from_cache_file=not data_args.overwrite_cache,
    )

    if training_args.do_train:
        if "train" not in tokenized_datasets:
            raise ValueError("--do_train requires a train dataset")
        train_dataset = lm_datasets["train"]
        if data_args.max_train_samples is not None:
            max_train_samples = min(len(train_dataset), data_args.max_train_samples)
            train_dataset = train_dataset.select(range(max_train_samples))

    if training_args.do_eval:
        if "validation" not in tokenized_datasets:
            raise ValueError("--do_eval requires a validation dataset")
        eval_dataset = lm_datasets["validation"]
        if data_args.max_eval_samples is not None:
            max_eval_samples = min(len(eval_dataset), data_args.max_eval_samples)
            eval_dataset = eval_dataset.select(range(max_eval_samples))

    
    has_tensorboard = is_tensorboard_available()
    if has_tensorboard and jax.process_index() == 0:
        try:
            from flax.metrics.tensorboard import SummaryWriter

            summary_writer = SummaryWriter(log_dir=Path(training_args.output_dir))
        except ImportError as ie:
            has_tensorboard = False
            logger.warning(
                f"Unable to display metrics through TensorBoard because some package are not installed: {ie}"
            )
    else:
        logger.warning(
            "Unable to display metrics through TensorBoard because the package is not installed: "
            "Please run pip install tensorboard to enable."
        )

    
    rng = jax.random.PRNGKey(training_args.seed)
    rng, dropout_rng = jax.random.split(rng)

    
    num_epochs = int(training_args.num_train_epochs)
    train_batch_size = int(training_args.per_device_train_batch_size) * jax.device_count()
    eval_batch_size = int(training_args.per_device_eval_batch_size) * jax.device_count()
    steps_per_epoch = len(train_dataset) // train_batch_size
    total_train_steps = steps_per_epoch * num_epochs

    
    
    
    model = FlaxAutoModelForCausalLM.from_pretrained(
        model_args.model_name_or_path, seed=training_args.seed, dtype=getattr(jnp, model_args.dtype)
    )

    
    linear_decay_lr_schedule_fn = create_learning_rate_fn(
        len(train_dataset),
        train_batch_size,
        training_args.num_train_epochs,
        training_args.warmup_steps,
        training_args.learning_rate,
    )

    optimizer = optax.adamw(
        learning_rate=linear_decay_lr_schedule_fn,
        b1=training_args.adam_beta1,
        b2=training_args.adam_beta2,
        eps=training_args.adam_epsilon,
        weight_decay=training_args.weight_decay,
    )

    def get_initial_state(params):
        state = optimizer.init(params)
        return tuple(state), params

    
    param_spec = set_partitions(unfreeze(model.params))

    
    params_shapes = jax.tree_util.tree_map(lambda x: x.shape, model.params)
    state_shapes = jax.eval_shape(get_initial_state, params_shapes)

    
    
    
    def get_opt_spec(x):
        if isinstance(x, dict):
            return param_spec
        return None

    opt_state_spec, param_spec = jax.tree_util.tree_map(
        get_opt_spec, state_shapes, is_leaf=lambda x: isinstance(x, (dict, optax.EmptyState))
    )

    
    
    p_get_initial_state = pjit(
        get_initial_state,
        in_axis_resources=None,
        out_axis_resources=(opt_state_spec, param_spec),
    )

    
    
    model.params = jax.tree_util.tree_map(lambda x: np.asarray(x), model.params)

    
    mesh_devices = np.array(jax.devices()).reshape(1, jax.local_device_count())

    
    with mesh(mesh_devices, ("dp", "mp")):
        opt_state, params = p_get_initial_state(freeze(model.params))

    
    def loss_fn(logits, labels, z_loss=0):
        shift_logits = logits[..., :-1, :]
        shift_labels = labels[..., 1:]

        shift_labels = onehot(shift_labels, shift_logits.shape[-1])

        shift_logits = shift_logits - jax.lax.stop_gradient(shift_logits.max(axis=-1, keepdims=True))
        log_z = jnp.log(jnp.sum(jnp.exp(shift_logits), axis=-1, keepdims=True))
        log_softmax = shift_logits - log_z
        loss = -jnp.sum(shift_labels * log_softmax, axis=-1)

        loss += (1e-4 * jnp.square(log_z.squeeze(-1))) * z_loss

        return loss.mean()

    
    
    def train_step(params, opt_state, dropout_rng, batch, step):
        dropout_rng, new_dropout_rng = jax.random.split(dropout_rng)

        def compute_loss(params):
            labels = batch.pop("labels")
            logits = model(**batch, params=params, dropout_rng=dropout_rng, train=True)[0]
            loss = loss_fn(logits, labels, z_loss=1.0)
            return loss

        grad_fn = jax.value_and_grad(compute_loss)
        loss, grads = grad_fn(params)

        updates, new_opt_state = optimizer.update(grads, opt_state, params)
        new_params = optax.apply_updates(params, updates)

        metrics = {"loss": loss, "learning_rate": linear_decay_lr_schedule_fn(step)}
        return new_params, tuple(new_opt_state), new_dropout_rng, metrics, step + 1

    
    def eval_step(input_ids, labels, params):
        logits = model(input_ids=input_ids, params=params, train=False)[0]
        loss = loss_fn(logits, labels)
        
        return {"loss": loss}

    p_train_step = pjit(
        train_step,
        in_axis_resources=(param_spec, opt_state_spec, None, None, None),
        out_axis_resources=(param_spec, opt_state_spec, None, None, None),
        donate_argnums=(0, 1),
    )

    p_eval_step = pjit(
        eval_step,
        in_axis_resources=(None, None, param_spec),
        out_axis_resources=None,
    )

    logger.info("***** Running training *****")
    logger.info(f"  Num examples = {len(train_dataset)}")
    logger.info(f"  Num Epochs = {num_epochs}")
    logger.info(f"  Instantaneous batch size per device = {training_args.per_device_train_batch_size}")
    logger.info(f"  Total train batch size (w. parallel & distributed) = {train_batch_size}")
    logger.info(f"  Total optimization steps = {total_train_steps}")

    train_time = 0
    train_metrics = []
    epochs = tqdm(range(num_epochs), desc=f"Epoch ... (1/{num_epochs})", position=0)
    global_step = 0
    
    with mesh(mesh_devices, ("dp", "mp")):
        for _ in epochs:
            
            train_start = time.time()

            
            rng, input_rng = jax.random.split(rng)

            
            train_metrics = []
            train_loader = data_loader(input_rng, train_dataset, train_batch_size, shuffle=True)
            steps_per_epoch = len(train_dataset) // train_batch_size

            
            for _ in tqdm(range(steps_per_epoch), desc="Training...", position=1, leave=False):
                batch = next(train_loader)
                params, opt_state, dropout_rng, train_metric, global_step = p_train_step(
                    params,
                    opt_state,
                    dropout_rng,
                    batch,
                    global_step,
                )
                train_metrics.append(train_metric)

                cur_step = global_step

                if cur_step % training_args.logging_steps == 0 and cur_step > 0:
                    
                    train_time += time.time() - train_start
                    if has_tensorboard and jax.process_index() == 0:
                        write_train_metric(summary_writer, train_metrics, train_time, cur_step)

                    epochs.write(
                        f"Step... ({cur_step} | Loss: {train_metric['loss']}, Learning Rate:"
                        f" {train_metric['learning_rate']})"
                    )

                    train_metrics = []

                if cur_step % training_args.eval_steps == 0 and cur_step > 0:
                    
                    eval_metrics = []
                    eval_loader = data_loader(input_rng, eval_dataset, eval_batch_size)
                    eval_steps = len(eval_dataset) // eval_batch_size

                    for _ in tqdm(range(eval_steps), desc="Evaluating...", position=2, leave=False):
                        batch = next(eval_loader)
                        metrics = p_eval_step(batch["input_ids"], batch["labels"], params)
                        eval_metrics.append(metrics)

                    
                    eval_metrics = stack_forest(eval_metrics)
                    eval_metrics = jax.tree_util.tree_map(jnp.mean, eval_metrics)

                    try:
                        eval_metrics["perplexity"] = math.exp(eval_metrics["loss"])
                    except OverflowError:
                        eval_metrics["perplexity"] = float("inf")

                    logger.info(
                        f"Step... ({cur_step} | Eval loss: {eval_metrics['loss']} | Eval Perplexity:"
                        f" {eval_metrics['perplexity']}"
                    )

                if cur_step % training_args.save_steps == 0 and cur_step > 0:
                    
                    if jax.process_index() == 0:
                        params = jax.device_get(params)
                        model.save_pretrained(
                            training_args.output_dir,
                            params=params,
                            push_to_hub=training_args.push_to_hub,
                            commit_message=f"Saving weights and logs of step {cur_step}",
                        )


if __name__ == "__main__":
    main()