codingutf8 2022 the google research s 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 selftraining for sequence classification import argparse import dataclasses import json import logging import os import shutil from typing import list optional import datasets from accelerate import accelerator from datasets import loaddataset from finetuning import finetune from tqdm auto import tqdm import transformers from transformers import autoconfig setseed from transformers trainerutils import intervalstrategy logger logging getloggername modelbinfile pytorchmodel bin dataclasses dataclass class stmodelarguments arguments pertaining to what data we are going to input our model for training and evaluation trainfile str dataclasses fieldmetadatahelp a csv or a json file containing the training data inferfile str dataclasses fieldmetadatahelp a csv or a json file containing the data to predict on evalfile optionalstr dataclasses field defaultnone metadatahelp a csv or a json file containing the validation data taskname optionalstr dataclasses field defaultnone metadatahelp the name of the task to train on labellist optionalliststr dataclasses field defaultnone metadatahelp the list of labels for the task dataclasses dataclass class sttrainingarguments create pseudeo labeled data for the next selftraining iteration dataset datasets concatenatedatasetsinferinput inferoutput axis1 if args dofilterbyconfidence dataset dataset filterlambda example exampleprobability args confidencethreshold if args dofilterbyvalperformance assert evalresult 0 0 and evalresult 1 0 numselectedrows intevalresult lendataset printnumselectedrows dataset dataset sortprobability reversetrue dataset dataset selectrangenumselectedrows dataset dataset removecolumnslabel probability dataset dataset renamecolumnprediction label dataset dataset maplambda example label id2labelexamplelabel dataset dataset shuffleseedargs seed pseudolabeleddatafile os path joinnextdatadir ftrainpseudo args datafileextension if args datafileextension csv dataset tocsvpseudolabeleddatafile indexfalse else dataset tojsonpseudolabeleddatafile def selftrainmodelnameorpath trainfile inferfile outputdir kwargs initialize the accelerator we will let the accelerator handle device placement for us accelerator 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 setup logging we only want one process per machine to log things on the screen accelerator islocalmainprocess is only true for one process per machine logger setlevellogging info if accelerator islocalmainprocess else logging error if accelerator islocalmainprocess datasets utils logging setverbositywarning transformers utils logging setverbosityinfo else datasets utils logging setverbosityerror transformers utils logging setverbosityerror modelargs stmodelargumentsmodelnameorpathmodelnameorpath dataargs stdataargumentstrainfiletrainfile inferfileinferfile trainingargs sttrainingargumentsoutputdiroutputdir args argparse namespace for argclass in modelargs dataargs trainingargs for key value in varsargclass items setattrargs key value for key value in kwargs items if hasattrargs key setattrargs key value sanity checks datafiles args datafileextension none you need to provide the training data and the data to predict on assert args trainfile is not none assert args inferfile is not none datafilestrain args trainfile datafilesinfer args inferfile if args evaluationstrategy intervalstrategy no value assert args evalfile is not none datafileseval args evalfile for key in datafiles extension datafileskey split 1 assert extension in csv json fkeyfile should be a csv or a json file if args datafileextension is none args datafileextension extension else assert extension args datafileextension fkeyfile should be a args datafileextension file assert args evalmetric in datasets listmetrics fargs evalmetric not in the list of supported metrics datasets listmetrics if passed along set the training seed now if args seed is not none setseedargs seed logger infocreating the initial data directory for selftraining datadirformat fargs outputdirselftrainiter format initialdatadir datadirformat0 if accelerator ismainprocess if args outputdir is not none os makedirsargs outputdir existoktrue os makedirsinitialdatadir existoktrue accelerator waitforeveryone bestiteration none bestevalresult none earlystoppingpatiencecounter 0 shouldtrainingstop false show the progress bar progressbar tqdmrangeargs maxselftrainiterations disablenot accelerator islocalmainprocess selftrain for iteration in range0 intargs maxselftrainiterations currentdatadir datadirformatiteration assert os path existscurrentdatadir stage 1 initial finetuning for iteration 0 or pseudotraining for iteration 0 currentoutputdir os path joincurrentdatadir stage1 argumentsdict accelerator accelerator modelnameorpath args modelnameorpath cachedir args cachedir dotrain true trainfile datafilestrain if iteration 0 else datafilestrainpseudo doeval true if args evalfile is not none else false evalfile datafileseval dopredict true inferfile datafilesinfer taskname args taskname labellist args labellist outputdir currentoutputdir evalmetric args evalmetric evaluationstrategy args evaluationstrategy earlystoppingpatience args earlystoppingpatience earlystoppingthreshold args earlystoppingthreshold seed args seed add additional training arguments for key value in kwargs items if key not in argumentsdict and not hasattrtrainingargs key argumentsdict updatekey value modelbinfilepath os path joincurrentoutputdir bestcheckpoint modelbinfile if os path existsmodelbinfilepath logger info found existing model checkpoint at s skipping selftraining iteration d stage 1 modelbinfilepath iteration else logger info running selftraining iteration d stage 1 iteration finetuneargumentsdict accelerator waitforeveryone assert os path existsmodelbinfilepath logger infoselftraining job completed iteration d stage 1 iteration if iteration 0 and args finetuneonlabeleddata stage 2 optional finetuning on the original labeled data modelpath os path joincurrentoutputdir bestcheckpoint currentoutputdir os path joincurrentdatadir stage2 update argumentsdict argumentsdictmodelnameorpath modelpath argumentsdicttrainfile datafilestrain argumentsdictoutputdir currentoutputdir modelbinfilepath os path joincurrentoutputdir bestcheckpoint modelbinfile if os path existsmodelbinfilepath logger info found existing model checkpoint at s skipping selftraining iteration d stage 2 modelbinfilepath iteration else logger info running selftraining iteration d stage 2 iteration finetuneargumentsdict accelerator waitforeveryone assert os path existsmodelbinfilepath logger infoselftraining job completed iteration d stage 2 iteration newiteration iteration nextdatadir datadirformatiteration 1 config autoconfig frompretrainedos path joincurrentoutputdir bestcheckpoint id2label config id2label evalresultsfile os path joincurrentoutputdir evalresultsbestcheckpoint json testresultsfile os path joincurrentoutputdir testresultsbestcheckpoint json assert os path existsevalresultsfile with openevalresultsfile r as f evalresult floatjson loadfargs evalmetric inferoutputfile os path joincurrentoutputdir inferoutputbestcheckpoint csv assert os path existsinferoutputfile loading the dataset from local csv or json files inferinput loaddatasetargs datafileextension datafilesdata datafilesinferdata inferoutput loaddatasetcsv datafilesdata inferoutputfiledata if accelerator ismainprocess os makedirsnextdatadir existoktrue shutil copyevalresultsfile os path joinoutputdir fevalresultsiteriteration json if os path existstestresultsfile shutil copyevalresultsfile os path joinoutputdir ftestresultsiteriteration json createpseudolabeleddataargs inferinput inferoutput evalresult id2label nextdatadir accelerator waitforeveryone datafilestrainpseudo os path joinnextdatadir ftrainpseudo args datafileextension if args evaluationstrategy intervalstrategy no value newevalresult evalresult if bestiteration is none bestiteration newiteration bestevalresult newevalresult else if newevalresult bestevalresult args earlystoppingthreshold bestiteration newiteration bestevalresult newevalresult earlystoppingpatiencecounter 0 else if newevalresult bestevalresult bestiteration newiteration bestevalresult newevalresult earlystoppingpatiencecounter 1 if earlystoppingpatiencecounter args earlystoppingpatience shouldtrainingstop true progressbar update1 if shouldtrainingstop break if bestiteration is not none save the best iteration logger infobest iteration d bestiteration logger infobest evaluation result s f args evalmetric bestevalresult accelerator waitforeveryone if accelerator ismainprocess shutil copy os path joinoutputdir fevalresultsiteriteration json os path joinoutputdir evalresultsbestiteration json else assume that the last iteration is the best logger infobest iteration d args maxselftrainiterations 1 logger infobest evaluation result s f args evalmetric evalresult accelerator waitforeveryone if accelerator ismainprocess shutil copy os path joinoutputdir fevalresultsiterargs maxselftrainiterations 1 json os path joinoutputdir evalresultsbestiteration json coding utf 8 2022 the google research s 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 self training for sequence classification arguments pertaining to which config tokenizer model we are going to fine tune from arguments pertaining to what data we are going to input our model for training and evaluation training arguments pertaining to the training loop itself create pseudeo labeled data for the next self training iteration self training a pre trained model on a downstream task args model_name_or_path path to pretrained model or model identifier from huggingface co models train_file a csv or a json file containing the training data infer_file a csv or a json file containing the data to predict on output_dir the output directory where the model predictions and checkpoints will be written kwargs dictionary of key value pairs with which to update the configuration object after loading the values in kwargs of any keys which are configuration attributes will be used to override the loaded values initialize the accelerator we will let the accelerator handle device placement for us make one log on every process with the configuration for debugging setup logging we only want one process per machine to log things on the screen accelerator is_local_main_process is only true for one process per machine sanity checks you need to provide the training data and the data to predict on if passed along set the training seed now show the progress bar self train stage 1 initial fine tuning for iteration 0 or pseudo training for iteration 0 add additional training arguments stage 2 optional fine tuning on the original labeled data update arguments_dict loading the dataset from local csv or json files save the best iteration assume that the last iteration is the best
import argparse
import dataclasses
import json
import logging
import os
import shutil
from typing import List, Optional

