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 pretraining the library models for denoising language modeling on a text file or a dataset here is the full list of checkpoints on the hub that can be pretrained by this script https huggingface comodels filterbart you can also adapt this script on your own denoising language modeling task pointers for this are left as comments serializes this instance while replace enum by their values for json serialization support it obfuscates the token values by removing their value 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 data collator used for bart denoising language modeling the code is largely copied from https github commorganmcg1rotobartblobmaindatacollator pyl223 for more information on how bart denoising language modeling works one can take a look at the official paper https arxiv orgpdf1910 13461 pdf or the official code for preprocessing https github comfacebookresearchfairseqblobmainfairseqdatadenoisingdataset py args tokenizer class transformers pretrainedtokenizer or class transformers pretrainedtokenizerfast the tokenizer used for encoding the data maskratio obj float the probability with which to randomly mask tokens in the input poissonlambda obj float mean parameter of poisson distribution used to generate spanlengths to be masked permutesentenceratio obj float ratio of sentences to be permuted in each document decoderstarttokenid obj int the decoder start token id of the model convert list to dict and tensorize input permuting sentences masking span of tokens text infilling in the paper ignore pad tokens shuffle sentences in each document find end locations of sentences write shuffled sentences into results sampling text spans with span lengths drawn from a poisson distribution and masking them determine how many tokens we need to mask in total generate a sufficient number of span lengths remove all spans of length 0 note that bart inserts additional mask tokens where length 0 which we do not implement for now as it adds additional complexity trim to about numtokenstomask tokens randomly choose starting positions for masking prepare mask mask starting positions fill up spans place the mask tokens remove mask tokens that are not starts of spans generate batches of data for a specified batch size from sample indices if the dataset size is not divisible by the batch size and droplast is true the last incomplete batch is dropped else it is returned numsamples lensamplesidx if droplast samplestoremove numsamples batchsize if samplestoremove 0 samplesidx samplesidx samplestoremove sectionssplit numsamples batchsize samplesidx samplesidx reshapesectionssplit batchsize else sectionssplit math ceilnumsamples batchsize samplesidx np arraysplitsamplesidx sectionssplit return samplesidx def writetrainmetricsummarywriter trainmetrics traintime step summarywriter scalartraintime traintime step trainmetrics getmetricstrainmetrics for key vals in trainmetrics items tag ftrainkey for i val in enumeratevals summarywriter scalartag val step lenvals i 1 def writeevalmetricsummarywriter evalmetrics step for metricname value in evalmetrics items summarywriter scalarfevalmetricname value step 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 modelargs useauthtoken is not none warnings warn the useauthtoken argument is deprecated and will be removed in v4 34 please use token instead futurewarning if modelargs token is not none raise valueerrortoken and useauthtoken are both specified please set only the argument token modelargs token modelargs useauthtoken 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 pythonpytorch versions sendexampletelemetryrunbartdlm modelargs dataargs frameworkflax 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 setup logging logging basicconfig formatasctimes levelnames names messages levellogging info datefmtx log on each process the small summary logger logging getloggername set the verbosity to info of the transformers logger on main process only logger infoftrainingevaluation parameters trainingargs set seed before initializing model setseedtrainingargs seed handle the repository creation if trainingargs pushtohub retrieve of infer reponame reponame trainingargs hubmodelid if reponame is none reponame pathtrainingargs outputdir absolute name create repo and retrieve repoid repoid createreporeponame existoktrue tokentrainingargs hubtoken repoid clone repo locally repo repositorytrainingargs outputdir clonefromrepoid tokentrainingargs hubtoken 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 datasets loaddataset dataargs datasetname dataargs datasetconfigname cachedirmodelargs cachedir tokenmodelargs token numprocdataargs preprocessingnumworkers if validation not in datasets keys datasetsvalidation loaddataset dataargs datasetname dataargs datasetconfigname splitftrain dataargs validationsplitpercentage cachedirmodelargs cachedir tokenmodelargs token numprocdataargs preprocessingnumworkers datasetstrain loaddataset dataargs datasetname dataargs datasetconfigname splitftraindataargs validationsplitpercentage cachedirmodelargs cachedir tokenmodelargs token numprocdataargs preprocessingnumworkers 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 datasets loaddataset extension datafilesdatafiles cachedirmodelargs cachedir tokenmodelargs token numprocdataargs preprocessingnumworkers if validation not in datasets keys datasetsvalidation loaddataset extension datafilesdatafiles splitftrain dataargs validationsplitpercentage cachedirmodelargs cachedir tokenmodelargs token numprocdataargs preprocessingnumworkers datasetstrain loaddataset extension datafilesdatafiles splitftraindataargs validationsplitpercentage cachedirmodelargs cachedir tokenmodelargs token numprocdataargs preprocessingnumworkers see more about loading any type of standard or custom dataset from files python dict pandas dataframe etc at https huggingface codocsdatasetsloadingdatasets load pretrained model and tokenizer if modelargs tokenizername tokenizer autotokenizer frompretrained modelargs tokenizername cachedirmodelargs cachedir usefastmodelargs usefasttokenizer tokenmodelargs token elif modelargs modelnameorpath tokenizer autotokenizer frompretrained modelargs modelnameorpath cachedirmodelargs cachedir usefastmodelargs usefasttokenizer tokenmodelargs token 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 modelargs configname config bartconfig frompretrained modelargs configname cachedirmodelargs cachedir vocabsizelentokenizer tokenmodelargs token elif modelargs modelnameorpath config bartconfig frompretrained modelargs modelnameorpath cachedirmodelargs cachedir tokenmodelargs token else config configmappingmodelargs modeltype logger warningyou are instantiating a new config instance from scratch preprocessing the datasets first we tokenize all the texts if trainingargs dotrain columnnames datasetstrain columnnames else columnnames datasetsvalidation columnnames textcolumnname text if text in columnnames else columnnames0 maxseqlength mindataargs maxseqlength tokenizer modelmaxlength use punkt sentence tokenizer to divide a document into a list of sentences nltk downloadpunkt sentencetokenizer nltk data loadtokenizerspunktenglish pickle def sentencesplitfunctionexample sents sentencetokenizer tokenizeexampletext use pad token as end of sentence indicator newtext tokenizer bostoken ftokenizer padtoken joinsents tokenizer eostoken return text newtext splitdatasets datasets map sentencesplitfunction batchedfalse numprocdataargs preprocessingnumworkers removecolumnscolumnnames loadfromcachefilenot dataargs overwritecache tokenize every text then concatenate them together before splitting them in smaller parts since we make sure that all sequences are of the same length no attentionmask is needed def tokenizefunctionexamples return tokenizerexamplestextcolumnname addspecialtokensfalse returnattentionmaskfalse tokenizeddatasets splitdatasets map tokenizefunction batchedtrue numprocdataargs preprocessingnumworkers removecolumnstextcolumnname loadfromcachefilenot dataargs overwritecache main data processing function that will concatenate all texts from our dataset and generate chunks of maxseqlength 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 maxseqlength totallength totallength maxseqlength maxseqlength split by chunks of maxlen result k ti i maxseqlength for i in range0 totallength maxseqlength for k t in concatenatedexamples items 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 tokenizeddatasets tokenizeddatasets map grouptexts batchedtrue numprocdataargs preprocessingnumworkers loadfromcachefilenot dataargs overwritecache 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 dropoutrngs jax random splitrng jax localdevicecount if modelargs modelnameorpath model flaxbartforconditionalgeneration frompretrained modelargs modelnameorpath configconfig seedtrainingargs seed dtypegetattrjnp modelargs dtype tokenmodelargs token else config vocabsize lentokenizer model flaxbartforconditionalgeneration config seedtrainingargs seed dtypegetattrjnp modelargs dtype data collator this one will take care of randomly masking the tokens and permuting the sentences datacollator flaxdatacollatorforbartdenoisinglm tokenizertokenizer decoderstarttokenidmodel config decoderstarttokenid maskratiodataargs mlmprobability poissonlambdadataargs poissonlambda permutesentenceratiodataargs permutesentenceratio store some constant numepochs inttrainingargs numtrainepochs trainbatchsize inttrainingargs perdevicetrainbatchsize jax devicecount perdeviceevalbatchsize inttrainingargs perdeviceevalbatchsize evalbatchsize perdeviceevalbatchsize jax devicecount numtrainsteps lentokenizeddatasetstrain trainbatchsize numepochs create learning rate schedule warmupfn optax linearschedule initvalue0 0 endvaluetrainingargs learningrate transitionstepstrainingargs warmupsteps decayfn optax linearschedule initvaluetrainingargs learningrate endvalue0 transitionstepsnumtrainsteps trainingargs warmupsteps lineardecaylrschedulefn optax joinschedules scheduleswarmupfn decayfn boundariestrainingargs warmupsteps we use optax s masking functionality to not apply weight decay to bias and layernorm scale parameters decaymaskfn returns a mask boolean with the same structure as the parameters the mask is true for parameters that should be decayed def decaymaskfnparams flatparams traverseutil flattendictparams find out all layernorm parameters layernormcandidates layernorm layernorm ln layernormnamedparams layer2 for layernormname in layernormcandidates for layer in flatparams keys if layernormname in joinlayer lower flatmask path path1 bias and path2 not in layernormnamedparams for path in flatparams return traverseutil unflattendictflatmask create adam optimizer if trainingargs adafactor we use the default parameters here to initialize adafactor for more details about the parameters please check https github comdeepmindoptaxblobed02befef9bf81cbbf236be3d2b0e032e9ed4a40optaxsrcalias pyl74 optimizer optax adafactor learningratelineardecaylrschedulefn else optimizer optax adamw learningratelineardecaylrschedulefn b1trainingargs adambeta1 b2trainingargs adambeta2 weightdecaytrainingargs weightdecay maskdecaymaskfn setup train state state trainstate trainstate createapplyfnmodel call paramsmodel params txoptimizer define gradient update step fn def trainstepstate batch dropoutrng dropoutrng newdropoutrng jax random splitdropoutrng def lossfnparams labels batch poplabels logits state applyfnbatch paramsparams dropoutrngdropoutrng traintrue0 compute loss ignore padded input tokens and special tokens labelmask jnp wherelabels 0 1 0 0 0 loss optax softmaxcrossentropylogits onehotlabels logits shape1 labelmask take average loss loss sum numlabels labelmask sum return loss numlabels gradfn jax valueandgradlossfn hasauxtrue loss numlabels grad gradfnstate params numlabels jax lax psumnumlabels batch true loss total loss total samples loss jax lax psumloss batch loss jax treeutil treemaplambda x x numlabels loss true grad total grad total samples grad jax lax psumgrad batch grad jax treeutil treemaplambda x x numlabels grad newstate state applygradientsgradsgrad metrics loss loss learningrate lineardecaylrschedulefnstate step return newstate metrics newdropoutrng create parallel version of the train step ptrainstep jax pmaptrainstep batch donateargnums0 define eval fn def evalstepparams batch labels batch poplabels logits modelbatch paramsparams trainfalse0 compute loss ignore padded input tokens and special tokens labelmask jnp wherelabels 0 1 0 0 0 loss optax softmaxcrossentropylogits onehotlabels logits shape1 labelmask compute accuracy accuracy jnp equaljnp argmaxlogits axis1 labels labelmask summarize metrics metrics loss loss sum accuracy accuracy sum normalizer labelmask sum metrics jax lax psummetrics axisnamebatch return metrics pevalstep jax pmapevalstep batch donateargnums0 replicate the train state on each device state jaxutils replicatestate traintime 0 epochs tqdmrangenumepochs descepoch position0 for epoch in epochs training trainstart time time trainmetrics create sampling rng rng inputrng jax random splitrng generate an epoch by shuffling sampling indices from the train dataset numtrainsamples lentokenizeddatasetstrain avoid using jax numpy here in case of tpu training trainsamplesidx np random permutationnp arangenumtrainsamples trainbatchidx generatebatchsplitstrainsamplesidx trainbatchsize gather the indexes for creating the batch and do a training step for step batchidx in enumeratetqdmtrainbatchidx desctraining position1 samples tokenizeddatasetstrainintidx for idx in batchidx modelinputs datacollatorsamples model forward modelinputs shardmodelinputs data state trainmetric dropoutrngs ptrainstepstate modelinputs dropoutrngs trainmetrics appendtrainmetric curstep epoch numtrainsamples trainbatchsize step if curstep trainingargs loggingsteps 0 and curstep 0 save metrics trainmetric jaxutils unreplicatetrainmetric 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 numevalsamples lentokenizeddatasetsvalidation avoid using jax numpy here in case of tpu training evalsamplesidx np arangenumevalsamples evalbatchidx generatebatchsplitsevalsamplesidx evalbatchsize evalmetrics for i batchidx in enumeratetqdmevalbatchidx descevaluating position2 samples tokenizeddatasetsvalidationintidx for idx in batchidx modelinputs datacollatorsamples model forward metrics padshardunpadpevalstep staticreturntrue state params modelinputs data mindevicebatchperdeviceevalbatchsize evalmetrics appendmetrics normalize eval metrics evalmetrics getmetricsevalmetrics evalmetrics jax treeutil treemapjnp sum evalmetrics evalnormalizer evalmetrics popnormalizer evalmetrics jax treeutil treemaplambda x x evalnormalizer evalmetrics update progress bar epochs desc fstep curstep loss evalmetrics loss acc evalmetrics accuracy save metrics if hastensorboard and jax processindex 0 writeevalmetricsummarywriter evalmetrics curstep 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 devicegetjax treeutil treemaplambda x x0 state params model savepretrainedtrainingargs outputdir paramsparams tokenizer savepretrainedtrainingargs outputdir if trainingargs pushtohub repo pushtohubcommitmessagefsaving weights and logs of step curstep blockingfalse eval after training if trainingargs doeval numevalsamples lentokenizeddatasetsvalidation avoid using jax numpy here in case of tpu training evalsamplesidx np arangenumevalsamples evalbatchidx generatebatchsplitsevalsamplesidx evalbatchsize evalmetrics for batchidx in enumeratetqdmevalbatchidx descevaluating position2 samples tokenizeddatasetsvalidationintidx for idx in batchidx modelinputs datacollatorsamples model forward metrics padshardunpadpevalstep staticreturntrue state params modelinputs data mindevicebatchperdeviceevalbatchsize evalmetrics appendmetrics normalize eval metrics evalmetrics getmetricsevalmetrics evalmetrics jax treeutil treemaplambda metric jnp summetric item evalmetrics evalnormalizer evalmetrics popnormalizer evalmetrics jax treeutil treemaplambda x x evalnormalizer evalmetrics try perplexity math expevalmetricsloss except overflowerror perplexity floatinf evalmetricsperplexity perplexity if jax processindex 0 evalmetrics fevalmetricname value for metricname value in evalmetrics items path os path jointrainingargs outputdir evalresults json with openpath w as f json dumpevalmetrics f indent4 sortkeystrue 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 pretraining the library models for denoising language modeling on a text file or a dataset here is the full list of checkpoints on the hub that can be pretrained by this script https huggingface co models filter bart you can also adapt this script on your own denoising language modeling task pointers for this are left as comments serializes this instance while replace enum by their values for json serialization support it obfuscates the token values by removing their value 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 data collator used for bart denoising language modeling the code is largely copied from https github com morganmcg1 rotobart blob main data_collator py l223 __ for more information on how bart denoising language modeling works one can take a look at the official paper https arxiv org pdf 1910 13461 pdf __ or the official code for preprocessing https github com facebookresearch fairseq blob main fairseq data denoising_dataset py __ args tokenizer class transformers pretrainedtokenizer or class transformers pretrainedtokenizerfast the tokenizer used for encoding the data mask_ratio obj float the probability with which to randomly mask tokens in the input poisson_lambda obj float mean parameter of poisson distribution used to generate span lengths to be masked permute_sentence_ratio obj float ratio of sentences to be permuted in each document decoder_start_token_id obj int the decoder start token id of the model convert list to dict and tensorize input permuting sentences masking span of tokens text infilling in the paper ignore pad tokens shuffle sentences in each document find end locations of sentences write shuffled sentences into results sampling text spans with span lengths drawn from a poisson distribution and masking them determine how many tokens we need to mask in total generate a sufficient number of span lengths remove all spans of length 0 note that bart inserts additional mask tokens where length 0 which we do not implement for now as it adds additional complexity trim to about num_tokens_to_mask tokens randomly choose starting positions for masking prepare mask mask starting positions fill up spans place the mask tokens remove mask tokens that are not starts of spans generate batches of data for a specified batch size from sample indices if the dataset size is not divisible by the batch size and drop_last is true the last incomplete batch is dropped else it is returned 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 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 setup logging log on each process the small summary set the verbosity to info of the transformers logger on main process only set seed before initializing model handle the repository creation retrieve of infer repo_name create repo and retrieve repo_id clone repo locally 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 model and tokenizer preprocessing the datasets first we tokenize all the texts use punkt sentence tokenizer to divide a document into a list of sentences use pad token as end of sentence indicator tokenize every text then concatenate them together before splitting them in smaller parts since we make sure that all sequences are of the same length no attention_mask is needed main data processing function that will concatenate all texts from our dataset and generate chunks of max_seq_length 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 data collator this one will take care of randomly masking the tokens and permuting the sentences store some constant create learning rate schedule we use optax s masking functionality to not apply weight decay to bias and layernorm scale parameters decay_mask_fn returns a mask boolean with the same structure as the parameters the mask is true for parameters that should be decayed find out all layernorm parameters create adam optimizer we use the default parameters here to initialize adafactor for more details about the parameters please check https github com deepmind optax blob ed02befef9bf81cbbf236be3d2b0e032e9ed4a40 optax _src alias py l74 setup train state define gradient update step fn compute loss ignore padded input tokens and special tokens take average true loss total loss total samples true grad total grad total samples create parallel version of the train step define eval fn compute loss ignore padded input tokens and special tokens compute accuracy summarize metrics replicate the train state on each device training create sampling rng generate an epoch by shuffling sampling indices from the train dataset avoid using jax numpy here in case of tpu training gather the indexes for creating the batch and do a training step model forward save metrics evaluating avoid using jax numpy here in case of tpu training model forward normalize eval metrics update progress bar save metrics save checkpoint after each epoch and push checkpoint to the hub eval after training avoid using jax numpy here in case of tpu training model forward normalize eval metrics
import json
import logging
import math
import os
import sys
import time
import warnings
from dataclasses import asdict, dataclass, field
from enum import Enum
from itertools import chain
from pathlib import Path
from typing import Dict, List, Optional

