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 transformers model for sequence classification on glue import argparse import json import logging import math import os import random from pathlib import path import datasets import evaluate import torch from accelerate import accelerator from accelerate logging import getlogger from accelerate utils import setseed from datasets import loaddataset from huggingfacehub import repository createrepo from torch utils data import dataloader from tqdm auto import tqdm import transformers from transformers import autoconfig automodelforsequenceclassification autotokenizer datacollatorwithpadding pretrainedconfig schedulertype defaultdatacollator getscheduler from transformers utils import checkminversion sendexampletelemetry from transformers utils versions import requireversion will error if the minimal version of transformers is not installed remove at your own risks checkminversion4 36 0 dev0 logger getloggername requireversiondatasets1 8 0 to fix pip install r examplespytorchtextclassificationrequirements txt 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 def parseargs parser argparse argumentparserdescriptionfinetune a transformers model on a text classification task parser addargument taskname typestr defaultnone helpthe name of the glue task to train on choiceslisttasktokeys keys parser addargument trainfile typestr defaultnone helpa csv or a json file containing the training data parser addargument validationfile typestr defaultnone helpa csv or a json file containing the validation data parser addargument maxlength typeint default128 help the maximum total input sequence length after tokenization sequences longer than this will be truncated sequences shorter will be padded if padtomaxlength is passed parser addargument padtomaxlength actionstoretrue helpif passed pad all samples to maxlength otherwise dynamic padding is used parser addargument modelnameorpath typestr helppath to pretrained model or model identifier from huggingface comodels requiredtrue parser addargument useslowtokenizer actionstoretrue helpif passed will use a slow tokenizer not backed by the tokenizers library parser addargument perdevicetrainbatchsize typeint default8 helpbatch size per device for the training dataloader parser addargument perdeviceevalbatchsize typeint default8 helpbatch size per device for the evaluation dataloader parser addargument learningrate typefloat default5e5 helpinitial learning rate after the potential warmup period to use parser addargumentweightdecay typefloat default0 0 helpweight decay to use parser addargumentnumtrainepochs typeint default3 helptotal number of training epochs to perform parser addargument maxtrainsteps typeint defaultnone helptotal number of training steps to perform if provided overrides numtrainepochs parser addargument gradientaccumulationsteps typeint default1 helpnumber of updates steps to accumulate before performing a backwardupdate pass parser addargument lrschedulertype typeschedulertype defaultlinear helpthe scheduler type to use choiceslinear cosine cosinewithrestarts polynomial constant constantwithwarmup parser addargument numwarmupsteps typeint default0 helpnumber of steps for the warmup in the lr scheduler parser addargumentoutputdir typestr defaultnone helpwhere to store the final model parser addargumentseed typeint defaultnone helpa seed for reproducible training parser addargumentpushtohub actionstoretrue helpwhether or not to push the model to the hub parser addargument hubmodelid typestr helpthe name of the repository to keep in sync with the local outputdir parser addargumenthubtoken typestr helpthe token to use to push to the model hub parser addargument trustremotecode typebool defaultfalse 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 parser addargument checkpointingsteps typestr defaultnone helpwhether the various states should be saved at the end of every n steps or epoch for each epoch parser addargument resumefromcheckpoint typestr defaultnone helpif the training should continue from a checkpoint folder parser addargument withtracking actionstoretrue helpwhether to enable experiment trackers for logging parser addargument reportto typestr defaultall help the integration to report the results and logs to supported platforms are tensorboard wandb cometml and clearml use all default to report to all integrations only applicable when withtracking is passed parser addargument ignoremismatchedsizes actionstoretrue helpwhether or not to enable to load a pretrained model whose head dimensions are different args parser parseargs sanity checks if args taskname is none and args trainfile is none and args validationfile is none raise valueerrorneed either a task name or a trainingvalidation file else if args trainfile is not none extension args trainfile split 1 assert extension in csv json trainfile should be a csv or a json file if args validationfile is not none extension args validationfile split 1 assert extension in csv json validationfile should be a csv or a json file if args pushtohub assert args outputdir is not none need an outputdir to create a repo when pushtohub is passed return args def main args parseargs 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 sendexampletelemetryrungluenotrainer args initialize the accelerator we will let the accelerator handle device placement for us in this example if we re using tracking we also need to initialize it here and it will by default pick up all supported trackers in the environment accelerator acceleratorlogwithargs reportto projectdirargs outputdir if args withtracking else accelerator make one log on every process with the configuration for debugging logging basicconfig formatasctimes levelnames names messages datefmtmdy h m s levellogging info logger infoaccelerator state mainprocessonlyfalse if accelerator islocalmainprocess datasets utils logging setverbositywarning transformers utils logging setverbosityinfo else datasets utils logging setverbosityerror transformers utils logging setverbosityerror if passed along set the training seed now if args seed is not none setseedargs seed handle the repository creation if accelerator ismainprocess if args pushtohub retrieve of infer reponame reponame args hubmodelid if reponame is none reponame pathargs outputdir absolute name create repo and retrieve repoid repoid createreporeponame existoktrue tokenargs hubtoken repoid clone repo locally repo repositoryargs outputdir clonefromrepoid tokenargs hubtoken with openos path joinargs outputdir gitignore w as gitignore if step not in gitignore gitignore writestepn if epoch not in gitignore gitignore writeepochn elif args outputdir is not none os makedirsargs outputdir existoktrue accelerator waitforeveryone 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 args taskname is not none downloading and loading a dataset from the hub rawdatasets loaddatasetglue args taskname else loading the dataset from local csv or json file datafiles if args trainfile is not none datafilestrain args trainfile if args validationfile is not none datafilesvalidation args validationfile extension args trainfile if args trainfile is not none else args validationfile split 1 rawdatasets loaddatasetextension datafilesdatafiles see more about loading any type of standard or custom dataset at https huggingface codocsdatasetsloadingdatasets labels if args taskname is not none isregression args 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 codocsdatasetspackagereferencemainclasses htmldatasets dataset unique labellist rawdatasetstrain uniquelabel labellist sort let s sort it for determinism numlabels lenlabellist load pretrained model and tokenizer in distributed training the frompretrained methods guarantee that only one local process can concurrently download model vocab config autoconfig frompretrained args modelnameorpath numlabelsnumlabels finetuningtaskargs taskname trustremotecodeargs trustremotecode tokenizer autotokenizer frompretrained args modelnameorpath usefastnot args useslowtokenizer trustremotecodeargs trustremotecode model automodelforsequenceclassification frompretrained args modelnameorpath fromtfbool ckpt in args modelnameorpath configconfig ignoremismatchedsizesargs ignoremismatchedsizes trustremotecodeargs trustremotecode preprocessing the datasets if args taskname is not none sentence1key sentence2key tasktokeysargs 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 args 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 args taskname is none and not isregression labeltoid v i for i v in enumeratelabellist if labeltoid is not none model config label2id labeltoid model config id2label id label for label id in config label2id items elif args taskname is not none and not isregression model config label2id l i for i l in enumeratelabellist model config id2label id label for label id in config label2id items padding maxlength if args padtomaxlength else false def preprocessfunctionexamples tokenize the texts texts examplessentence1key if sentence2key is none else examplessentence1key examplessentence2key result tokenizertexts paddingpadding maxlengthargs maxlength 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 with accelerator mainprocessfirst processeddatasets rawdatasets map preprocessfunction batchedtrue removecolumnsrawdatasetstrain columnnames descrunning tokenizer on dataset traindataset processeddatasetstrain evaldataset processeddatasetsvalidationmatched if args 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 dataloaders creation if args padtomaxlength if padding was already done ot max length we use the default data collator that will just convert everything to tensors datacollator defaultdatacollator else otherwise datacollatorwithpadding will apply dynamic padding for us by padding to the maximum length of the samples passed when using mixed precision we add padtomultipleof8 to pad all tensors to multiple of 8s which will enable the use of tensor cores on nvidia hardware with compute capability 7 5 volta datacollator datacollatorwithpaddingtokenizer padtomultipleof8 if accelerator usefp16 else none traindataloader dataloader traindataset shuffletrue collatefndatacollator batchsizeargs perdevicetrainbatchsize evaldataloader dataloaderevaldataset collatefndatacollator batchsizeargs perdeviceevalbatchsize optimizer split weights in two groups one with weight decay and the other not nodecay bias layernorm weight optimizergroupedparameters params p for n p in model namedparameters if not anynd in n for nd in nodecay weightdecay args weightdecay params p for n p in model namedparameters if anynd in n for nd in nodecay weightdecay 0 0 optimizer torch optim adamwoptimizergroupedparameters lrargs learningrate scheduler and math around the number of training steps overrodemaxtrainsteps false numupdatestepsperepoch math ceillentraindataloader args gradientaccumulationsteps if args maxtrainsteps is none args maxtrainsteps args numtrainepochs numupdatestepsperepoch overrodemaxtrainsteps true lrscheduler getscheduler nameargs lrschedulertype optimizeroptimizer numwarmupstepsargs numwarmupsteps numtrainingstepsargs maxtrainsteps prepare everything with our accelerator model optimizer traindataloader evaldataloader lrscheduler accelerator prepare model optimizer traindataloader evaldataloader lrscheduler we need to recalculate our total training steps as the size of the training dataloader may have changed numupdatestepsperepoch math ceillentraindataloader args gradientaccumulationsteps if overrodemaxtrainsteps args maxtrainsteps args numtrainepochs numupdatestepsperepoch afterwards we recalculate our number of training epochs args numtrainepochs math ceilargs maxtrainsteps numupdatestepsperepoch figure out how many steps we should save the accelerator states checkpointingsteps args checkpointingsteps if checkpointingsteps is not none and checkpointingsteps isdigit checkpointingsteps intcheckpointingsteps we need to initialize the trackers we use and also store our configuration the trackers initializes automatically on the main process if args withtracking experimentconfig varsargs tensorboard cannot log enums need the raw value experimentconfiglrschedulertype experimentconfiglrschedulertype value accelerator inittrackersgluenotrainer experimentconfig get the metric function if args taskname is not none metric evaluate loadglue args taskname else metric evaluate loadaccuracy train totalbatchsize args perdevicetrainbatchsize accelerator numprocesses args gradientaccumulationsteps logger info running training logger infof num examples lentraindataset logger infof num epochs args numtrainepochs logger infof instantaneous batch size per device args perdevicetrainbatchsize logger infof total train batch size w parallel distributed accumulation totalbatchsize logger infof gradient accumulation steps args gradientaccumulationsteps logger infof total optimization steps args maxtrainsteps only show the progress bar once on each machine progressbar tqdmrangeargs maxtrainsteps disablenot accelerator islocalmainprocess completedsteps 0 startingepoch 0 potentially load in the weights and states from a previous save if args resumefromcheckpoint if args resumefromcheckpoint is not none or args resumefromcheckpoint checkpointpath args resumefromcheckpoint path os path basenameargs resumefromcheckpoint else get the most recent checkpoint dirs f name for f in os scandiros getcwd if f isdir dirs sortkeyos path getctime path dirs1 sorts folders by date modified most recent checkpoint is the last checkpointpath path path os path basenamecheckpointpath accelerator printfresumed from checkpoint checkpointpath accelerator loadstatecheckpointpath extract epochi or stepi trainingdifference os path splitextpath0 if epoch in trainingdifference startingepoch inttrainingdifference replaceepoch 1 resumestep none completedsteps startingepoch numupdatestepsperepoch else need to multiply gradientaccumulationsteps to reflect real steps resumestep inttrainingdifference replacestep args gradientaccumulationsteps startingepoch resumestep lentraindataloader completedsteps resumestep args gradientaccumulationsteps resumestep startingepoch lentraindataloader update the progressbar if load from checkpoint progressbar updatecompletedsteps for epoch in rangestartingepoch args numtrainepochs model train if args withtracking totalloss 0 if args resumefromcheckpoint and epoch startingepoch and resumestep is not none we skip the first n batches in the dataloader when resuming from a checkpoint activedataloader accelerator skipfirstbatchestraindataloader resumestep else activedataloader traindataloader for step batch in enumerateactivedataloader outputs modelbatch loss outputs loss we keep track of the loss at each epoch if args withtracking totalloss loss detach float loss loss args gradientaccumulationsteps accelerator backwardloss if step args gradientaccumulationsteps 0 or step lentraindataloader 1 optimizer step lrscheduler step optimizer zerograd progressbar update1 completedsteps 1 if isinstancecheckpointingsteps int if completedsteps checkpointingsteps 0 outputdir fstepcompletedsteps if args outputdir is not none outputdir os path joinargs outputdir outputdir accelerator savestateoutputdir if completedsteps args maxtrainsteps break model eval samplesseen 0 for step batch in enumerateevaldataloader with torch nograd outputs modelbatch predictions outputs logits argmaxdim1 if not isregression else outputs logits squeeze predictions references accelerator gatherpredictions batchlabels if we are in a multiprocess environment the last batch has duplicates if accelerator numprocesses 1 if step lenevaldataloader 1 predictions predictions lenevaldataloader dataset samplesseen references references lenevaldataloader dataset samplesseen else samplesseen references shape0 metric addbatch predictionspredictions referencesreferences evalmetric metric compute logger infofepoch epoch evalmetric if args withtracking accelerator log accuracy if args taskname is not none else glue evalmetric trainloss totalloss item lentraindataloader epoch epoch step completedsteps stepcompletedsteps if args pushtohub and epoch args numtrainepochs 1 accelerator waitforeveryone unwrappedmodel accelerator unwrapmodelmodel unwrappedmodel savepretrained args outputdir ismainprocessaccelerator ismainprocess savefunctionaccelerator save if accelerator ismainprocess tokenizer savepretrainedargs outputdir repo pushtohub commitmessageftraining in progress epoch epoch blockingfalse autolfsprunetrue if args checkpointingsteps epoch outputdir fepochepoch if args outputdir is not none outputdir os path joinargs outputdir outputdir accelerator savestateoutputdir if args withtracking accelerator endtraining if args outputdir is not none accelerator waitforeveryone unwrappedmodel accelerator unwrapmodelmodel unwrappedmodel savepretrained args outputdir ismainprocessaccelerator ismainprocess savefunctionaccelerator save if accelerator ismainprocess tokenizer savepretrainedargs outputdir if args pushtohub repo pushtohubcommitmessageend of training autolfsprunetrue if args taskname mnli final evaluation on mismatched validation set evaldataset processeddatasetsvalidationmismatched evaldataloader dataloader evaldataset collatefndatacollator batchsizeargs perdeviceevalbatchsize evaldataloader accelerator prepareevaldataloader model eval for step batch in enumerateevaldataloader outputs modelbatch predictions outputs logits argmaxdim1 metric addbatch predictionsaccelerator gatherpredictions referencesaccelerator gatherbatchlabels evalmetric metric compute logger infofmnlimm evalmetric if args outputdir is not none allresults fevalk v for k v in evalmetric items with openos path joinargs outputdir allresults json w as f json dumpallresults f if name main main 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 transformers model for sequence classification on glue will error if the minimal version of transformers is not installed remove at your own risks sanity checks 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 initialize the accelerator we will let the accelerator handle device placement for us in this example if we re using tracking we also need to initialize it here and it will by default pick up all supported trackers in the environment make one log on every process with the configuration for debugging if passed along set the training seed now 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 html datasets dataset unique let s sort it for determinism load pretrained model and tokenizer in distributed training the from_pretrained methods guarantee that only one local process can concurrently download model vocab 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 dataloaders creation if padding was already done ot max length we use the default data collator that will just convert everything to tensors otherwise datacollatorwithpadding will apply dynamic padding for us by padding to the maximum length of the samples passed when using mixed precision we add pad_to_multiple_of 8 to pad all tensors to multiple of 8s which will enable the use of tensor cores on nvidia hardware with compute capability 7 5 volta optimizer split weights in two groups one with weight decay and the other not scheduler and math around the number of training steps prepare everything with our accelerator we need to recalculate our total training steps as the size of the training dataloader may have changed afterwards we recalculate our number of training epochs figure out how many steps we should save the accelerator states we need to initialize the trackers we use and also store our configuration the trackers initializes automatically on the main process tensorboard cannot log enums need the raw value get the metric function train only show the progress bar once on each machine potentially load in the weights and states from a previous save get the most recent checkpoint sorts folders by date modified most recent checkpoint is the last extract epoch_ i or step_ i need to multiply gradient_accumulation_steps to reflect real steps update the progress_bar if load from checkpoint we skip the first n batches in the dataloader when resuming from a checkpoint we keep track of the loss at each epoch if we are in a multiprocess environment the last batch has duplicates final evaluation on mismatched validation set
import argparse
import json
import logging
import math
import os
import random
from pathlib import Path