import datasets
from accelerate import Accelerator
from datasets import load_dataset
from finetuning import finetune
from tqdm.auto import tqdm

import transformers
from transformers import AutoConfig, set_seed
from transformers.trainer_utils import IntervalStrategy


logger = logging.getLogger(__name__)

MODEL_BIN_FILE = "pytorch_model.bin"


@dataclasses.dataclass
class STModelArguments:
    

    model_name_or_path: str = dataclasses.field(
        metadata={"help": "Path to pretrained model or model identifier from huggingface.co/models."}
    )
    cache_dir: Optional[str] = dataclasses.field(
        default=None,
        metadata={"help": "Where do you want to store the pretrained models downloaded from huggingface.co."},
    )


@dataclasses.dataclass
class STDataArguments:
    

    train_file: str = dataclasses.field(metadata={"help": "A csv or a json file containing the training data."})
    infer_file: str = dataclasses.field(metadata={"help": "A csv or a json file containing the data to predict on."})
    eval_file: Optional[str] = dataclasses.field(
        default=None, metadata={"help": "A csv or a json file containing the validation data."}
    )
    task_name: Optional[str] = dataclasses.field(
        default=None,
        metadata={"help": "The name of the task to train on."},
    )
    label_list: Optional[List[str]] = dataclasses.field(
        default=None, metadata={"help": "The list of labels for the task."}
    )


