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 the library models for sequence classification you can also adapt this script on your own text classification task pointers for this are left as comments import json import logging import os import sys import warnings from dataclasses import dataclass field from pathlib import path from typing import optional import numpy as np from datasets import loaddataset from transformers import autoconfig autotokenizer hfargumentparser pretrainedconfig pushtohubcallback tfautomodelforsequenceclassification tftrainingarguments createoptimizer setseed from transformers utils import configname tf2weightsname sendexampletelemetry os environtfcppminloglevel 1 reduce the amount of console output from tf import tensorflow as tf noqa e402 logger logging getloggername region helper classes class savepretrainedcallbacktf keras callbacks callback hugging face models have a savepretrained method that saves both the weights and the necessary metadata to allow them to be loaded as a pretrained model in future this is a simple keras callback that saves the model with this method after each epoch def initself outputdir kwargs super init self outputdir outputdir def onepochendself epoch logsnone self model savepretrainedself outputdir endregion region commandline arguments dataclass class datatrainingarguments trainfile optionalstr field defaultnone metadatahelp a csv or a json file containing the training data validationfile optionalstr field defaultnone metadatahelp a csv or a json file containing the validation data testfile optionalstr fielddefaultnone metadatahelp a csv or a json file containing the test data maxseqlength int field default128 metadata help the maximum total input sequence length after tokenization sequences longer than this will be truncated sequences shorter will be padded overwritecache bool field defaultfalse metadatahelp overwrite the cached preprocessed datasets or not padtomaxlength bool field defaultfalse metadata help whether to pad all samples to maxseqlength if false will pad the samples dynamically when batching to the maximum length in the batch data will always be padded when using tpus maxtrainsamples optionalint field defaultnone metadata help for debugging purposes or quicker training truncate the number of training examples to this value if set maxvalsamples optionalint field defaultnone metadata help for debugging purposes or quicker training truncate the number of validation examples to this value if set maxtestsamples optionalint field defaultnone metadata help for debugging purposes or quicker training truncate the number of test examples to this value if set def postinitself trainextension self trainfile split 1 lower if self trainfile is not none else none validationextension self validationfile split 1 lower if self validationfile is not none else none testextension self testfile split 1 lower if self testfile is not none else none extensions trainextension validationextension testextension extensions discardnone assert lenextensions 0 need to supply at least one of trainfile validationfile or testfile assert lenextensions 1 all input files should have the same file extension either csv or json assert csv in extensions or json in extensions input files should have either csv or json extensions self inputfileextension extensions pop 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 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 endregion def main region argument parsing see all possible arguments in srctransformerstrainingargs py or by passing the help flag to this script we now keep distinct sets of args for a cleaner separation of concerns parser hfargumentparsermodelarguments datatrainingarguments tftrainingarguments 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 sendexampletelemetryruntextclassification modelargs dataargs frameworktensorflow outputdir pathtrainingargs outputdir outputdir mkdirparentstrue existoktrue endregion region checkpoints detecting last checkpoint checkpoint none if lenos listdirtrainingargs outputdir 0 and not trainingargs overwriteoutputdir if outputdir configname isfile and outputdir tf2weightsname isfile checkpoint outputdir logger info fcheckpoint detected resuming training from checkpoint in trainingargs outputdir to avoid this behavior change the outputdir or add overwriteoutputdir to train from scratch else raise valueerror foutput directory trainingargs outputdir already exists and is not empty use overwriteoutputdir to continue regardless endregion region logging logging basicconfig formatasctimes levelnames names messages datefmtmdy h m s handlerslogging streamhandlersys stdout logger setlevellogging info logger infoftrainingevaluation parameters trainingargs endregion region loading data for csvjson files this script will use the label field as the label and the sentence1 and optionally sentence2 fields as inputs if they exist if not the first two fields not named label are used if at least two columns are provided note that the term sentence can be slightly misleading as they often contain more than a single grammatical sentence when the task requires it 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 datafiles train dataargs trainfile validation dataargs validationfile test dataargs testfile datafiles key file for key file in datafiles items if file is not none for key in datafiles keys logger infofloading a local file for key datafileskey if dataargs inputfileextension csv loading a dataset from local csv files datasets loaddataset csv datafilesdatafiles cachedirmodelargs cachedir tokenmodelargs token else loading a dataset from local json files datasets loaddatasetjson datafilesdatafiles cachedirmodelargs cachedir see more about loading any type of standard or custom dataset at https huggingface codocsdatasetsloadingdatasets endregion region label preprocessing if you ve passed us a training set we try to infer your labels from it if train in datasets by default we assume that if your label column looks like a float then you re doing regression and if not then you re doing classification this is something you may want to change isregression datasetstrain featureslabel dtype in float32 float64 if isregression numlabels 1 else a useful fast method https huggingface codocsdatasetspackagereferencemainclasses htmldatasets dataset unique labellist datasetstrain uniquelabel labellist sort let s sort it for determinism numlabels lenlabellist if you haven t passed a training set we read label info from the saved model this happens later else numlabels none labellist none isregression none endregion region load model config and tokenizer if checkpoint is not none configpath trainingargs outputdir elif modelargs configname configpath modelargs configname else configpath modelargs modelnameorpath if numlabels is not none config autoconfig frompretrained configpath numlabelsnumlabels cachedirmodelargs cachedir revisionmodelargs modelrevision tokenmodelargs token trustremotecodemodelargs trustremotecode else config autoconfig frompretrained configpath cachedirmodelargs cachedir revisionmodelargs modelrevision tokenmodelargs token trustremotecodemodelargs trustremotecode tokenizer autotokenizer frompretrained modelargs tokenizername if modelargs tokenizername else modelargs modelnameorpath cachedirmodelargs cachedir revisionmodelargs modelrevision tokenmodelargs token trustremotecodemodelargs trustremotecode endregion region dataset preprocessing again we try to have some nice defaults but don t hesitate to tweak to your use case columnnames col for cols in datasets columnnames values for col in cols nonlabelcolumnnames name for name in columnnames if name label if sentence1 in nonlabelcolumnnames and sentence2 in nonlabelcolumnnames sentence1key sentence2key sentence1 sentence2 elif sentence1 in nonlabelcolumnnames sentence1key sentence2key sentence1 none else if lennonlabelcolumnnames 2 sentence1key sentence2key nonlabelcolumnnames 2 else sentence1key sentence2key nonlabelcolumnnames0 none if dataargs maxseqlength tokenizer modelmaxlength logger warning fthe maxseqlength passed dataargs maxseqlength is larger than the maximum length for the fmodel tokenizer modelmaxlength using maxseqlengthtokenizer modelmaxlength maxseqlength mindataargs maxseqlength tokenizer modelmaxlength ensure that our labels match the model s if it has some prespecified if train in datasets if not isregression and config label2id pretrainedconfignumlabelsnumlabels label2id labelnametoid config label2id if sortedlabelnametoid keys sortedlabellist labeltoid labelnametoid use the model s labels 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 f sortedlabellist nignoring the model labels as a result labeltoid v i for i v in enumeratelabellist elif not isregression labeltoid v i for i v in enumeratelabellist else labeltoid none now we ve established our label2id let s overwrite the model config with it config label2id labeltoid if config label2id is not none config id2label id label for label id in labeltoid items else config id2label none else labeltoid config label2id just load the data from the model if validation in datasets and config label2id is not none validationlabellist datasetsvalidation uniquelabel for vallabel in validationlabellist assert vallabel in labeltoid flabel vallabel is in the validation set but not the training set def preprocessfunctionexamples tokenize the texts args examplessentence1key if sentence2key is none else examplessentence1key examplessentence2key result tokenizerargs maxlengthmaxseqlength truncationtrue map labels to ids if config label2id is not none and label in examples resultlabel config label2idl if l 1 else 1 for l in exampleslabel return result datasets datasets mappreprocessfunction batchedtrue loadfromcachefilenot dataargs overwritecache endregion with trainingargs strategy scope region load pretrained model set seed before initializing model setseedtrainingargs seed in distributed training the frompretrained methods guarantee that only one local process can concurrently download model vocab if checkpoint is none modelpath modelargs modelnameorpath else modelpath checkpoint model tfautomodelforsequenceclassification frompretrained modelpath configconfig cachedirmodelargs cachedir revisionmodelargs modelrevision tokenmodelargs token trustremotecodemodelargs trustremotecode endregion region convert data to a tf data dataset datasetoptions tf data options datasetoptions experimentaldistribute autoshardpolicy tf data experimental autoshardpolicy off numreplicas trainingargs strategy numreplicasinsync tfdata maxsamples train dataargs maxtrainsamples validation dataargs maxvalsamples test dataargs maxtestsamples for key in train validation test if key not in datasets tfdatakey none continue if key train and not trainingargs dotrain or key validation and not trainingargs doeval or key test and not trainingargs dopredict tfdatakey none continue if key in train validation assert label in datasetskey features fmissing labels from key data if key train shuffle true batchsize trainingargs perdevicetrainbatchsize numreplicas else shuffle false batchsize trainingargs perdeviceevalbatchsize numreplicas sampleslimit maxsampleskey dataset datasetskey if sampleslimit is not none dataset dataset selectrangesampleslimit model preparetfdataset wraps a hugging face dataset in a tf data dataset which is ready to use in training this is the recommended way to use a hugging face dataset when training with keras you can also use the lowerlevel dataset totfdataset method but you will have to specify things like column names yourself if you use this method whereas they are automatically inferred from the model input names when using model preparetfdataset for more info see the docs https huggingface codocstransformersmainenmainclassesmodeltransformers tfpretrainedmodel preparetfdataset https huggingface codocsdatasetsmainenpackagereferencemainclassesdatasets dataset totfdataset data model preparetfdataset dataset shuffleshuffle batchsizebatchsize tokenizertokenizer data data withoptionsdatasetoptions tfdatakey data endregion region optimizer loss and compilation if trainingargs dotrain numtrainsteps lentfdatatrain trainingargs numtrainepochs if trainingargs warmupsteps 0 numwarmupsteps trainingargs warmupsteps elif trainingargs warmupratio 0 numwarmupsteps intnumtrainsteps trainingargs warmupratio else numwarmupsteps 0 optimizer schedule createoptimizer initlrtrainingargs learningrate numtrainstepsnumtrainsteps numwarmupstepsnumwarmupsteps adambeta1trainingargs adambeta1 adambeta2trainingargs adambeta2 adamepsilontrainingargs adamepsilon weightdecayratetrainingargs weightdecay adamglobalclipnormtrainingargs maxgradnorm else optimizer none if isregression metrics else metrics accuracy transformers models compute the right loss for their task by default when labels are passed and will use this for training unless you specify your own loss function in compile model compileoptimizeroptimizer metricsmetrics endregion region preparing pushtohub and model card pushtohubmodelid trainingargs pushtohubmodelid modelname modelargs modelnameorpath split1 if not pushtohubmodelid pushtohubmodelid fmodelnamefinetunedtextclassification modelcardkwargs finetunedfrom modelargs modelnameorpath tasks textclassification if trainingargs pushtohub callbacks pushtohubcallback outputdirtrainingargs outputdir hubmodelidpushtohubmodelid hubtokentrainingargs pushtohubtoken tokenizertokenizer modelcardkwargs else callbacks endregion region training and validation if tfdatatrain is not none model fit tfdatatrain validationdatatfdatavalidation epochsinttrainingargs numtrainepochs callbackscallbacks if tfdatavalidation is not none logger infocomputing metrics on validation data if isregression loss model evaluatetfdatavalidation logger infofeval loss loss 5f else loss accuracy model evaluatetfdatavalidation logger infofeval loss loss 5f eval accuracy accuracy 100 4f if trainingargs outputdir is not none outputevalfile os path jointrainingargs outputdir allresults json evaldict evalloss loss if not isregression evaldictevalaccuracy accuracy with openoutputevalfile w as writer writer writejson dumpsevaldict endregion region prediction if tfdatatest is not none logger infodoing predictions on test dataset predictions model predicttfdatatestlogits predictedclass np squeezepredictions if isregression else np argmaxpredictions axis1 outputtestfile os path jointrainingargs outputdir testresults txt with openoutputtestfile w as writer writer writeindextpredictionn for index item in enumeratepredictedclass if isregression writer writefindextitem 3 3fn else item config id2labelitem writer writefindextitemn logger infofwrote predictions to outputtestfile endregion if trainingargs outputdir is not none and not trainingargs pushtohub if we re not pushing to hub at least save a local copy when we re done model savepretrainedtrainingargs outputdir if name main main 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 fine tuning the library models for sequence classification you can also adapt this script on your own text classification task pointers for this are left as comments reduce the amount of console output from tf noqa e402 region helper classes hugging face models have a save_pretrained method that saves both the weights and the necessary metadata to allow them to be loaded as a pretrained model in future this is a simple keras callback that saves the model with this method after each epoch endregion region command line arguments arguments pertaining to what data we are going to input our model for training and eval using hfargumentparser we can turn this class into argparse arguments to be able to specify them on the command line arguments pertaining to which model config tokenizer we are going to fine tune from endregion region argument parsing see all possible arguments in src transformers training_args py or by passing the help flag to this script we now keep distinct sets of args for a cleaner separation of concerns if we pass only one argument to the script and it s the path to a json file let s parse it to get our arguments sending telemetry tracking the example usage helps us better allocate resources to maintain them the information sent is the one passed as arguments along with your python pytorch versions endregion region checkpoints detecting last checkpoint endregion region logging endregion region loading data for csv json files this script will use the label field as the label and the sentence1 and optionally sentence2 fields as inputs if they exist if not the first two fields not named label are used if at least two columns are provided note that the term sentence can be slightly misleading as they often contain more than a single grammatical sentence when the task requires it 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 loading a dataset from local csv files loading a dataset from local json files see more about loading any type of standard or custom dataset at https huggingface co docs datasets loading_datasets endregion region label preprocessing if you ve passed us a training set we try to infer your labels from it by default we assume that if your label column looks like a float then you re doing regression and if not then you re doing classification this is something you may want to change a useful fast method https huggingface co docs datasets package_reference main_classes html datasets dataset unique let s sort it for determinism if you haven t passed a training set we read label info from the saved model this happens later endregion region load model config and tokenizer endregion region dataset preprocessing again we try to have some nice defaults but don t hesitate to tweak to your use case ensure that our labels match the model s if it has some pre specified use the model s labels now we ve established our label2id let s overwrite the model config with it just load the data from the model tokenize the texts map labels to ids endregion region load pretrained model set seed before initializing model in distributed training the from_pretrained methods guarantee that only one local process can concurrently download model vocab endregion region convert data to a tf data dataset model prepare_tf_dataset wraps a hugging face dataset in a tf data dataset which is ready to use in training this is the recommended way to use a hugging face dataset when training with keras you can also use the lower level dataset to_tf_dataset method but you will have to specify things like column names yourself if you use this method whereas they are automatically inferred from the model input names when using model prepare_tf_dataset for more info see the docs https huggingface co docs transformers main en main_classes model transformers tfpretrainedmodel prepare_tf_dataset https huggingface co docs datasets main en package_reference main_classes datasets dataset to_tf_dataset endregion region optimizer loss and compilation transformers models compute the right loss for their task by default when labels are passed and will use this for training unless you specify your own loss function in compile endregion region preparing push_to_hub and model card endregion region training and validation endregion region prediction endregion if we re not pushing to hub at least save a local copy when we re done
import json
import logging
import os
import sys
import warnings
from dataclasses import dataclass, field
from pathlib import Path
from typing import Optional

