usrbinenv python codingutf8 2021 the huggingface inc 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 finetuning a flax transformers model for sequence classification on glue import json import logging import math import os import random import sys import time import warnings from dataclasses import dataclass field from pathlib import path from typing import any callable dict optional tuple import datasets import evaluate import jax import jax numpy as jnp import numpy as np import optax from datasets import loaddataset from flax import struct traverseutil from flax jaxutils import padshardunpad replicate unreplicate from flax training import trainstate from flax training commonutils import getmetrics onehot shard from huggingfacehub import repository createrepo from tqdm import tqdm import transformers from transformers import autoconfig autotokenizer flaxautomodelforsequenceclassification hfargumentparser pretrainedconfig trainingarguments istensorboardavailable from transformers utils import checkminversion sendexampletelemetry logger logging getloggername will error if the minimal version of transformers is not installed remove at your own risks checkminversion4 36 0 dev0 array any dataset datasets arrowdataset dataset prngkey any tasktokeys cola sentence none mnli premise hypothesis mrpc sentence1 sentence2 qnli question sentence qqp question1 question2 rte sentence1 sentence2 sst2 sentence none stsb sentence1 sentence2 wnli sentence1 sentence2 dataclass class modelarguments modelnameorpath str field metadatahelp path to pretrained model or model identifier from huggingface comodels configname optionalstr field defaultnone metadatahelp pretrained config name or path if not the same as modelname tokenizername optionalstr field defaultnone metadatahelp pretrained tokenizer name or path if not the same as modelname useslowtokenizer optionalbool field defaultfalse metadatahelp if passed will use a slow tokenizer not backed by the tokenizers library cachedir optionalstr field defaultnone metadatahelp where do you want to store the pretrained models downloaded from huggingface co modelrevision str field defaultmain metadatahelp the specific model version to use can be a branch name tag name or commit id token str field defaultnone metadata help the token to use as http bearer ization for remote files if not specified will use the token generated when running huggingfacecli login stored in huggingface useauthtoken bool field defaultnone metadata help the useauthtoken argument is deprecated and will be removed in v4 34 please use token instead trustremotecode bool field defaultfalse metadata help whether or not to allow for custom models defined on the hub in their own modeling files this option should only be set to true for repositories you trust and in which you have read the code as it will execute code present on the hub on your local machine dataclass class datatrainingarguments taskname optionalstr field defaultnone metadatahelp fthe name of the glue task to train on choices listtasktokeys keys datasetconfigname optionalstr field defaultnone metadatahelp the configuration name of the dataset to use via the datasets library trainfile optionalstr field defaultnone metadatahelp the input training data file a csv or json file validationfile optionalstr field defaultnone metadatahelp an optional input evaluation data file to evaluate on a csv or json file testfile optionalstr field defaultnone metadatahelp an optional input test data file to predict on a csv or json file textcolumnname optionalstr field defaultnone metadatahelp the column name of text to input in the file a csv or json file labelcolumnname optionalstr field defaultnone metadatahelp the column name of label to input in the file a csv or json file overwritecache bool field defaultfalse metadatahelp overwrite the cached training and evaluation sets preprocessingnumworkers optionalint field defaultnone metadatahelp the number of processes to use for the preprocessing maxseqlength int field defaultnone metadata help the maximum total input sequence length after tokenization if set sequences longer than this will be truncated sequences shorter will be padded maxtrainsamples optionalint field defaultnone metadata help for debugging purposes or quicker training truncate the number of training examples to this value if set maxevalsamples optionalint field defaultnone metadata help for debugging purposes or quicker training truncate the number of evaluation examples to this value if set maxpredictsamples optionalint field defaultnone metadata help for debugging purposes or quicker training truncate the number of prediction examples to this value if set def postinitself if self taskname is none and self trainfile is none and self validationfile is none raise valueerrorneed either a dataset name or a trainingvalidation file else if self trainfile is not none extension self trainfile split 1 assert extension in csv json trainfile should be a csv or a json file if self validationfile is not none extension self validationfile split 1 assert extension in csv json validationfile should be a csv or a json file self taskname self taskname lower if isinstanceself taskname str else self taskname def createtrainstate model flaxautomodelforsequenceclassification learningratefn callableint float isregression bool numlabels int weightdecay float trainstate trainstate train state with an optax optimizer the two functions below differ depending on whether the task is classification or regression args logitsfn applied to last layer to obtain the logits lossfn function to compute the loss 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 find out all layernorm parameters 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 gluetraindatacollatorrng prngkey dataset dataset batchsize int returns batches of size batchsize from eval dataset sharding handled by padshardunpad in the eval loop batchidx np arangelendataset stepsperepoch math ceillendataset batchsize batchidx np arraysplitbatchidx stepsperepoch for idx in batchidx batch datasetidx batch k np arrayv for k v in batch items yield batch def main 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 sendexampletelemetryrunglue modelargs dataargs frameworkflax 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 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 csvjson training and evaluation files see below or specify a glue benchmark task the dataset will be downloaded automatically from the datasets hub for csvjson files this script will use as labels the column called label and as pair of sentences the sentences in columns called sentence1 and sentence2 if such column exists or the first two columns not named label if at least two columns are provided if the csvsjsons contain only one nonlabel column the script does single sentence classification on this single column you can easily tweak this behavior see below in distributed training the loaddataset function guarantee that only one local process can concurrently download the dataset if dataargs taskname is not none downloading and loading a dataset from the hub rawdatasets loaddataset glue dataargs taskname tokenmodelargs token else loading the dataset from local csv or json file datafiles if dataargs trainfile is not none datafilestrain dataargs trainfile if dataargs validationfile is not none datafilesvalidation dataargs validationfile extension dataargs trainfile if dataargs trainfile is not none else dataargs validfile split 1 rawdatasets loaddataset extension datafilesdatafiles tokenmodelargs token see more about loading any type of standard or custom dataset at https huggingface codocsdatasetsloadingdatasets labels if dataargs taskname is not none isregression dataargs taskname stsb if not isregression labellist rawdatasetstrain featureslabel names numlabels lenlabellist else numlabels 1 else trying to have good defaults here don t hesitate to tweak to your needs isregression rawdatasetstrain featureslabel dtype in float32 float64 if isregression numlabels 1 else a useful fast method https huggingface codocsdatasetspackagereferencemainclassesdatasets dataset unique labellist rawdatasetstrain uniquelabel labellist sort let s sort it for determinism numlabels lenlabellist load pretrained model and tokenizer config autoconfig frompretrained modelargs modelnameorpath numlabelsnumlabels finetuningtaskdataargs taskname tokenmodelargs token trustremotecodemodelargs trustremotecode tokenizer autotokenizer frompretrained modelargs modelnameorpath usefastnot modelargs useslowtokenizer tokenmodelargs token trustremotecodemodelargs trustremotecode model flaxautomodelforsequenceclassification frompretrained modelargs modelnameorpath configconfig tokenmodelargs token trustremotecodemodelargs trustremotecode preprocessing the datasets if dataargs taskname is not none sentence1key sentence2key tasktokeysdataargs taskname else again we try to have some nice defaults but don t hesitate to tweak to your use case nonlabelcolumnnames name for name in rawdatasetstrain columnnames if name label if sentence1 in nonlabelcolumnnames and sentence2 in nonlabelcolumnnames sentence1key sentence2key sentence1 sentence2 else if lennonlabelcolumnnames 2 sentence1key sentence2key nonlabelcolumnnames 2 else sentence1key sentence2key nonlabelcolumnnames0 none some models have set the order of the labels to use so let s make sure we do use it labeltoid none if model config label2id pretrainedconfignumlabelsnumlabels label2id and dataargs taskname is not none and not isregression some have all caps in their config some don t labelnametoid k lower v for k v in model config label2id items if sortedlabelnametoid keys sortedlabellist logger info fthe configuration of the model provided the following label correspondence labelnametoid using it labeltoid i labelnametoidlabellisti for i in rangenumlabels else logger warning your model seems to have been trained with labels but they don t match the dataset fmodel labels sortedlabelnametoid keys dataset labels sortedlabellist nignoring the model labels as a result elif dataargs taskname is none labeltoid v i for i v in enumeratelabellist def preprocessfunctionexamples tokenize the texts texts examplessentence1key if sentence2key is none else examplessentence1key examplessentence2key result tokenizertexts paddingmaxlength maxlengthdataargs maxseqlength truncationtrue if label in examples if labeltoid is not none map labels to ids not necessary for glue tasks resultlabels labeltoidl for l in exampleslabel else in all cases rename the column to labels because the model will expect that resultlabels exampleslabel return result processeddatasets rawdatasets map preprocessfunction batchedtrue removecolumnsrawdatasetstrain columnnames traindataset processeddatasetstrain evaldataset processeddatasetsvalidationmatched if dataargs taskname mnli else validation log a few random samples from the training set for index in random samplerangelentraindataset 3 logger infofsample index of the training set traindatasetindex define a summary writer hastensorboard istensorboardavailable if hastensorboard and jax processindex 0 try from flax metrics tensorboard import summarywriter summarywriter summarywritertrainingargs outputdir summarywriter hparamstrainingargs todict varsmodelargs varsdataargs 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 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 numepochs inttrainingargs numtrainepochs rng jax random prngkeytrainingargs seed dropoutrngs jax random splitrng jax localdevicecount trainbatchsize inttrainingargs perdevicetrainbatchsize jax localdevicecount perdeviceevalbatchsize inttrainingargs perdeviceevalbatchsize evalbatchsize perdeviceevalbatchsize jax devicecount learningratefn createlearningratefn lentraindataset trainbatchsize trainingargs numtrainepochs trainingargs warmupsteps trainingargs learningrate state createtrainstate model learningratefn isregression numlabelsnumlabels weightdecaytrainingargs weightdecay define step functions def trainstep state trainstate trainstate batch dictstr array dropoutrng prngkey tupletrainstate trainstate float make sure weights are replicated on each device create sampling rng train save metrics evaluate save checkpoint after each epoch and push checkpoint to the hub save the eval metrics in json usr bin env python coding utf 8 2021 the huggingface inc 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 finetuning a flax transformers model for sequence classification on glue will error if the minimal version of transformers is not installed remove at your own risks arguments pertaining to which model config tokenizer we are going to fine tune from arguments pertaining to what data we are going to input our model for training and eval create initial training state train state with an optax optimizer the two functions below differ depending on whether the task is classification or regression args logits_fn applied to last layer to obtain the logits loss_fn function to compute the loss 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 classification returns a linear warmup linear_decay learning rate function returns shuffled batches of size batch_size from truncated train dataset sharded over all local devices skip incomplete batch returns batches of size batch_size from eval dataset sharding handled by pad_shard_unpad in the eval loop 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 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 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 training and evaluation files see below or specify a glue benchmark task the dataset will be downloaded automatically from the datasets hub for csv json files this script will use as labels the column called label and as pair of sentences the sentences in columns called sentence1 and sentence2 if such column exists or the first two columns not named label if at least two columns are provided if the csvs jsons contain only one non label column the script does single sentence classification on this single column you can easily tweak this behavior see below in distributed training the load_dataset function guarantee that only one local process can concurrently download the dataset downloading and loading a dataset from the hub loading the dataset from local csv or json file see more about loading any type of standard or custom dataset at https huggingface co docs datasets loading_datasets labels trying to have good defaults here don t hesitate to tweak to your needs a useful fast method https huggingface co docs datasets package_reference main_classes datasets dataset unique let s sort it for determinism load pretrained model and tokenizer preprocessing the datasets again we try to have some nice defaults but don t hesitate to tweak to your use case some models have set the order of the labels to use so let s make sure we do use it some have all caps in their config some don t tokenize the texts map labels to ids not necessary for glue tasks in all cases rename the column to labels because the model will expect that log a few random samples from the training set define a summary writer define step functions trains model with an optimizer both in state on batch returning a pair new_state loss make sure weights are replicated on each device create sampling rng train save metrics evaluate save checkpoint after each epoch and push checkpoint to the hub save the eval metrics in json
import json
import logging
import math
import os
import random
import sys
import time
import warnings
from dataclasses import dataclass, field
from pathlib import Path
from typing import Any, Callable, Dict, Optional, Tuple