@dataclasses.dataclass
class STTrainingArguments:
    

    output_dir: str = dataclasses.field(
        metadata={"help": "The output directory where the model predictions and checkpoints will be written."}
    )
    eval_metric: Optional[str] = dataclasses.field(
        default="accuracy", metadata={"help": "The evaluation metric used for the task."}
    )
    evaluation_strategy: Optional[str] = dataclasses.field(
        default="no",
        metadata={
            "help": 'The evaluation strategy to adopt during training. Possible values are: ["no", "step", "epoch]'
        },
    )
    early_stopping_patience: Optional[int] = dataclasses.field(
        default=10,
        metadata={"help": "Number of evaluation calls with no improvement after which training will be stopped."},
    )
    early_stopping_threshold: Optional[float] = dataclasses.field(
        default=0.0,
        metadata={
            "help": "How much the specified evaluation metric must improve to satisfy early stopping conditions."
        },
    )
    do_filter_by_confidence: Optional[bool] = dataclasses.field(
        default=False,
        metadata={"help": "Whether to filter the pseudo-labeled data based on the confidence score."},
    )
    do_filter_by_val_performance: Optional[bool] = dataclasses.field(
        default=False,
        metadata={"help": "Whether to filter the pseudo-labeled data based on the validation performance."},
    )
    finetune_on_labeled_data: Optional[bool] = dataclasses.field(
        default=False,
        metadata={"help": "Whether to fine-tune on labeled data after pseudo training."},
    )
    confidence_threshold: Optional[float] = dataclasses.field(
        default=0.0,
        metadata={"help": "Confidence threshold for pseudo-labeled data filtering."},
    )
    max_selftrain_iterations: Optional[int] = dataclasses.field(
        default=100,
        metadata={"help": "Number of evaluation calls with no improvement after which training will be stopped."},
    )
    seed: Optional[int] = dataclasses.field(
        default=None,
        metadata={"help": "Random seed for initialization."},
    )