import flax
import jax
import jax.numpy as jnp
import nltk
import numpy as np
import optax
from datasets import load_dataset
from flax import jax_utils, traverse_util
from flax.jax_utils import pad_shard_unpad
from flax.training import train_state
from flax.training.common_utils import get_metrics, onehot, shard
from huggingface_hub import Repository, create_repo
from tqdm import tqdm

from transformers import (
    CONFIG_MAPPING,
    FLAX_MODEL_FOR_MASKED_LM_MAPPING,
    AutoTokenizer,
    BartConfig,
    BatchEncoding,
    FlaxBartForConditionalGeneration,
    HfArgumentParser,
    PreTrainedTokenizerBase,
    is_tensorboard_available,
    set_seed,
)
from transformers.models.bart.modeling_flax_bart import shift_tokens_right
from transformers.utils import send_example_telemetry


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


@dataclass
class TrainingArguments:
    output_dir: str = field(
        metadata={"help": "The output directory where the model predictions and checkpoints will be written."},
    )
    overwrite_output_dir: bool = field(
        default=False,
        metadata={
            "help": (
                "Overwrite the content of the output directory. "
                "Use this to continue training if output_dir points to a checkpoint directory."
            )
        },
    )
    do_train: bool = field(default=False, metadata={"help": "Whether to run training."})
    do_eval: bool = field(default=False, metadata={"help": "Whether to run eval on the dev set."})
    per_device_train_batch_size: int = field(
        default=8, metadata={"help": "Batch size per GPU/TPU core/CPU for training."}
    )
    per_device_eval_batch_size: int = field(
        default=8, metadata={"help": "Batch size per GPU/TPU core/CPU for evaluation."}
    )
    learning_rate: float = field(default=5e-5, metadata={"help": "The initial learning rate for AdamW."})
    weight_decay: float = field(default=0.0, metadata={"help": "Weight decay for AdamW if we apply some."})
    adam_beta1: float = field(default=0.9, metadata={"help": "Beta1 for AdamW optimizer"})
    adam_beta2: float = field(default=0.999, metadata={"help": "Beta2 for AdamW optimizer"})
    adam_epsilon: float = field(default=1e-8, metadata={"help": "Epsilon for AdamW optimizer."})
    adafactor: bool = field(default=False, metadata={"help": "Whether or not to replace AdamW by Adafactor."})
    num_train_epochs: float = field(default=3.0, metadata={"help": "Total number of training epochs to perform."})
    warmup_steps: int = field(default=0, metadata={"help": "Linear warmup over warmup_steps."})
    logging_steps: int = field(default=500, metadata={"help": "Log every X updates steps."})
    save_steps: int = field(default=500, metadata={"help": "Save checkpoint every X updates steps."})
    eval_steps: int = field(default=None, metadata={"help": "Run an evaluation every X steps."})
    seed: int = field(default=42, metadata={"help": "Random seed that will be set at the beginning of training."})
    push_to_hub: bool = field(
        default=False, metadata={"help": "Whether or not to upload the trained model to the model hub after training."}
    )
    hub_model_id: str = field(
        default=None, metadata={"help": "The name of the repository to keep in sync with the local `output_dir`."}
    )
    hub_token: str = field(default=None, metadata={"help": "The token to use to push to the Model Hub."})

    def __post_init__(self):
        if self.output_dir is not None:
            self.output_dir = os.path.expanduser(self.output_dir)

    def to_dict(self):
        
        d = asdict(self)
        for k, v in d.items():
            if isinstance(v, Enum):
                d[k] = v.value
            if isinstance(v, list) and len(v) > 0 and isinstance(v[0], Enum):
                d[k] = [x.value for x in v]
            if k.endswith("_token"):
                d[k] = f"<{k.upper()}>"
        return d