import numpy as np
from datasets import load_dataset

from transformers import (
    AutoConfig,
    AutoTokenizer,
    HfArgumentParser,
    PretrainedConfig,
    PushToHubCallback,
    TFAutoModelForSequenceClassification,
    TFTrainingArguments,
    create_optimizer,
    set_seed,
)
from transformers.utils import CONFIG_NAME, TF2_WEIGHTS_NAME, send_example_telemetry


os.environ["TF_CPP_MIN_LOG_LEVEL"] = "1"  
import tensorflow as tf  


logger = logging.getLogger(__name__)



class SavePretrainedCallback(tf.keras.callbacks.Callback):
    
    
    
    def __init__(self, output_dir, **kwargs):
        super().__init__()
        self.output_dir = output_dir

    def on_epoch_end(self, epoch, logs=None):
        self.model.save_pretrained(self.output_dir)






@dataclass
class DataTrainingArguments:
    

    train_file: Optional[str] = field(
        default=None, metadata={"help": "A csv or a json file containing the training data."}
    )
    validation_file: Optional[str] = field(
        default=None, metadata={"help": "A csv or a json file containing the validation data."}
    )
    test_file: Optional[str] = field(default=None, metadata={"help": "A csv or a json file containing the test data."})

    max_seq_length: int = field(
        default=128,
        metadata={
            "help": (
                "The maximum total input sequence length after tokenization. Sequences longer "
                "than this will be truncated, sequences shorter will be padded."
            )
        },
    )
    overwrite_cache: bool = field(
        default=False, metadata={"help": "Overwrite the cached preprocessed datasets or not."}
    )
    pad_to_max_length: bool = field(
        default=False,
        metadata={
            "help": (
                "Whether to pad all samples to `max_seq_length`. "
                "If False, will pad the samples dynamically when batching to the maximum length in the batch. "
                "Data will always be padded when using TPUs."
            )
        },
    )
    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_val_samples: Optional[int] = field(
        default=None,
        metadata={
            "help": (
                "For debugging purposes or quicker training, truncate the number of validation examples to this "
                "value if set."
            )
        },
    )
    max_test_samples: Optional[int] = field(
        default=None,
        metadata={
            "help": (
                "For debugging purposes or quicker training, truncate the number of test examples to this "
                "value if set."
            )
        },
    )

    def __post_init__(self):
        train_extension = self.train_file.split(".")[-1].lower() if self.train_file is not None else None
        validation_extension = (
            self.validation_file.split(".")[-1].lower() if self.validation_file is not None else None
        )
        test_extension = self.test_file.split(".")[-1].lower() if self.test_file is not None else None
        extensions = {train_extension, validation_extension, test_extension}
        extensions.discard(None)
        assert len(extensions) != 0, "Need to supply at least one of --train_file, --validation_file or --test_file!"
        assert len(extensions) == 1, "All input files should have the same file extension, either csv or json!"
        assert "csv" in extensions or "json" in extensions, "Input files should have either .csv or .json extensions!"
        self.input_file_extension = extensions.pop()


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