import datasets
import evaluate
import torch
from accelerate import Accelerator
from accelerate.logging import get_logger
from accelerate.utils import set_seed
from datasets import load_dataset
from huggingface_hub import Repository, create_repo
from torch.utils.data import DataLoader
from tqdm.auto import tqdm

import transformers
from transformers import (
    AutoConfig,
    AutoModelForSequenceClassification,
    AutoTokenizer,
    DataCollatorWithPadding,
    PretrainedConfig,
    SchedulerType,
    default_data_collator,
    get_scheduler,
)
from transformers.utils import check_min_version, send_example_telemetry
from transformers.utils.versions import require_version



check_min_version("4.36.0.dev0")

logger = get_logger(__name__)

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

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"),
}


def parse_args():
    parser = argparse.ArgumentParser(description="Finetune a transformers model on a text classification task")
    parser.add_argument(
        "--task_name",
        type=str,
        default=None,
        help="The name of the glue task to train on.",
        choices=list(task_to_keys.keys()),
    )
    parser.add_argument(
        "--train_file", type=str, default=None, help="A csv or a json file containing the training data."
    )
    parser.add_argument(
        "--validation_file", type=str, default=None, help="A csv or a json file containing the validation data."
    )
    parser.add_argument(
        "--max_length",
        type=int,
        default=128,
        help=(
            "The maximum total input sequence length after tokenization. Sequences longer than this will be truncated,"
            " sequences shorter will be padded if `--pad_to_max_length` is passed."
        ),
    )
    parser.add_argument(
        "--pad_to_max_length",
        action="store_true",
        help="If passed, pad all samples to `max_length`. Otherwise, dynamic padding is used.",
    )
    parser.add_argument(
        "--model_name_or_path",
        type=str,
        help="Path to pretrained model or model identifier from huggingface.co/models.",
        required=True,
    )
    parser.add_argument(
        "--use_slow_tokenizer",
        action="store_true",
        help="If passed, will use a slow tokenizer (not backed by the 🤗 Tokenizers library).",
    )
    parser.add_argument(
        "--per_device_train_batch_size",
        type=int,
        default=8,
        help="Batch size (per device) for the training dataloader.",
    )
    parser.add_argument(
        "--per_device_eval_batch_size",
        type=int,
        default=8,
        help="Batch size (per device) for the evaluation dataloader.",
    )
    parser.add_argument(
        "--learning_rate",
        type=float,
        default=5e-5,
        help="Initial learning rate (after the potential warmup period) to use.",
    )
    parser.add_argument("--weight_decay", type=float, default=0.0, help="Weight decay to use.")
    parser.add_argument("--num_train_epochs", type=int, default=3, help="Total number of training epochs to perform.")
    parser.add_argument(
        "--max_train_steps",
        type=int,
        default=None,
        help="Total number of training steps to perform. If provided, overrides num_train_epochs.",
    )
    parser.add_argument(
        "--gradient_accumulation_steps",
        type=int,
        default=1,
        help="Number of updates steps to accumulate before performing a backward/update pass.",
    )
    parser.add_argument(
        "--lr_scheduler_type",
        type=SchedulerType,
        default="linear",
        help="The scheduler type to use.",
        choices=["linear", "cosine", "cosine_with_restarts", "polynomial", "constant", "constant_with_warmup"],
    )
    parser.add_argument(
        "--num_warmup_steps", type=int, default=0, help="Number of steps for the warmup in the lr scheduler."
    )
    parser.add_argument("--output_dir", type=str, default=None, help="Where to store the final model.")
    parser.add_argument("--seed", type=int, default=None, help="A seed for reproducible training.")
    parser.add_argument("--push_to_hub", action="store_true", help="Whether or not to push the model to the Hub.")
    parser.add_argument(
        "--hub_model_id", type=str, help="The name of the repository to keep in sync with the local `output_dir`."
    )
    parser.add_argument("--hub_token", type=str, help="The token to use to push to the Model Hub.")
    parser.add_argument(
        "--trust_remote_code",
        type=bool,
        default=False,
        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."
        ),
    )
    parser.add_argument(
        "--checkpointing_steps",
        type=str,
        default=None,
        help="Whether the various states should be saved at the end of every n steps, or 'epoch' for each epoch.",
    )
    parser.add_argument(
        "--resume_from_checkpoint",
        type=str,
        default=None,
        help="If the training should continue from a checkpoint folder.",
    )
    parser.add_argument(
        "--with_tracking",
        action="store_true",
        help="Whether to enable experiment trackers for logging.",
    )
    parser.add_argument(
        "--report_to",
        type=str,
        default="all",
        help=(
            'The integration to report the results and logs to. Supported platforms are `"tensorboard"`,'
            ' `"wandb"`, `"comet_ml"` and `"clearml"`. Use `"all"` (default) to report to all integrations. '
            "Only applicable when `--with_tracking` is passed."
        ),
    )
    parser.add_argument(
        "--ignore_mismatched_sizes",
        action="store_true",
        help="Whether or not to enable to load a pretrained model whose head dimensions are different.",
    )
    args = parser.parse_args()

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

    if args.push_to_hub:
        assert args.output_dir is not None, "Need an `output_dir` to create a repo when `--push_to_hub` is passed."

    return args