import datasets
import evaluate
import jax
import jax.numpy as jnp
import numpy as np
import optax
from datasets import load_dataset
from flax import struct, traverse_util
from flax.jax_utils import pad_shard_unpad, replicate, unreplicate
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

import transformers
from transformers import (
    AutoConfig,
    AutoTokenizer,
    FlaxAutoModelForSequenceClassification,
    HfArgumentParser,
    PretrainedConfig,
    TrainingArguments,
    is_tensorboard_available,
)
from transformers.utils import check_min_version, send_example_telemetry


logger = logging.getLogger(__name__)

check_min_version("4.36.0.dev0")

Array = Any
Dataset = datasets.arrow_dataset.Dataset
PRNGKey = Any


task_to_keys = {
    "cola": ("sentence", None),
    "mnli": ("premise", "hypothesis"),
    "mrpc": ("sentence1", "sentence2"),
    "qnli": ("question", "sentence"),
    "qqp": ("question1", "question2"),
    "rte": ("sentence1", "sentence2"),
    "sst2": ("sentence", None),
    "stsb": ("sentence1", "sentence2"),
    "wnli": ("sentence1", "sentence2"),
}


@dataclass
class ModelArguments:
    

    model_name_or_path: str = field(
        metadata={"help": "Path to pretrained model or model identifier from huggingface.co/models"}
    )
    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"}
    )
    use_slow_tokenizer: Optional[bool] = field(
        default=False,
        metadata={"help": "If passed, will use a slow tokenizer (not backed by the 🤗 Tokenizers library)."},
    )
    cache_dir: Optional[str] = field(
        default=None,
        metadata={"help": "Where do you want to store the pretrained models downloaded from huggingface.co"},
    )
    model_revision: str = field(
        default="main",
        metadata={"help": "The specific model version to use (can be a branch name, tag name or commit id)."},
    )
    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."
        },
    )
    trust_remote_code: bool = field(
        default=False,
        metadata={
            "help": (
                "Whether or not to allow for custom models defined on the Hub in their own modeling files. This option"
                "should only be set to `True` for repositories you trust and in which you have read the code, as it will "
                "execute code present on the Hub on your local machine."
            )
        },
    )