def create_pseudo_labeled_data(args, infer_input, infer_output, eval_result, id2label, next_data_dir):
    

    dataset = datasets.concatenate_datasets([infer_input, infer_output], axis=1)

    if args.do_filter_by_confidence:
        dataset = dataset.filter(lambda example: example["probability"] > args.confidence_threshold)

    if args.do_filter_by_val_performance:
        assert eval_result >= 0.0 and eval_result <= 1.0
        num_selected_rows = int(eval_result * len(dataset))
        print(num_selected_rows)
        dataset = dataset.sort("probability", reverse=True)
        dataset = dataset.select(range(num_selected_rows))

    dataset = dataset.remove_columns(["label", "probability"])
    dataset = dataset.rename_column("prediction", "label")
    dataset = dataset.map(lambda example: {"label": id2label[example["label"]]})
    dataset = dataset.shuffle(seed=args.seed)

    pseudo_labeled_data_file = os.path.join(next_data_dir, f"train_pseudo.{args.data_file_extension}")
    if args.data_file_extension == "csv":
        dataset.to_csv(pseudo_labeled_data_file, index=False)
    else:
        dataset.to_json(pseudo_labeled_data_file)


def selftrain(model_name_or_path, train_file, infer_file, output_dir, **kwargs):
    
    
    
    accelerator = 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)

    
    
    
    logger.setLevel(logging.INFO if accelerator.is_local_main_process else logging.ERROR)

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

    model_args = STModelArguments(model_name_or_path=model_name_or_path)
    data_args = STDataArguments(train_file=train_file, infer_file=infer_file)
    training_args = STTrainingArguments(output_dir=output_dir)
    args = argparse.Namespace()

    for arg_class in (model_args, data_args, training_args):
        for key, value in vars(arg_class).items():
            setattr(args, key, value)

    for key, value in kwargs.items():
        if hasattr(args, key):
            setattr(args, key, value)

    
    data_files = {}
    args.data_file_extension = None

    
    assert args.train_file is not None
    assert args.infer_file is not None
    data_files["train"] = args.train_file
    data_files["infer"] = args.infer_file

    if args.evaluation_strategy != IntervalStrategy.NO.value:
        assert args.eval_file is not None
        data_files["eval"] = args.eval_file

    for key in data_files:
        extension = data_files[key].split(".")[-1]
        assert extension in ["csv", "json"], f"`{key}_file` should be a csv or a json file."
        if args.data_file_extension is None:
            args.data_file_extension = extension
        else:
            assert extension == args.data_file_extension, f"`{key}_file` should be a {args.data_file_extension} file`."

    assert (
        args.eval_metric in datasets.list_metrics()
    ), f"{args.eval_metric} not in the list of supported metrics {datasets.list_metrics()}."

    
    if args.seed is not None:
        set_seed(args.seed)

    logger.info("Creating the initial data directory for self-training...")
    data_dir_format = f"{args.output_dir}/self-train_iter-{{}}".format
    initial_data_dir = data_dir_format(0)

    if accelerator.is_main_process:
        if args.output_dir is not None:
            os.makedirs(args.output_dir, exist_ok=True)
            os.makedirs(initial_data_dir, exist_ok=True)
    accelerator.wait_for_everyone()

    best_iteration = None
    best_eval_result = None
    early_stopping_patience_counter = 0
    should_training_stop = False
    
    progress_bar = tqdm(range(args.max_selftrain_iterations), disable=not accelerator.is_local_main_process)

    
    for iteration in range(0, int(args.max_selftrain_iterations)):
        current_data_dir = data_dir_format(iteration)
        assert os.path.exists(current_data_dir)

        
        
        current_output_dir = os.path.join(current_data_dir, "stage-1")
        arguments_dict = {
            "accelerator": accelerator,
            "model_name_or_path": args.model_name_or_path,
            "cache_dir": args.cache_dir,
            "do_train": True,
            "train_file": data_files["train"] if iteration == 0 else data_files["train_pseudo"],
            "do_eval": True if args.eval_file is not None else False,
            "eval_file": data_files["eval"],
            "do_predict": True,
            "infer_file": data_files["infer"],
            "task_name": args.task_name,
            "label_list": args.label_list,
            "output_dir": current_output_dir,
            "eval_metric": args.eval_metric,
            "evaluation_strategy": args.evaluation_strategy,
            "early_stopping_patience": args.early_stopping_patience,
            "early_stopping_threshold": args.early_stopping_threshold,
            "seed": args.seed,
        }
        
        for key, value in kwargs.items():
            if key not in arguments_dict and not hasattr(training_args, key):
                arguments_dict.update({key: value})

        model_bin_file_path = os.path.join(current_output_dir, "best-checkpoint", MODEL_BIN_FILE)
        if os.path.exists(model_bin_file_path):
            logger.info(
                "Found existing model checkpoint at %s. Skipping self-training: iteration: %d, stage: 1.",
                model_bin_file_path,
                iteration,
            )
        else:
            logger.info("***** Running self-training: iteration: %d, stage: 1 *****", iteration)
            finetune(**arguments_dict)
            accelerator.wait_for_everyone()
            assert os.path.exists(model_bin_file_path)
            logger.info("Self-training job completed: iteration: %d, stage: 1.", iteration)

        if iteration > 0 and args.finetune_on_labeled_data:
            
            model_path = os.path.join(current_output_dir, "best-checkpoint")
            current_output_dir = os.path.join(current_data_dir, "stage-2")
            
            arguments_dict["model_name_or_path"] = model_path
            arguments_dict["train_file"] = data_files["train"]
            arguments_dict["output_dir"] = current_output_dir

            model_bin_file_path = os.path.join(current_output_dir, "best-checkpoint", MODEL_BIN_FILE)
            if os.path.exists(model_bin_file_path):
                logger.info(
                    "Found existing model checkpoint at %s. Skipping self-training: iteration: %d, stage: 2.",
                    model_bin_file_path,
                    iteration,
                )
            else:
                logger.info("***** Running self-training: iteration: %d, stage: 2 *****", iteration)
                finetune(**arguments_dict)
                accelerator.wait_for_everyone()
                assert os.path.exists(model_bin_file_path)
                logger.info("Self-training job completed: iteration: %d, stage: 2.", iteration)

        new_iteration = iteration
        next_data_dir = data_dir_format(iteration + 1)

        config = AutoConfig.from_pretrained(os.path.join(current_output_dir, "best-checkpoint"))
        id2label = config.id2label
        eval_results_file = os.path.join(current_output_dir, "eval_results_best-checkpoint.json")
        test_results_file = os.path.join(current_output_dir, "test_results_best-checkpoint.json")
        assert os.path.exists(eval_results_file)

        with open(eval_results_file, "r") as f:
            eval_result = float(json.load(f)[args.eval_metric])
        infer_output_file = os.path.join(current_output_dir, "infer_output_best-checkpoint.csv")
        assert os.path.exists(infer_output_file)
        
        infer_input = load_dataset(args.data_file_extension, data_files={"data": data_files["infer"]})["data"]
        infer_output = load_dataset("csv", data_files={"data": infer_output_file})["data"]

        if accelerator.is_main_process:
            os.makedirs(next_data_dir, exist_ok=True)
            shutil.copy(eval_results_file, os.path.join(output_dir, f"eval_results_iter-{iteration}.json"))
            if os.path.exists(test_results_file):
                shutil.copy(eval_results_file, os.path.join(output_dir, f"test_results_iter-{iteration}.json"))
            create_pseudo_labeled_data(args, infer_input, infer_output, eval_result, id2label, next_data_dir)
        accelerator.wait_for_everyone()

        data_files["train_pseudo"] = os.path.join(next_data_dir, f"train_pseudo.{args.data_file_extension}")

        if args.evaluation_strategy != IntervalStrategy.NO.value:
            new_eval_result = eval_result

            if best_iteration is None:
                best_iteration = new_iteration
                best_eval_result = new_eval_result
            else:
                if new_eval_result - best_eval_result > args.early_stopping_threshold:
                    best_iteration = new_iteration
                    best_eval_result = new_eval_result
                    early_stopping_patience_counter = 0
                else:
                    if new_eval_result == best_eval_result:
                        best_iteration = new_iteration
                        best_eval_result = new_eval_result
                    early_stopping_patience_counter += 1

                if early_stopping_patience_counter >= args.early_stopping_patience:
                    should_training_stop = True

        progress_bar.update(1)

        if should_training_stop:
            break

    if best_iteration is not None:
        
        logger.info("Best iteration: %d", best_iteration)
        logger.info("Best evaluation result: %s = %f", args.eval_metric, best_eval_result)
        accelerator.wait_for_everyone()
        if accelerator.is_main_process:
            shutil.copy(
                os.path.join(output_dir, f"eval_results_iter-{iteration}.json"),
                os.path.join(output_dir, "eval_results_best-iteration.json"),
            )
    else:
        
        logger.info("Best iteration: %d", args.max_selftrain_iterations - 1)
        logger.info("Best evaluation result: %s = %f", args.eval_metric, eval_result)
        accelerator.wait_for_everyone()
        if accelerator.is_main_process:
            shutil.copy(
                os.path.join(output_dir, f"eval_results_iter-{args.max_selftrain_iterations - 1}.json"),
                os.path.join(output_dir, "eval_results_best-iteration.json"),
            )