def main():
    args = parse_args()
    
    
    send_example_telemetry("run_glue_no_trainer", args)

    
    
    
    accelerator = (
        Accelerator(log_with=args.report_to, project_dir=args.output_dir) if args.with_tracking else Accelerator()
    )
    
    logging.basicConfig(
        format="%(asctime)s - %(levelname)s - %(name)s - %(message)s",
        datefmt="%m/%d/%Y %H:%M:%S",
        level=logging.INFO,
    )
    logger.info(accelerator.state, main_process_only=False)
    if accelerator.is_local_main_process:
        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 args.seed is not None:
        set_seed(args.seed)

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

            with open(os.path.join(args.output_dir, ".gitignore"), "w+") as gitignore:
                if "step_*" not in gitignore:
                    gitignore.write("step_*\n")
                if "epoch_*" not in gitignore:
                    gitignore.write("epoch_*\n")
        elif args.output_dir is not None:
            os.makedirs(args.output_dir, exist_ok=True)
    accelerator.wait_for_everyone()

    
    

    
    
    

    
    

    
    
    if args.task_name is not None:
        
        raw_datasets = load_dataset("glue", args.task_name)
    else:
        
        data_files = {}
        if args.train_file is not None:
            data_files["train"] = args.train_file
        if args.validation_file is not None:
            data_files["validation"] = args.validation_file
        extension = (args.train_file if args.train_file is not None else args.validation_file).split(".")[-1]
        raw_datasets = load_dataset(extension, data_files=data_files)
    
    

    
    if args.task_name is not None:
        is_regression = 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(
        args.model_name_or_path,
        num_labels=num_labels,
        finetuning_task=args.task_name,
        trust_remote_code=args.trust_remote_code,
    )
    tokenizer = AutoTokenizer.from_pretrained(
        args.model_name_or_path, use_fast=not args.use_slow_tokenizer, trust_remote_code=args.trust_remote_code
    )
    model = AutoModelForSequenceClassification.from_pretrained(
        args.model_name_or_path,
        from_tf=bool(".ckpt" in args.model_name_or_path),
        config=config,
        ignore_mismatched_sizes=args.ignore_mismatched_sizes,
        trust_remote_code=args.trust_remote_code,
    )

    
    if args.task_name is not None:
        sentence1_key, sentence2_key = task_to_keys[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 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 args.task_name is None and not is_regression:
        label_to_id = {v: i for i, v in enumerate(label_list)}

    if label_to_id is not None:
        model.config.label2id = label_to_id
        model.config.id2label = {id: label for label, id in config.label2id.items()}
    elif args.task_name is not None and not is_regression:
        model.config.label2id = {l: i for i, l in enumerate(label_list)}
        model.config.id2label = {id: label for label, id in config.label2id.items()}

    padding = "max_length" if args.pad_to_max_length else False

    def preprocess_function(examples):
        
        texts = (
            (examples[sentence1_key],) if sentence2_key is None else (examples[sentence1_key], examples[sentence2_key])
        )
        result = tokenizer(*texts, padding=padding, max_length=args.max_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

    with accelerator.main_process_first():
        processed_datasets = raw_datasets.map(
            preprocess_function,
            batched=True,
            remove_columns=raw_datasets["train"].column_names,
            desc="Running tokenizer on dataset",
        )

    train_dataset = processed_datasets["train"]
    eval_dataset = processed_datasets["validation_matched" if 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]}.")

    
    if args.pad_to_max_length:
        
        
        data_collator = default_data_collator
    else:
        
        
        
        data_collator = DataCollatorWithPadding(tokenizer, pad_to_multiple_of=(8 if accelerator.use_fp16 else None))

    train_dataloader = DataLoader(
        train_dataset, shuffle=True, collate_fn=data_collator, batch_size=args.per_device_train_batch_size
    )
    eval_dataloader = DataLoader(eval_dataset, collate_fn=data_collator, batch_size=args.per_device_eval_batch_size)

    
    
    no_decay = ["bias", "LayerNorm.weight"]
    optimizer_grouped_parameters = [
        {
            "params": [p for n, p in model.named_parameters() if not any(nd in n for nd in no_decay)],
            "weight_decay": args.weight_decay,
        },
        {
            "params": [p for n, p in model.named_parameters() if any(nd in n for nd in no_decay)],
            "weight_decay": 0.0,
        },
    ]
    optimizer = torch.optim.AdamW(optimizer_grouped_parameters, lr=args.learning_rate)

    
    overrode_max_train_steps = False
    num_update_steps_per_epoch = math.ceil(len(train_dataloader) / args.gradient_accumulation_steps)
    if args.max_train_steps is None:
        args.max_train_steps = args.num_train_epochs * num_update_steps_per_epoch
        overrode_max_train_steps = True

    lr_scheduler = get_scheduler(
        name=args.lr_scheduler_type,
        optimizer=optimizer,
        num_warmup_steps=args.num_warmup_steps,
        num_training_steps=args.max_train_steps,
    )

    
    model, optimizer, train_dataloader, eval_dataloader, lr_scheduler = accelerator.prepare(
        model, optimizer, train_dataloader, eval_dataloader, lr_scheduler
    )

    
    num_update_steps_per_epoch = math.ceil(len(train_dataloader) / args.gradient_accumulation_steps)
    if overrode_max_train_steps:
        args.max_train_steps = args.num_train_epochs * num_update_steps_per_epoch
    
    args.num_train_epochs = math.ceil(args.max_train_steps / num_update_steps_per_epoch)

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

    
    
    if args.with_tracking:
        experiment_config = vars(args)
        
        experiment_config["lr_scheduler_type"] = experiment_config["lr_scheduler_type"].value
        accelerator.init_trackers("glue_no_trainer", experiment_config)

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

    
    total_batch_size = args.per_device_train_batch_size * accelerator.num_processes * args.gradient_accumulation_steps

    logger.info("***** Running training *****")
    logger.info(f"  Num examples = {len(train_dataset)}")
    logger.info(f"  Num Epochs = {args.num_train_epochs}")
    logger.info(f"  Instantaneous batch size per device = {args.per_device_train_batch_size}")
    logger.info(f"  Total train batch size (w. parallel, distributed & accumulation) = {total_batch_size}")
    logger.info(f"  Gradient Accumulation steps = {args.gradient_accumulation_steps}")
    logger.info(f"  Total optimization steps = {args.max_train_steps}")
    
    progress_bar = tqdm(range(args.max_train_steps), disable=not accelerator.is_local_main_process)
    completed_steps = 0
    starting_epoch = 0
    
    if args.resume_from_checkpoint:
        if args.resume_from_checkpoint is not None or args.resume_from_checkpoint != "":
            checkpoint_path = args.resume_from_checkpoint
            path = os.path.basename(args.resume_from_checkpoint)
        else:
            
            dirs = [f.name for f in os.scandir(os.getcwd()) if f.is_dir()]
            dirs.sort(key=os.path.getctime)
            path = dirs[-1]  
            checkpoint_path = path
            path = os.path.basename(checkpoint_path)

        accelerator.print(f"Resumed from checkpoint: {checkpoint_path}")
        accelerator.load_state(checkpoint_path)
        
        training_difference = os.path.splitext(path)[0]

        if "epoch" in training_difference:
            starting_epoch = int(training_difference.replace("epoch_", "")) + 1
            resume_step = None
            completed_steps = starting_epoch * num_update_steps_per_epoch
        else:
            
            resume_step = int(training_difference.replace("step_", "")) * args.gradient_accumulation_steps
            starting_epoch = resume_step // len(train_dataloader)
            completed_steps = resume_step // args.gradient_accumulation_steps
            resume_step -= starting_epoch * len(train_dataloader)

    
    progress_bar.update(completed_steps)

    for epoch in range(starting_epoch, args.num_train_epochs):
        model.train()
        if args.with_tracking:
            total_loss = 0
        if args.resume_from_checkpoint and epoch == starting_epoch and resume_step is not None:
            
            active_dataloader = accelerator.skip_first_batches(train_dataloader, resume_step)
        else:
            active_dataloader = train_dataloader
        for step, batch in enumerate(active_dataloader):
            outputs = model(**batch)
            loss = outputs.loss
            
            if args.with_tracking:
                total_loss += loss.detach().float()
            loss = loss / args.gradient_accumulation_steps
            accelerator.backward(loss)
            if step % args.gradient_accumulation_steps == 0 or step == len(train_dataloader) - 1:
                optimizer.step()
                lr_scheduler.step()
                optimizer.zero_grad()
                progress_bar.update(1)
                completed_steps += 1

            if isinstance(checkpointing_steps, int):
                if completed_steps % checkpointing_steps == 0:
                    output_dir = f"step_{completed_steps}"
                    if args.output_dir is not None:
                        output_dir = os.path.join(args.output_dir, output_dir)
                    accelerator.save_state(output_dir)

            if completed_steps >= args.max_train_steps:
                break

        model.eval()
        samples_seen = 0
        for step, batch in enumerate(eval_dataloader):
            with torch.no_grad():
                outputs = model(**batch)
            predictions = outputs.logits.argmax(dim=-1) if not is_regression else outputs.logits.squeeze()
            predictions, references = accelerator.gather((predictions, batch["labels"]))
            
            if accelerator.num_processes > 1:
                if step == len(eval_dataloader) - 1:
                    predictions = predictions[: len(eval_dataloader.dataset) - samples_seen]
                    references = references[: len(eval_dataloader.dataset) - samples_seen]
                else:
                    samples_seen += references.shape[0]
            metric.add_batch(
                predictions=predictions,
                references=references,
            )

        eval_metric = metric.compute()
        logger.info(f"epoch {epoch}: {eval_metric}")

        if args.with_tracking:
            accelerator.log(
                {
                    "accuracy" if args.task_name is not None else "glue": eval_metric,
                    "train_loss": total_loss.item() / len(train_dataloader),
                    "epoch": epoch,
                    "step": completed_steps,
                },
                step=completed_steps,
            )

        if args.push_to_hub and epoch < args.num_train_epochs - 1:
            accelerator.wait_for_everyone()
            unwrapped_model = accelerator.unwrap_model(model)
            unwrapped_model.save_pretrained(
                args.output_dir, is_main_process=accelerator.is_main_process, save_function=accelerator.save
            )
            if accelerator.is_main_process:
                tokenizer.save_pretrained(args.output_dir)
                repo.push_to_hub(
                    commit_message=f"Training in progress epoch {epoch}", blocking=False, auto_lfs_prune=True
                )

        if args.checkpointing_steps == "epoch":
            output_dir = f"epoch_{epoch}"
            if args.output_dir is not None:
                output_dir = os.path.join(args.output_dir, output_dir)
            accelerator.save_state(output_dir)

    if args.with_tracking:
        accelerator.end_training()

    if args.output_dir is not None:
        accelerator.wait_for_everyone()
        unwrapped_model = accelerator.unwrap_model(model)
        unwrapped_model.save_pretrained(
            args.output_dir, is_main_process=accelerator.is_main_process, save_function=accelerator.save
        )
        if accelerator.is_main_process:
            tokenizer.save_pretrained(args.output_dir)
            if args.push_to_hub:
                repo.push_to_hub(commit_message="End of training", auto_lfs_prune=True)

    if args.task_name == "mnli":
        
        eval_dataset = processed_datasets["validation_mismatched"]
        eval_dataloader = DataLoader(
            eval_dataset, collate_fn=data_collator, batch_size=args.per_device_eval_batch_size
        )
        eval_dataloader = accelerator.prepare(eval_dataloader)

        model.eval()
        for step, batch in enumerate(eval_dataloader):
            outputs = model(**batch)
            predictions = outputs.logits.argmax(dim=-1)
            metric.add_batch(
                predictions=accelerator.gather(predictions),
                references=accelerator.gather(batch["labels"]),
            )

        eval_metric = metric.compute()
        logger.info(f"mnli-mm: {eval_metric}")

    if args.output_dir is not None:
        all_results = {f"eval_{k}": v for k, v in eval_metric.items()}
        with open(os.path.join(args.output_dir, "all_results.json"), "w") as f:
            json.dump(all_results, f)


if __name__ == "__main__":
    main()