@dataclass
class DataTrainingArguments:
    

    task_name: Optional[str] = field(
        default=None, metadata={"help": f"The name of the glue task to train on. choices {list(task_to_keys.keys())}"}
    )
    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 csv or JSON file)."}
    )
    validation_file: Optional[str] = field(
        default=None,
        metadata={"help": "An optional input evaluation data file to evaluate on (a csv or JSON file)."},
    )
    test_file: Optional[str] = field(
        default=None,
        metadata={"help": "An optional input test data file to predict on (a csv or JSON file)."},
    )
    text_column_name: Optional[str] = field(
        default=None, metadata={"help": "The column name of text to input in the file (a csv or JSON file)."}
    )
    label_column_name: Optional[str] = field(
        default=None, metadata={"help": "The column name of label to input in the file (a csv or JSON file)."}
    )
    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."},
    )
    max_seq_length: int = field(
        default=None,
        metadata={
            "help": (
                "The maximum total input sequence length after tokenization. If set, sequences longer "
                "than this will be truncated, sequences shorter will be padded."
            )
        },
    )
    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."
            )
        },
    )
    max_predict_samples: Optional[int] = field(
        default=None,
        metadata={
            "help": (
                "For debugging purposes or quicker training, truncate the number of prediction examples to this "
                "value if set."
            )
        },
    )

    def __post_init__(self):
        if self.task_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"], "`train_file` should be a csv or a json file."
            if self.validation_file is not None:
                extension = self.validation_file.split(".")[-1]
                assert extension in ["csv", "json"], "`validation_file` should be a csv or a json file."
        self.task_name = self.task_name.lower() if isinstance(self.task_name, str) else self.task_name