@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]`."
            )
        },
    )
    token: str = field(
        default=None,
        metadata={
            "help": (
                "The token to use as HTTP bearer authorization for remote files. If not specified, will use the token "
                "generated when running `huggingface-cli login` (stored in `~/.huggingface`)."
            )
        },
    )
    use_auth_token: bool = field(
        default=None,
        metadata={
            "help": "The `use_auth_token` argument is deprecated and will be removed in v4.34. Please use `token` instead."
        },
    )


@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)."},
    )
    train_ref_file: Optional[str] = field(
        default=None,
        metadata={"help": "An optional input train ref data file for whole word masking in Chinese."},
    )
    validation_ref_file: Optional[str] = field(
        default=None,
        metadata={"help": "An optional input validation ref data file for whole word masking in Chinese."},
    )
    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"
        },
    )
    max_seq_length: Optional[int] = field(
        default=None,
        metadata={
            "help": (
                "The maximum total input sequence length after tokenization and masking. Sequences longer than this"
                " will be truncated. Default to the max input length of the model."
            )
        },
    )
    preprocessing_num_workers: Optional[int] = field(
        default=None,
        metadata={"help": "The number of processes to use for the preprocessing."},
    )
    mlm_probability: float = field(
        default=0.3, metadata={"help": "Ratio of tokens to mask for span masked language modeling loss"}
    )
    permute_sentence_ratio: float = field(
        default=1.0, metadata={"help": "Ratio of sentences to be permuted in each document"}
    )
    poisson_lambda: float = field(
        default=3.0, metadata={"help": "Mean of Poisson distribution used to generate span-lengths to be masked"}
    )

    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]
                if extension not in ["csv", "json", "txt"]:
                    raise ValueError("train_file` should be a csv, json or text file.")
            if self.validation_file is not None:
                extension = self.validation_file.split(".")[-1]
                if extension not in ["csv", "json", "txt"]:
                    raise ValueError("`validation_file` should be a csv, json or text file.")