def main():
    
    
    
    

    parser = HfArgumentParser((ModelArguments, DataTrainingArguments, TFTrainingArguments))
    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_text_classification", model_args, data_args, framework="tensorflow")

    output_dir = Path(training_args.output_dir)
    output_dir.mkdir(parents=True, exist_ok=True)
    

    
    
    checkpoint = None
    if len(os.listdir(training_args.output_dir)) > 0 and not training_args.overwrite_output_dir:
        if (output_dir / CONFIG_NAME).is_file() and (output_dir / TF2_WEIGHTS_NAME).is_file():
            checkpoint = output_dir
            logger.info(
                f"Checkpoint detected, resuming training from checkpoint in {training_args.output_dir}. To avoid this"
                " behavior, change the `--output_dir` or add `--overwrite_output_dir` to train from scratch."
            )
        else:
            raise ValueError(
                f"Output directory ({training_args.output_dir}) already exists and is not empty. "
                "Use --overwrite_output_dir to continue regardless."
            )

    

    
    logging.basicConfig(
        format="%(asctime)s - %(levelname)s - %(name)s - %(message)s",
        datefmt="%m/%d/%Y %H:%M:%S",
        handlers=[logging.StreamHandler(sys.stdout)],
    )
    logger.setLevel(logging.INFO)

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

    
    
    
    
    
    
    
    
    
    
    
    data_files = {"train": data_args.train_file, "validation": data_args.validation_file, "test": data_args.test_file}
    data_files = {key: file for key, file in data_files.items() if file is not None}

    for key in data_files.keys():
        logger.info(f"Loading a local file for {key}: {data_files[key]}")

    if data_args.input_file_extension == "csv":
        
        datasets = load_dataset(
            "csv",
            data_files=data_files,
            cache_dir=model_args.cache_dir,
            token=model_args.token,
        )
    else:
        
        datasets = load_dataset("json", data_files=data_files, cache_dir=model_args.cache_dir)
    
    
    

    
    
    if "train" in datasets:
        
        
        is_regression = datasets["train"].features["label"].dtype in ["float32", "float64"]
        if is_regression:
            num_labels = 1
        else:
            
            
            label_list = datasets["train"].unique("label")
            label_list.sort()  
            num_labels = len(label_list)
    
    else:
        num_labels = None
        label_list = None
        is_regression = None
    

    
    if checkpoint is not None:
        config_path = training_args.output_dir
    elif model_args.config_name:
        config_path = model_args.config_name
    else:
        config_path = model_args.model_name_or_path
    if num_labels is not None:
        config = AutoConfig.from_pretrained(
            config_path,
            num_labels=num_labels,
            cache_dir=model_args.cache_dir,
            revision=model_args.model_revision,
            token=model_args.token,
            trust_remote_code=model_args.trust_remote_code,
        )
    else:
        config = AutoConfig.from_pretrained(
            config_path,
            cache_dir=model_args.cache_dir,
            revision=model_args.model_revision,
            token=model_args.token,
            trust_remote_code=model_args.trust_remote_code,
        )
    tokenizer = AutoTokenizer.from_pretrained(
        model_args.tokenizer_name if model_args.tokenizer_name else model_args.model_name_or_path,
        cache_dir=model_args.cache_dir,
        revision=model_args.model_revision,
        token=model_args.token,
        trust_remote_code=model_args.trust_remote_code,
    )
    

    
    
    column_names = {col for cols in datasets.column_names.values() for col in cols}
    non_label_column_names = [name for name in 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"
    elif "sentence1" in non_label_column_names:
        sentence1_key, sentence2_key = "sentence1", None
    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

    if data_args.max_seq_length > tokenizer.model_max_length:
        logger.warning(
            f"The max_seq_length passed ({data_args.max_seq_length}) is larger than the maximum length for the "
            f"model ({tokenizer.model_max_length}). Using max_seq_length={tokenizer.model_max_length}."
        )
    max_seq_length = min(data_args.max_seq_length, tokenizer.model_max_length)

    
    if "train" in datasets:
        if not is_regression and config.label2id != PretrainedConfig(num_labels=num_labels).label2id:
            label_name_to_id = config.label2id
            if sorted(label_name_to_id.keys()) == sorted(label_list):
                label_to_id = label_name_to_id  
            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:"
                    f" {sorted(label_list)}.\nIgnoring the model labels as a result.",
                )
                label_to_id = {v: i for i, v in enumerate(label_list)}
        elif not is_regression:
            label_to_id = {v: i for i, v in enumerate(label_list)}
        else:
            label_to_id = None
        
        config.label2id = label_to_id
        if config.label2id is not None:
            config.id2label = {id: label for label, id in label_to_id.items()}
        else:
            config.id2label = None
    else:
        label_to_id = config.label2id  

    if "validation" in datasets and config.label2id is not None:
        validation_label_list = datasets["validation"].unique("label")
        for val_label in validation_label_list:
            assert val_label in label_to_id, f"Label {val_label} is in the validation set but not the training set!"

    def preprocess_function(examples):
        
        args = (
            (examples[sentence1_key],) if sentence2_key is None else (examples[sentence1_key], examples[sentence2_key])
        )
        result = tokenizer(*args, max_length=max_seq_length, truncation=True)

        
        if config.label2id is not None and "label" in examples:
            result["label"] = [(config.label2id[l] if l != -1 else -1) for l in examples["label"]]
        return result

    datasets = datasets.map(preprocess_function, batched=True, load_from_cache_file=not data_args.overwrite_cache)

    

    with training_args.strategy.scope():
        
        
        set_seed(training_args.seed)
        
        
        
        if checkpoint is None:
            model_path = model_args.model_name_or_path
        else:
            model_path = checkpoint
        model = TFAutoModelForSequenceClassification.from_pretrained(
            model_path,
            config=config,
            cache_dir=model_args.cache_dir,
            revision=model_args.model_revision,
            token=model_args.token,
            trust_remote_code=model_args.trust_remote_code,
        )
        

        
        dataset_options = tf.data.Options()
        dataset_options.experimental_distribute.auto_shard_policy = tf.data.experimental.AutoShardPolicy.OFF
        num_replicas = training_args.strategy.num_replicas_in_sync

        tf_data = {}
        max_samples = {
            "train": data_args.max_train_samples,
            "validation": data_args.max_val_samples,
            "test": data_args.max_test_samples,
        }
        for key in ("train", "validation", "test"):
            if key not in datasets:
                tf_data[key] = None
                continue
            if (
                (key == "train" and not training_args.do_train)
                or (key == "validation" and not training_args.do_eval)
                or (key == "test" and not training_args.do_predict)
            ):
                tf_data[key] = None
                continue
            if key in ("train", "validation"):
                assert "label" in datasets[key].features, f"Missing labels from {key} data!"
            if key == "train":
                shuffle = True
                batch_size = training_args.per_device_train_batch_size * num_replicas
            else:
                shuffle = False
                batch_size = training_args.per_device_eval_batch_size * num_replicas
            samples_limit = max_samples[key]
            dataset = datasets[key]
            if samples_limit is not None:
                dataset = dataset.select(range(samples_limit))

            
            
            
            
            
            
            
            

            data = model.prepare_tf_dataset(
                dataset,
                shuffle=shuffle,
                batch_size=batch_size,
                tokenizer=tokenizer,
            )
            data = data.with_options(dataset_options)
            tf_data[key] = data
        

        

        if training_args.do_train:
            num_train_steps = len(tf_data["train"]) * training_args.num_train_epochs
            if training_args.warmup_steps > 0:
                num_warmup_steps = training_args.warmup_steps
            elif training_args.warmup_ratio > 0:
                num_warmup_steps = int(num_train_steps * training_args.warmup_ratio)
            else:
                num_warmup_steps = 0

            optimizer, schedule = create_optimizer(
                init_lr=training_args.learning_rate,
                num_train_steps=num_train_steps,
                num_warmup_steps=num_warmup_steps,
                adam_beta1=training_args.adam_beta1,
                adam_beta2=training_args.adam_beta2,
                adam_epsilon=training_args.adam_epsilon,
                weight_decay_rate=training_args.weight_decay,
                adam_global_clipnorm=training_args.max_grad_norm,
            )
        else:
            optimizer = None
        if is_regression:
            metrics = []
        else:
            metrics = ["accuracy"]
        
        
        model.compile(optimizer=optimizer, metrics=metrics)
        

        
        push_to_hub_model_id = training_args.push_to_hub_model_id
        model_name = model_args.model_name_or_path.split("/")[-1]
        if not push_to_hub_model_id:
            push_to_hub_model_id = f"{model_name}-finetuned-text-classification"

        model_card_kwargs = {"finetuned_from": model_args.model_name_or_path, "tasks": "text-classification"}

        if training_args.push_to_hub:
            callbacks = [
                PushToHubCallback(
                    output_dir=training_args.output_dir,
                    hub_model_id=push_to_hub_model_id,
                    hub_token=training_args.push_to_hub_token,
                    tokenizer=tokenizer,
                    **model_card_kwargs,
                )
            ]
        else:
            callbacks = []
        

        
        if tf_data["train"] is not None:
            model.fit(
                tf_data["train"],
                validation_data=tf_data["validation"],
                epochs=int(training_args.num_train_epochs),
                callbacks=callbacks,
            )
        if tf_data["validation"] is not None:
            logger.info("Computing metrics on validation data...")
            if is_regression:
                loss = model.evaluate(tf_data["validation"])
                logger.info(f"Eval loss: {loss:.5f}")
            else:
                loss, accuracy = model.evaluate(tf_data["validation"])
                logger.info(f"Eval loss: {loss:.5f}, Eval accuracy: {accuracy * 100:.4f}%")
            if training_args.output_dir is not None:
                output_eval_file = os.path.join(training_args.output_dir, "all_results.json")
                eval_dict = {"eval_loss": loss}
                if not is_regression:
                    eval_dict["eval_accuracy"] = accuracy
                with open(output_eval_file, "w") as writer:
                    writer.write(json.dumps(eval_dict))
        

        
        if tf_data["test"] is not None:
            logger.info("Doing predictions on test dataset...")
            predictions = model.predict(tf_data["test"])["logits"]
            predicted_class = np.squeeze(predictions) if is_regression else np.argmax(predictions, axis=1)
            output_test_file = os.path.join(training_args.output_dir, "test_results.txt")
            with open(output_test_file, "w") as writer:
                writer.write("index\tprediction\n")
                for index, item in enumerate(predicted_class):
                    if is_regression:
                        writer.write(f"{index}\t{item:3.3f}\n")
                    else:
                        item = config.id2label[item]
                        writer.write(f"{index}\t{item}\n")
            logger.info(f"Wrote predictions to {output_test_file}!")
        

        if training_args.output_dir is not None and not training_args.push_to_hub:
            
            model.save_pretrained(training_args.output_dir)


if __name__ == "__main__":
    main()