def create_train_state(
    model: FlaxAutoModelForSequenceClassification,
    learning_rate_fn: Callable[[int], float],
    is_regression: bool,
    num_labels: int,
    weight_decay: float,
) -> train_state.TrainState:
    

    class TrainState(train_state.TrainState):
        

        logits_fn: Callable = struct.field(pytree_node=False)
        loss_fn: Callable = struct.field(pytree_node=False)

    
    
    
    
    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)

    tx = optax.adamw(
        learning_rate=learning_rate_fn, b1=0.9, b2=0.999, eps=1e-6, weight_decay=weight_decay, mask=decay_mask_fn
    )

    if is_regression:

        def mse_loss(logits, labels):
            return jnp.mean((logits[..., 0] - labels) ** 2)

        return TrainState.create(
            apply_fn=model.__call__,
            params=model.params,
            tx=tx,
            logits_fn=lambda logits: logits[..., 0],
            loss_fn=mse_loss,
        )
    else:  

        def cross_entropy_loss(logits, labels):
            xentropy = optax.softmax_cross_entropy(logits, onehot(labels, num_classes=num_labels))
            return jnp.mean(xentropy)

        return TrainState.create(
            apply_fn=model.__call__,
            params=model.params,
            tx=tx,
            logits_fn=lambda logits: logits.argmax(-1),
            loss_fn=cross_entropy_loss,
        )


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 glue_train_data_collator(rng: PRNGKey, dataset: Dataset, batch_size: int):
    
    steps_per_epoch = len(dataset) // batch_size
    perms = jax.random.permutation(rng, len(dataset))
    perms = perms[: steps_per_epoch * batch_size]  
    perms = perms.reshape((steps_per_epoch, batch_size))

    for perm in perms:
        batch = dataset[perm]
        batch = {k: np.array(v) for k, v in batch.items()}
        batch = shard(batch)

        yield batch