@flax.struct.dataclass
class FlaxDataCollatorForBartDenoisingLM:
    

    tokenizer: PreTrainedTokenizerBase
    decoder_start_token_id: int
    mask_ratio: float = 0.3
    poisson_lambda: float = 3.0
    permute_sentence_ratio: float = 1.0

    def __post_init__(self):
        if self.tokenizer.mask_token is None or self.tokenizer.eos_token is None:
            raise ValueError(
                "This tokenizer does not have a mask token or eos token token which is necessary for denoising"
                " language modeling. "
            )

    def __call__(self, examples: List[Dict[str, List[int]]]) -> BatchEncoding:
        
        batch = BatchEncoding(
            {k: np.array([examples[i][k] for i in range(len(examples))]) for k, v in examples[0].items()}
        )
        batch["labels"] = batch["input_ids"].copy()
        batch["decoder_input_ids"] = shift_tokens_right(
            batch["labels"], self.tokenizer.pad_token_id, self.decoder_start_token_id
        )
        
        do_permute = False
        if self.permute_sentence_ratio > 0.0:
            batch["input_ids"] = self.permute_sentences(batch["input_ids"])
            do_permute = True

        
        if self.mask_ratio:
            batch["input_ids"], batch["labels"] = self.span_mask_tokens(
                batch["input_ids"], batch["labels"], do_permute
            )

        
        batch["attention_mask"] = (batch["input_ids"] != self.tokenizer.pad_token_id).astype(int)
        batch["decoder_attention_mask"] = (batch["decoder_input_ids"] != self.tokenizer.pad_token_id).astype(int)
        return batch

    def permute_sentences(self, input_ids):
        
        results = input_ids.copy()

        
        end_sentence_mask = input_ids == self.tokenizer.pad_token_id
        sentence_ends = np.argwhere(end_sentence_mask)
        sentence_ends[:, 1] += 1
        example_has_multiple_sentences, num_sentences = np.unique(sentence_ends[:, 0], return_counts=True)
        num_sentences_map = dict(zip(example_has_multiple_sentences, num_sentences))

        num_to_permute = np.ceil(num_sentences * self.permute_sentence_ratio).astype(int)
        num_to_permute_map = dict(zip(example_has_multiple_sentences, num_to_permute))

        sentence_ends = np.split(sentence_ends[:, 1], np.unique(sentence_ends[:, 0], return_index=True)[1][1:])
        sentence_ends_map = dict(zip(example_has_multiple_sentences, sentence_ends))

        for i in range(input_ids.shape[0]):
            if i not in example_has_multiple_sentences:
                continue
            substitutions = np.random.permutation(num_sentences_map[i])[: num_to_permute_map[i]]
            ordering = np.arange(0, num_sentences_map[i])
            ordering[substitutions] = substitutions[np.random.permutation(num_to_permute_map[i])]

            
            index = 0
            for j in ordering:
                sentence = input_ids[i, (sentence_ends_map[i][j - 1] if j > 0 else 0) : sentence_ends_map[i][j]]
                results[i, index : index + sentence.shape[0]] = sentence
                index += sentence.shape[0]
        return results

    def span_mask_tokens(self, input_ids, labels, do_permute):
        
        special_tokens_mask_labels = [
            self.tokenizer.get_special_tokens_mask(val, already_has_special_tokens=True) for val in labels.tolist()
        ]
        special_tokens_mask_inputs = [
            self.tokenizer.get_special_tokens_mask(val, already_has_special_tokens=True) for val in input_ids.tolist()
        ]
        special_tokens_mask_labels = np.array(special_tokens_mask_labels, dtype=bool)
        special_tokens_mask_inputs = np.array(special_tokens_mask_inputs, dtype=bool)

        
        is_token_mask = ~(input_ids == self.tokenizer.pad_token_id) & ~special_tokens_mask_inputs
        num_tokens_to_mask = int(math.ceil(is_token_mask.astype(float).sum() * self.mask_ratio))
        if num_tokens_to_mask == 0:
            return input_ids, labels

        
        span_lengths = np.random.poisson(lam=self.poisson_lambda, size=(num_tokens_to_mask,))
        while np.cumsum(span_lengths, 0)[-1] < num_tokens_to_mask:
            span_lengths = np.concatenate(
                [span_lengths, np.random.poisson(lam=self.poisson_lambda, size=(num_tokens_to_mask,))]
            )

        
        
        
        span_lengths = span_lengths[span_lengths > 0]

        
        cutoff_idx = np.argmin(np.abs(np.cumsum(span_lengths, 0) - num_tokens_to_mask)) + 1
        span_lengths = span_lengths[:cutoff_idx]

        
        token_indices = np.argwhere(is_token_mask == 1)
        span_starts = np.random.permutation(token_indices.shape[0])[: span_lengths.shape[0]]
        
        masked_indices = np.array(token_indices[span_starts])
        mask = np.full_like(input_ids, fill_value=False)

        
        for mi in masked_indices:
            mask[tuple(mi)] = True
        span_lengths -= 1

        
        max_index = input_ids.shape[1] - 1
        remaining = (span_lengths > 0) & (masked_indices[:, 1] < max_index)
        while np.any(remaining):
            masked_indices[remaining, 1] += 1
            for mi in masked_indices:
                mask[tuple(mi)] = True
            span_lengths -= 1
            remaining = (span_lengths > 0) & (masked_indices[:, 1] < max_index)

        
        mask[np.where(special_tokens_mask_inputs)] = False
        input_ids[np.where(mask)] = self.tokenizer.mask_token_id
        if not do_permute:
            labels[np.where(mask == 0)] = -100
        else:
            labels[np.where(special_tokens_mask_labels)] = -100

        
        to_remove = (mask == 1) & np.roll((mask == 1), 1, 1)
        new_input_ids = np.full_like(input_ids, fill_value=self.tokenizer.pad_token_id)
        for i, example in enumerate(input_ids):
            new_example = example[~to_remove[i]]
            new_input_ids[i, : new_example.shape[0]] = new_example

        return new_input_ids, labels