def glue_eval_data_collator(dataset: Dataset, batch_size: int):
    
    batch_idx = np.arange(len(dataset))

    steps_per_epoch = math.ceil(len(dataset) / batch_size)
    batch_idx = np.array_split(batch_idx, steps_per_epoch)

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

        yield batch


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_glue", model_args, data_args, framework="flax")

    
    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()

    
    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.task_name is not None:
        
        raw_datasets = load_dataset(
            "glue",
            data_args.task_name,
            token=model_args.token,
        )
    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 if data_args.train_file is not None else data_args.valid_file).split(".")[-1]
        raw_datasets = load_dataset(
            extension,
            data_files=data_files,
            token=model_args.token,
        )
    
    

    
    if data_args.task_name is not None:
        is_regression = data_args.task_name == "stsb"
        if not is_regression:
            label_list = raw_datasets["train"].features["label"].names
            num_labels = len(label_list)
        else:
            num_labels = 1
    else:
        
        is_regression = raw_datasets["train"].features["label"].dtype in ["float32", "float64"]
        if is_regression:
            num_labels = 1
        else:
            
            
            label_list = raw_datasets["train"].unique("label")
            label_list.sort()  
            num_labels = len(label_list)

    
    config = AutoConfig.from_pretrained(
        model_args.model_name_or_path,
        num_labels=num_labels,
        finetuning_task=data_args.task_name,
        token=model_args.token,
        trust_remote_code=model_args.trust_remote_code,
    )
    tokenizer = AutoTokenizer.from_pretrained(
        model_args.model_name_or_path,
        use_fast=not model_args.use_slow_tokenizer,
        token=model_args.token,
        trust_remote_code=model_args.trust_remote_code,
    )
    model = FlaxAutoModelForSequenceClassification.from_pretrained(
        model_args.model_name_or_path,
        config=config,
        token=model_args.token,
        trust_remote_code=model_args.trust_remote_code,
    )

    
    if data_args.task_name is not None:
        sentence1_key, sentence2_key = task_to_keys[data_args.task_name]
    else:
        
        non_label_column_names = [name for name in raw_datasets["train"].column_names if name != "label"]
        if "sentence1" in non_label_column_names and "sentence2" in non_label_column_names:
            sentence1_key, sentence2_key = "sentence1", "sentence2"
        else:
            if len(non_label_column_names) >= 2:
                sentence1_key, sentence2_key = non_label_column_names[:2]
            else:
                sentence1_key, sentence2_key = non_label_column_names[0], None

    
    label_to_id = None
    if (
        model.config.label2id != PretrainedConfig(num_labels=num_labels).label2id
        and data_args.task_name is not None
        and not is_regression
    ):
        
        label_name_to_id = {k.lower(): v for k, v in model.config.label2id.items()}
        if sorted(label_name_to_id.keys()) == sorted(label_list):
            logger.info(
                f"The configuration of the model provided the following label correspondence: {label_name_to_id}. "
                "Using it!"
            )
            label_to_id = {i: label_name_to_id[label_list[i]] for i in range(num_labels)}
        else:
            logger.warning(
                "Your model seems to have been trained with labels, but they don't match the dataset: ",
                f"model labels: {sorted(label_name_to_id.keys())}, dataset labels: {sorted(label_list)}."
                "\nIgnoring the model labels as a result.",
            )
    elif data_args.task_name is None:
        label_to_id = {v: i for i, v in enumerate(label_list)}

    def preprocess_function(examples):
        
        texts = (
            (examples[sentence1_key],) if sentence2_key is None else (examples[sentence1_key], examples[sentence2_key])
        )
        result = tokenizer(*texts, padding="max_length", max_length=data_args.max_seq_length, truncation=True)

        if "label" in examples:
            if label_to_id is not None:
                
                result["labels"] = [label_to_id[l] for l in examples["label"]]
            else:
                
                result["labels"] = examples["label"]
        return result

    processed_datasets = raw_datasets.map(
        preprocess_function, batched=True, remove_columns=raw_datasets["train"].column_names
    )

    train_dataset = processed_datasets["train"]
    eval_dataset = processed_datasets["validation_matched" if data_args.task_name == "mnli" else "validation"]

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

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

            summary_writer = SummaryWriter(training_args.output_dir)
            summary_writer.hparams({**training_args.to_dict(), **vars(model_args), **vars(data_args)})
        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."
        )

    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)

    num_epochs = int(training_args.num_train_epochs)
    rng = jax.random.PRNGKey(training_args.seed)
    dropout_rngs = jax.random.split(rng, jax.local_device_count())

    train_batch_size = int(training_args.per_device_train_batch_size) * jax.local_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()

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

    state = create_train_state(
        model, learning_rate_fn, is_regression, num_labels=num_labels, weight_decay=training_args.weight_decay
    )

    
    def train_step(
        state: train_state.TrainState, batch: Dict[str, Array], dropout_rng: PRNGKey
    ) -> Tuple[train_state.TrainState, float]:
        
        dropout_rng, new_dropout_rng = jax.random.split(dropout_rng)
        targets = batch.pop("labels")

        def loss_fn(params):
            logits = state.apply_fn(**batch, params=params, dropout_rng=dropout_rng, train=True)[0]
            loss = state.loss_fn(logits, targets)
            return loss

        grad_fn = jax.value_and_grad(loss_fn)
        loss, grad = grad_fn(state.params)
        grad = jax.lax.pmean(grad, "batch")
        new_state = state.apply_gradients(grads=grad)
        metrics = jax.lax.pmean({"loss": loss, "learning_rate": learning_rate_fn(state.step)}, axis_name="batch")
        return new_state, metrics, new_dropout_rng

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

    def eval_step(state, batch):
        logits = state.apply_fn(**batch, params=state.params, train=False)[0]
        return state.logits_fn(logits)

    p_eval_step = jax.pmap(eval_step, axis_name="batch")

    if data_args.task_name is not None:
        metric = evaluate.load("glue", data_args.task_name)
    else:
        metric = evaluate.load("accuracy")

    logger.info(f"===== Starting training ({num_epochs} epochs) =====")
    train_time = 0

    
    state = replicate(state)

    steps_per_epoch = len(train_dataset) // train_batch_size
    total_steps = steps_per_epoch * num_epochs
    epochs = tqdm(range(num_epochs), desc=f"Epoch ... (0/{num_epochs})", position=0)
    for epoch in epochs:
        train_start = time.time()
        train_metrics = []

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

        
        train_loader = glue_train_data_collator(input_rng, train_dataset, train_batch_size)
        for step, batch in enumerate(
            tqdm(
                train_loader,
                total=steps_per_epoch,
                desc="Training...",
                position=1,
            ),
        ):
            state, train_metric, dropout_rngs = p_train_step(state, batch, dropout_rngs)
            train_metrics.append(train_metric)

            cur_step = (epoch * steps_per_epoch) + (step + 1)

            if cur_step % training_args.logging_steps == 0 and cur_step > 0:
                
                train_metric = 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}/{total_steps} | Training Loss: {train_metric['loss']}, Learning Rate:"
                    f" {train_metric['learning_rate']})"
                )

                train_metrics = []

            if (cur_step % training_args.eval_steps == 0 or cur_step % steps_per_epoch == 0) and cur_step > 0:
                
                eval_loader = glue_eval_data_collator(eval_dataset, eval_batch_size)
                for batch in tqdm(
                    eval_loader,
                    total=math.ceil(len(eval_dataset) / eval_batch_size),
                    desc="Evaluating ...",
                    position=2,
                ):
                    labels = batch.pop("labels")
                    predictions = pad_shard_unpad(p_eval_step)(
                        state, batch, min_device_batch=per_device_eval_batch_size
                    )
                    metric.add_batch(predictions=np.array(predictions), references=labels)

                eval_metric = metric.compute()

                logger.info(f"Step... ({cur_step}/{total_steps} | Eval metrics: {eval_metric})")

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

            if (cur_step % training_args.save_steps == 0 and cur_step > 0) or (cur_step == total_steps):
                
                if jax.process_index() == 0:
                    params = jax.device_get(unreplicate(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)
            epochs.desc = f"Epoch ... {epoch + 1}/{num_epochs}"

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


if __name__ == "__main__":
    main()