def generate_batch_splits(samples_idx: np.ndarray, batch_size: int, drop_last=True) -> np.ndarray:
    
    num_samples = len(samples_idx)
    if drop_last:
        samples_to_remove = num_samples % batch_size
        if samples_to_remove != 0:
            samples_idx = samples_idx[:-samples_to_remove]
        sections_split = num_samples // batch_size
        samples_idx = samples_idx.reshape((sections_split, batch_size))
    else:
        sections_split = math.ceil(num_samples / batch_size)
        samples_idx = np.array_split(samples_idx, sections_split)
    return samples_idx


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

    train_metrics = get_metrics(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 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 model_args.use_auth_token is not None:
        warnings.warn(
            "The `use_auth_token` argument is deprecated and will be removed in v4.34. Please use `token` instead.",
            FutureWarning,
        )
        if model_args.token is not None:
            raise ValueError("`token` and `use_auth_token` are both specified. Please set only the argument `token`.")
        model_args.token = model_args.use_auth_token

    
    
    send_example_telemetry("run_bart_dlm", model_args, data_args, framework="flax")

    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",
        level=logging.INFO,
        datefmt="[%X]",
    )

    
    logger = logging.getLogger(__name__)

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

    
    set_seed(training_args.seed)

    
    if training_args.push_to_hub:
        
        repo_name = training_args.hub_model_id
        if repo_name is None:
            repo_name = Path(training_args.output_dir).absolute().name
        
        repo_id = create_repo(repo_name, exist_ok=True, token=training_args.hub_token).repo_id
        
        repo = Repository(training_args.output_dir, clone_from=repo_id, token=training_args.hub_token)

    
    
    
    
    
    
    if data_args.dataset_name is not None:
        
        datasets = load_dataset(
            data_args.dataset_name,
            data_args.dataset_config_name,
            cache_dir=model_args.cache_dir,
            token=model_args.token,
            num_proc=data_args.preprocessing_num_workers,
        )

        if "validation" not in datasets.keys():
            datasets["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,
                token=model_args.token,
                num_proc=data_args.preprocessing_num_workers,
            )
            datasets["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,
                token=model_args.token,
                num_proc=data_args.preprocessing_num_workers,
            )
    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"
        datasets = load_dataset(
            extension,
            data_files=data_files,
            cache_dir=model_args.cache_dir,
            token=model_args.token,
            num_proc=data_args.preprocessing_num_workers,
        )

        if "validation" not in datasets.keys():
            datasets["validation"] = load_dataset(
                extension,
                data_files=data_files,
                split=f"train[:{data_args.validation_split_percentage}%]",
                cache_dir=model_args.cache_dir,
                token=model_args.token,
                num_proc=data_args.preprocessing_num_workers,
            )
            datasets["train"] = load_dataset(
                extension,
                data_files=data_files,
                split=f"train[{data_args.validation_split_percentage}%:]",
                cache_dir=model_args.cache_dir,
                token=model_args.token,
                num_proc=data_args.preprocessing_num_workers,
            )
    
    

    

    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,
            token=model_args.token,
        )
    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,
            token=model_args.token,
        )
    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 model_args.config_name:
        config = BartConfig.from_pretrained(
            model_args.config_name,
            cache_dir=model_args.cache_dir,
            vocab_size=len(tokenizer),
            token=model_args.token,
        )
    elif model_args.model_name_or_path:
        config = BartConfig.from_pretrained(
            model_args.model_name_or_path,
            cache_dir=model_args.cache_dir,
            token=model_args.token,
        )
    else:
        config = CONFIG_MAPPING[model_args.model_type]()
        logger.warning("You are instantiating a new config instance from scratch.")

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

    max_seq_length = min(data_args.max_seq_length, tokenizer.model_max_length)

    
    nltk.download("punkt")
    sentence_tokenizer = nltk.data.load("tokenizers/punkt/english.pickle")

    def sentence_split_function(example):
        sents = sentence_tokenizer.tokenize(example["text"])
        
        new_text = tokenizer.bos_token + f"{tokenizer.pad_token}".join(sents) + tokenizer.eos_token
        return {"text": new_text}

    split_datasets = datasets.map(
        sentence_split_function,
        batched=False,
        num_proc=data_args.preprocessing_num_workers,
        remove_columns=column_names,
        load_from_cache_file=not data_args.overwrite_cache,
    )

    
    
    def tokenize_function(examples):
        return tokenizer(examples[text_column_name], add_special_tokens=False, return_attention_mask=False)

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

    
    
    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 >= max_seq_length:
            total_length = (total_length // max_seq_length) * max_seq_length
        
        result = {
            k: [t[i : i + max_seq_length] for i in range(0, total_length, max_seq_length)]
            for k, t in concatenated_examples.items()
        }
        return result

    
    
    
    
    
    
    tokenized_datasets = tokenized_datasets.map(
        group_texts,
        batched=True,
        num_proc=data_args.preprocessing_num_workers,
        load_from_cache_file=not data_args.overwrite_cache,
    )

    
    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)
    dropout_rngs = jax.random.split(rng, jax.local_device_count())

    if model_args.model_name_or_path:
        model = FlaxBartForConditionalGeneration.from_pretrained(
            model_args.model_name_or_path,
            config=config,
            seed=training_args.seed,
            dtype=getattr(jnp, model_args.dtype),
            token=model_args.token,
        )
    else:
        config.vocab_size = len(tokenizer)
        model = FlaxBartForConditionalGeneration(
            config,
            seed=training_args.seed,
            dtype=getattr(jnp, model_args.dtype),
        )

    
    
    data_collator = FlaxDataCollatorForBartDenoisingLM(
        tokenizer=tokenizer,
        decoder_start_token_id=model.config.decoder_start_token_id,
        mask_ratio=data_args.mlm_probability,
        poisson_lambda=data_args.poisson_lambda,
        permute_sentence_ratio=data_args.permute_sentence_ratio,
    )

    
    num_epochs = int(training_args.num_train_epochs)
    train_batch_size = int(training_args.per_device_train_batch_size) * jax.device_count()
    per_device_eval_batch_size = int(training_args.per_device_eval_batch_size)
    eval_batch_size = per_device_eval_batch_size * jax.device_count()

    num_train_steps = len(tokenized_datasets["train"]) // train_batch_size * num_epochs

    
    warmup_fn = optax.linear_schedule(
        init_value=0.0, end_value=training_args.learning_rate, transition_steps=training_args.warmup_steps
    )
    decay_fn = optax.linear_schedule(
        init_value=training_args.learning_rate,
        end_value=0,
        transition_steps=num_train_steps - training_args.warmup_steps,
    )
    linear_decay_lr_schedule_fn = optax.join_schedules(
        schedules=[warmup_fn, decay_fn], boundaries=[training_args.warmup_steps]
    )

    
    
    
    
    def decay_mask_fn(params):
        flat_params = traverse_util.flatten_dict(params)
        
        layer_norm_candidates = ["layernorm", "layer_norm", "ln"]
        layer_norm_named_params = {
            layer[-2:]
            for layer_norm_name in layer_norm_candidates
            for layer in flat_params.keys()
            if layer_norm_name in "".join(layer).lower()
        }
        flat_mask = {path: (path[-1] != "bias" and path[-2:] not in layer_norm_named_params) for path in flat_params}
        return traverse_util.unflatten_dict(flat_mask)

    
    if training_args.adafactor:
        
        
        optimizer = optax.adafactor(
            learning_rate=linear_decay_lr_schedule_fn,
        )
    else:
        optimizer = optax.adamw(
            learning_rate=linear_decay_lr_schedule_fn,
            b1=training_args.adam_beta1,
            b2=training_args.adam_beta2,
            weight_decay=training_args.weight_decay,
            mask=decay_mask_fn,
        )

    
    state = train_state.TrainState.create(apply_fn=model.__call__, params=model.params, tx=optimizer)

    
    def train_step(state, batch, dropout_rng):
        dropout_rng, new_dropout_rng = jax.random.split(dropout_rng)

        def loss_fn(params):
            labels = batch.pop("labels")

            logits = state.apply_fn(**batch, params=params, dropout_rng=dropout_rng, train=True)[0]

            
            label_mask = jnp.where(labels > 0, 1.0, 0.0)
            loss = optax.softmax_cross_entropy(logits, onehot(labels, logits.shape[-1])) * label_mask

            
            loss = loss.sum()
            num_labels = label_mask.sum()

            return loss, num_labels

        grad_fn = jax.value_and_grad(loss_fn, has_aux=True)
        (loss, num_labels), grad = grad_fn(state.params)
        num_labels = jax.lax.psum(num_labels, "batch")

        
        loss = jax.lax.psum(loss, "batch")
        loss = jax.tree_util.tree_map(lambda x: x / num_labels, loss)

        
        grad = jax.lax.psum(grad, "batch")
        grad = jax.tree_util.tree_map(lambda x: x / num_labels, grad)
        new_state = state.apply_gradients(grads=grad)

        metrics = {"loss": loss, "learning_rate": linear_decay_lr_schedule_fn(state.step)}
        return new_state, metrics, new_dropout_rng

    
    p_train_step = jax.pmap(train_step, "batch", donate_argnums=(0,))

    
    def eval_step(params, batch):
        labels = batch.pop("labels")

        logits = model(**batch, params=params, train=False)[0]

        
        label_mask = jnp.where(labels > 0, 1.0, 0.0)
        loss = optax.softmax_cross_entropy(logits, onehot(labels, logits.shape[-1])) * label_mask

        
        accuracy = jnp.equal(jnp.argmax(logits, axis=-1), labels) * label_mask

        
        metrics = {"loss": loss.sum(), "accuracy": accuracy.sum(), "normalizer": label_mask.sum()}
        metrics = jax.lax.psum(metrics, axis_name="batch")

        return metrics

    p_eval_step = jax.pmap(eval_step, "batch", donate_argnums=(0,))

    
    state = jax_utils.replicate(state)

    train_time = 0
    epochs = tqdm(range(num_epochs), desc="Epoch ... ", position=0)
    for epoch in epochs:
        
        train_start = time.time()
        train_metrics = []

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

        
        num_train_samples = len(tokenized_datasets["train"])
        
        train_samples_idx = np.random.permutation(np.arange(num_train_samples))
        train_batch_idx = generate_batch_splits(train_samples_idx, train_batch_size)

        
        for step, batch_idx in enumerate(tqdm(train_batch_idx, desc="Training...", position=1)):
            samples = [tokenized_datasets["train"][int(idx)] for idx in batch_idx]
            model_inputs = data_collator(samples)

            
            model_inputs = shard(model_inputs.data)
            state, train_metric, dropout_rngs = p_train_step(state, model_inputs, dropout_rngs)
            train_metrics.append(train_metric)

            cur_step = epoch * (num_train_samples // train_batch_size) + step

            if cur_step % training_args.logging_steps == 0 and cur_step > 0:
                
                train_metric = jax_utils.unreplicate(train_metric)
                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:
                
                num_eval_samples = len(tokenized_datasets["validation"])
                
                eval_samples_idx = np.arange(num_eval_samples)
                eval_batch_idx = generate_batch_splits(eval_samples_idx, eval_batch_size)

                eval_metrics = []
                for i, batch_idx in enumerate(tqdm(eval_batch_idx, desc="Evaluating ...", position=2)):
                    samples = [tokenized_datasets["validation"][int(idx)] for idx in batch_idx]
                    model_inputs = data_collator(samples)

                    
                    metrics = pad_shard_unpad(p_eval_step, static_return=True)(
                        state.params, model_inputs.data, min_device_batch=per_device_eval_batch_size
                    )
                    eval_metrics.append(metrics)

                
                eval_metrics = get_metrics(eval_metrics)
                eval_metrics = jax.tree_util.tree_map(jnp.sum, eval_metrics)
                eval_normalizer = eval_metrics.pop("normalizer")
                eval_metrics = jax.tree_util.tree_map(lambda x: x / eval_normalizer, eval_metrics)

                
                epochs.desc = f"Step... ({cur_step} | Loss: {eval_metrics['loss']}, Acc: {eval_metrics['accuracy']})"

                
                if has_tensorboard and jax.process_index() == 0:
                    write_eval_metric(summary_writer, eval_metrics, cur_step)

            if cur_step % training_args.save_steps == 0 and cur_step > 0:
                
                if jax.process_index() == 0:
                    params = jax.device_get(jax.tree_util.tree_map(lambda x: x[0], state.params))
                    model.save_pretrained(training_args.output_dir, params=params)
                    tokenizer.save_pretrained(training_args.output_dir)
                    if training_args.push_to_hub:
                        repo.push_to_hub(commit_message=f"Saving weights and logs of step {cur_step}", blocking=False)

    
    if training_args.do_eval:
        num_eval_samples = len(tokenized_datasets["validation"])
        
        eval_samples_idx = np.arange(num_eval_samples)
        eval_batch_idx = generate_batch_splits(eval_samples_idx, eval_batch_size)

        eval_metrics = []
        for _, batch_idx in enumerate(tqdm(eval_batch_idx, desc="Evaluating ...", position=2)):
            samples = [tokenized_datasets["validation"][int(idx)] for idx in batch_idx]
            model_inputs = data_collator(samples)

            
            metrics = pad_shard_unpad(p_eval_step, static_return=True)(
                state.params, model_inputs.data, min_device_batch=per_device_eval_batch_size
            )
            eval_metrics.append(metrics)

        
        eval_metrics = get_metrics(eval_metrics)
        eval_metrics = jax.tree_util.tree_map(lambda metric: jnp.sum(metric).item(), eval_metrics)
        eval_normalizer = eval_metrics.pop("normalizer")
        eval_metrics = jax.tree_util.tree_map(lambda x: x / eval_normalizer, eval_metrics)

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

        if jax.process_index() == 0:
            eval_metrics = {f"eval_{metric_name}": value for metric_name, value in eval_metrics.items()}
            path = os.path.join(training_args.output_dir, "eval_results.json")
            with open(path, "w") as f:
                json.dump(eval_metrics, f, indent=4, sort_keys=True)


if __name__ == "__main__":
    main()